research: combine current upstream with Qwen3.8 mmap PLE

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.

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__ = [

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),

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(

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);

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])