From 4b94bdc38a46a4dfe534e8793126160d56904c44 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Sun, 30 Aug 2026 01:00:07 -0700 Subject: [PATCH 1/3] docs: add SECURITY.md Added guidelines for reporting security issues. --- SECURITY.md | 7 +++++++ 1 file changed, 7 insertions(+) create mode 100644 SECURITY.md diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 000000000..b04e4029f --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,7 @@ +# Reporting Security Issues + +To report a security issue, please use the GitHub Security Advisory ["Report a Vulnerability"](https://github.com/FlashML-org/FreeToken/security/advisories/new) tab. Please do not report security issues as public issues or pull requests. + +We will send a response indicating the next steps in handling your report. After the initial reply to your report, the maintainers will keep you informed of the progress towards a fix and full announcement, and may ask for additional information or guidance. + +Report security bugs in third-party dependencies to the person or team maintaining the dependency. From 3a20a79038338c33bd051c52152e6d1faa4d9791 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Sun, 30 Aug 2026 22:43:50 -0700 Subject: [PATCH 2/3] fix(kernel): unbreak the nightly kernel-cache wheel build (#310) * fix(kernel): move fp8_block_scale_pad into aot_models to unbreak the kernel-cache build * fix(kernel): exclude bank rows fast_index_copy cannot compile from aot specs * fix(moe): fail fast when fused copy is off and a bank row cannot fall back --- python/freetoken/kernel/aot_models.py | 12 +++++++++--- python/freetoken/moe/offload_cache.py | 20 +++++++++++++++----- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index c9c2fb98e..6268ae338 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -63,6 +63,13 @@ class AotModel: arch_aliases: tuple[str, ...] = () +def fp8_block_scale_pad(rows: int, cols: int) -> int: + """Trailing scale-bank dim padded so per-expert row bytes are 16B-aligned (fused copy).""" + while (rows * cols * 2) % 16: + cols += 1 + return cols + + def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int) -> dict[str, int]: """Per-expert row bytes for each offload bank a format registers. @@ -77,8 +84,6 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int if fmt == "fp8_block": # qwen3_5_moe/weight.py _build_fp8_expert_banks: fp8 weights + bf16 128x128 block # scales, trailing scale dim 16B-padded (same helper as the loader) - from freetoken.moe.offload_cache import fp8_block_scale_pad - B = 128 return { "gate_up": 2 * I * H, @@ -417,7 +422,8 @@ def aggregate_fast_index_copy_feature_sizes() -> tuple[int, ...]: sizes: set[int] = set(TEST_FEATURE_SIZES) for model in SUPPORTED_MODELS: sizes.update(fast_index_copy_feature_sizes(model)) - return tuple(sorted(sizes)) + # the per-bank kernel copies rows in fixed 128-byte steps; other sizes cannot compile + return tuple(sorted(size for size in sizes if size % 128 == 0)) __all__ = [ diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 33a47553b..e1f20dd2f 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -77,11 +77,8 @@ "ds_fp4": ("gate_up_packed", "gate_up_scale", "down_packed", "down_scale"), } -def fp8_block_scale_pad(rows: int, cols: int) -> int: - """Trailing scale-bank dim padded so per-expert row bytes are 16B-aligned (fused copy).""" - while (rows * cols * 2) % 16: - cols += 1 - return cols +# lives in kernel/aot_models.py: the AOT row table shares it and must stay importable in the torch-only kernel-cache build env, which cannot import freetoken.moe +from freetoken.kernel.aot_models import fp8_block_scale_pad # bytes per (expert, layer) as f(hidden, moe_intermediate), from the bank shapes above; keep in sync with _BANK_SCHEMAS @@ -350,6 +347,19 @@ def set_bank_sources( self._init_prefill_overlap_buffers() def _build_copy_plan(self) -> None: + self._build_fused_copy_plan() + if self._copy_fused_ok or self.device.type != "cuda" or not self.banks: + return + for name in self.bank_schema: + cache = self.bank_caches[name] + feat = math.prod(cache.shape[1:]) * cache.element_size() + if feat % 128: + raise RuntimeError( + f"MoE bank {name!r} rows are {feat} bytes (not a multiple of 128): " + f"only the fused multi-bank copy can move them, but it is disabled" + ) + + def _build_fused_copy_plan(self) -> None: """Precompute the fused multi-bank copy descriptor (base addrs + per-row bytes). Built once here (and on :meth:`rebuild`, which reallocates the slot caches); From e05cff83a04b322fc7823678aa2d05c826aad26c Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Mon, 31 Aug 2026 20:21:11 -0700 Subject: [PATCH 3/3] perf(moe): route fused_topk through the in-repo triton router (#319) --- python/freetoken/kernel/backend.py | 11 ------ python/freetoken/moe/fused.py | 57 +++--------------------------- tests/moe/test_fused_moe.py | 6 ++-- 3 files changed, 8 insertions(+), 66 deletions(-) diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 3037ad8d7..fc42c49fd 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -31,17 +31,6 @@ def is_sgl_kernel_installed() -> bool: return _importable("sgl_kernel") -@functools.cache -def is_triton_kernels_installed() -> bool: - """OpenAI's ``triton_kernels`` (the fused MoE router used by ``moe.fused.fused_topk``). - - Distinct from the ``triton`` runtime we always depend on: it ships with the Triton - source tree and has no Windows wheel. It is also not one of the six ops - ``freetoken.kernel.triton`` reimplements, so its call-site carries its own fallback. - """ - return _importable("triton_kernels") - - @functools.cache def driver_cuda_version() -> int | None: """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0), diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index 4b9a4875f..ecc45eae4 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -6,11 +6,7 @@ import torch from freetoken.moe import BaseMoeBackend -from freetoken.utils import div_ceil, init_logger - -logger = init_logger(__name__) - -_warned_torch_topk = False +from freetoken.utils import div_ceil def _torch_fused_topk( @@ -19,7 +15,7 @@ def _torch_fused_topk( renormalize: bool, num_token_non_padded: torch.Tensor | None, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Pure-torch softmax router matching triton_kernels.topk (Windows fallback). + """Pure-torch reference for the fused softmax router; tests compare the kernel against it. Softmax over all experts, select the top-k, and (when ``renormalize``) rescale the selected weights to sum to 1 -- the standard fused-MoE routing convention. @@ -44,52 +40,9 @@ def fused_topk( ) -> Tuple[torch.Tensor, torch.Tensor]: assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" - from freetoken.kernel.backend import is_triton_kernels_installed - - # triton_kernels ships no Windows wheel, and unlike flashinfer/sgl_kernel it is not one - # of the six ops the in-repo triton kernels cover -- so this router needs its own fallback. - if not is_triton_kernels_installed(): - global _warned_torch_topk - if not _warned_torch_topk: - _warned_torch_topk = True - # Once, not per call: this runs every MoE forward. On Linux a missing - # triton_kernels used to fail fast with ImportError; keep the misconfiguration - # visible without giving up the fallback that Windows needs. - logger.warning_rank0( - "fused_topk: triton_kernels is not installed -> pure-torch router fallback " - "(numerically equivalent, slower). Expected on Windows (no wheel); on Linux " - "install triton_kernels to restore the fused router." - ) - return _torch_fused_topk(gating_output, topk, renormalize, num_token_non_padded) - - if topk & (topk - 1): - # triton_kernels.topk builds tl.arange(0, k), which must be a power of 2; a - # top-10 router (qwen4_exp) takes the equivalent vendored triton router instead. - from freetoken.kernel.triton.moe_router import fused_topk_softmax - - return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded) - - from triton_kernels.topk import topk as triton_kernels_topk - - logits = gating_output.float() - softmax_first = not renormalize - if softmax_first: - logits = torch.softmax(logits, dim=-1) - sparse_topk = triton_kernels_topk( - logits, - topk, - apply_softmax=not softmax_first, - ) - if hasattr(sparse_topk, "vals"): - topk_weights = sparse_topk.vals - topk_ids = sparse_topk.indx - else: - topk_weights, topk_ids = sparse_topk[:2] - topk_ids = topk_ids.to(torch.int32) - if num_token_non_padded is not None: - indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device) - topk_ids[indices >= num_token_non_padded, :] = -1 - return topk_weights, topk_ids + from freetoken.kernel.triton.moe_router import fused_topk_softmax + + return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded) def moe_align_block_size( diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py index 41a75c20e..cafe77844 100644 --- a/tests/moe/test_fused_moe.py +++ b/tests/moe/test_fused_moe.py @@ -277,8 +277,8 @@ def test_fused_experts_decode_activation_and_router_weight_modes( @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_fused_topk_non_power_of_2_k_routes_vendored_router(): - """triton_kernels.topk builds tl.arange(0, k) (power-of-2 only); k=10 must not reach it.""" +def test_fused_topk_handles_non_power_of_2_k(): + """A top-10 router (qwen4_exp) must route like any other k.""" from freetoken.moe.fused import _torch_fused_topk, fused_topk gating = torch.randn(5, 64, device="cuda") @@ -289,7 +289,7 @@ def test_fused_topk_non_power_of_2_k_routes_vendored_router(): torch.testing.assert_close(weights, ref_w, rtol=1e-5, atol=1e-6) -# The vendored triton router behind that k=10 branch; fp32 logits keep the reference top-k tie-free. +# The in-repo triton router behind fused_topk; fp32 logits keep the reference top-k tie-free. # Ties get their own case below, because torch.topk does not break them by expert id. @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @pytest.mark.parametrize("renormalize", [True, False])