feat(moe): load a short prefill chunk's experts on demand through the slot cache

docs/cli.md

@@ -94,6 +94,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does, and [perf
 | `--moe-locked-fetch` | auto | With `hybrid` and locked layers: share of a locked layer's misses fetched per step into its GPU slots through pinned staging. Auto takes p / (2p + c) from the `ft bench bw` row of the layer's kind (at most 1 per step without a row); 0 = locked layers decode on CPU only |
 | `--moe-locked-max-fetch` | 8 | Most experts a locked layer fetches per step; sizes the staging (that many rows of the widest expert, per bank) |
 | `--moe-collect-stats` | off | Log decode MoE stats with each decode log line: hit rate, and per layer-step the misses, experts fetched and experts computed on CPU, for pinned, locked and CPU-only layers |
+| `--moe-prefill-demand-tokens` | 0 | Prefill: a chunk of at most this many tokens loads only the experts it routes to, through the expert cache as decode does (GPU decode only); 0 streams whole layers |
 | `--moe-prefill-hit-d2d` | off | Prefill: copy cache-hit experts device-side, stream only misses (CUDA >= 13) |
 | `--disable-moe-prefill-overlap` | overlap on | Disable the two-buffer prefill copy overlap |
 

python/freetoken/engine/config.py

@@ -53,6 +53,7 @@ class EngineConfig:
     moe_cache_split: str = "equal"
     kv_reserve_tokens: int = 8192  # KV floor for --moe-cache-auto; small by design (MoE-priority)
     moe_cache_policy: str = "lru"
+    moe_prefill_demand_tokens: int = 0
     moe_prefill_overlap: bool = True
     # Prefill hit/miss split: serve cache-resident experts D2D during prefill
     # prefetch instead of re-streaming the full layer over PCIe. Needs CUDA >= 13.0

python/freetoken/engine/engine.py

@@ -830,6 +830,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
             cache_size=config.moe_cache_size,
             device=self.device,
             cache_policy=config.moe_cache_policy,
+            prefill_demand_tokens=config.moe_prefill_demand_tokens,
             prefill_overlap=config.moe_prefill_overlap,
             prefill_hit_d2d=config.moe_prefill_hit_d2d,
             quant_format=banks.quant_format,
@@ -1632,6 +1633,7 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int, *, reserved: int
     "moe_hybrid_max_fetch": -1,
     "moe_locked_fetch": -1.0,
     "moe_locked_max_fetch": 8,
+    "moe_prefill_demand_tokens": 0,
     "moe_prefill_overlap": True,
     "moe_prefill_hit_d2d": False,
     "expert_load": "auto",

python/freetoken/layers/moe.py

@@ -10,7 +10,7 @@
 
 
 from .base import BaseOP
-from .quantization import ExpertView, LayerKind, QuantConfig, quant_method_for
+from .quantization import ExpertView, LayerKind, QuantConfig, RoutedExperts, quant_method_for
 
 if TYPE_CHECKING:
     from freetoken.models.config import ModelConfig
@@ -252,8 +252,9 @@ def prefill_forward(
     # ------------------------------------------------------------------
     # Data movement -- one decision tree for every quant format (the banks
     # registry makes the cache machinery bank-count agnostic). Decode loads
-    # on demand; prefill streams whole layers, double-buffered when overlap
-    # is enabled. The kernels only ever see bank views plus row indices;
+    # on demand, and so does a short enough prefill chunk; a longer one
+    # streams whole layers, double-buffered when overlap is enabled. The
+    # kernels only ever see bank views plus row indices;
     # which kernel runs is decided afterwards, in ``_expert_gemm``.
     # ------------------------------------------------------------------
 
@@ -348,12 +349,26 @@ def _prefill_routed(
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
     ) -> torch.Tensor:
-        """Prefill movement: stream whole layers -- double-buffered behind the
-        previous layer's GEMMs when ``prefill_overlap`` is on, else a synchronous
-        ``materialize_layer``. In both, position == expert id, so the routing ids
-        pass through unmapped."""
+        """Prefill movement. A chunk within ``prefill_demand_tokens`` loads its routed experts
+        into the slot cache as decode does and runs the same grouped GEMM over the slots. A
+        longer one streams whole layers -- double-buffered behind the previous layer's GEMMs
+        when ``prefill_overlap`` is on, else a synchronous ``materialize_layer`` -- where
+        position == expert id, so the routing ids pass through unmapped."""
         cache = self.offload_cache
         assert cache is not None
+        if cache.loads_prefill_on_demand(hidden_states.shape[0]):
+            slots = cache.load_prefill_experts(self.layer_id, topk_ids)
+            return self._expert_gemm(
+                cache,
+                hidden_states,
+                topk_weights,
+                slots,
+                views=cache.pool_bank_views(self.layer_id),
+                n=cache.pool_slots(self.layer_id),
+                alphas=cache.alphas_for_slots(self.layer_id),
+                routed=RoutedExperts(self.num_experts, cache.alphas_for_layer(self.layer_id)),
+                is_prefill=True,
+            )
         if cache.prefill_overlap:
             views = self._wait_prefill_overlap(cache)
             out = self._expert_gemm(
@@ -410,12 +425,13 @@ def _expert_gemm(
         n: int | None,
         alphas: tuple[torch.Tensor, torch.Tensor] | None,
         is_prefill: bool,
+        routed: RoutedExperts | None = None,
     ) -> torch.Tensor:
         from freetoken.moe.legacy_format import canonical_role  # legacy_format imports this package
 
         view = ExpertView(
             {canonical_role(name): t for name, t in zip(cache.bank_schema, views)},
-            slots=None if n is not None else topk_ids, n=n, alphas=alphas,
+            slots=None if n is not None else topk_ids, n=n, alphas=alphas, routed=routed,
         )
         return self.quant_method.apply(
             hidden_states, topk_weights, topk_ids, view, layer=self, is_prefill=is_prefill

python/freetoken/layers/quantization/__init__.py

@@ -22,6 +22,7 @@
     Mxfp4MoEMethod,
     Mxfp8MoEMethod,
     Nvfp4MoEMethod,
+    RoutedExperts,
     UnquantizedMoEMethod,
 )
 from .quant_backend import QuantBackend, get_quant_backend, set_quant_backend
@@ -47,7 +48,7 @@
     "EmbeddingConfig", "EmbeddingKernel", "EmbeddingMethod", "UnquantizedEmbeddingMethod", "GGUFEmbeddingMethod",
     "LinearConfig", "LinearKernel", "LinearMethod",
     "UnquantizedLinearMethod", "Fp8TensorLinearMethod", "Fp8BlockLinearMethod", "Mxfp8LinearMethod", "Nvfp4LinearMethod",
-    "BankSpec", "ExpertView", "MoEConfig", "MoEKernel", "MoEMethod",
+    "BankSpec", "ExpertView", "MoEConfig", "MoEKernel", "MoEMethod", "RoutedExperts",
     "UnquantizedMoEMethod", "Fp8BlockMoEMethod", "Nvfp4MoEMethod", "Mxfp4MoEMethod", "Mxfp8MoEMethod",
     "QuantBackend", "set_quant_backend", "get_quant_backend", "NameMap",
     "QuantConfig", "NoQuantConfig", "ModelOptConfig", "CompressedTensorsConfig", "Fp8BlockConfig", "GGUFConfig", "Mxfp4Config",

python/freetoken/layers/quantization/moe/__init__.py

@@ -1,4 +1,4 @@
-from .base import BankSpec, ExpertView, MoEConfig, MoEKernel, MoEMethod
+from .base import BankSpec, ExpertView, MoEConfig, MoEKernel, MoEMethod, RoutedExperts
 from . import fp8_block, gguf, mxfp4, mxfp8, nvfp4, unquantized
 from .fp8_block import Fp8BlockMoEMethod
 from .mxfp4 import Mxfp4MoEMethod
@@ -8,7 +8,7 @@
 from .gguf import GGUFMoEMethod
 
 __all__ = [
-    "BankSpec", "ExpertView", "MoEConfig", "MoEKernel", "MoEMethod",
+    "BankSpec", "ExpertView", "MoEConfig", "MoEKernel", "MoEMethod", "RoutedExperts",
     "UnquantizedMoEMethod", "Fp8BlockMoEMethod", "Nvfp4MoEMethod", "Mxfp4MoEMethod", "Mxfp8MoEMethod", "GGUFMoEMethod",
     "fp8_block", "gguf", "mxfp4", "mxfp8", "nvfp4", "unquantized",
 ]

python/freetoken/layers/quantization/moe/base.py

@@ -77,6 +77,15 @@ class BankSpec:
     resident: bool = False
 
 
+@dataclass(frozen=True)
+class RoutedExperts:
+    """The layer behind the slot view of a demand prefill chunk: the prefill GEMMs size their tiles by its ``count`` of
+    experts and W4A4 takes its activation scales over their ``alphas``, as the stream of that layer would."""
+
+    count: int
+    alphas: tuple[torch.Tensor, torch.Tensor] | None = None
+
+
 @dataclass
 class ExpertView:
     """The weights an apply() call reads: layout roles -> [E, ...] (resident, materialized layer) or [S, ...] (slot cache)."""
@@ -85,6 +94,7 @@ class ExpertView:
     slots: torch.Tensor | None = None
     n: int | None = None
     alphas: tuple[torch.Tensor, torch.Tensor] | None = None
+    routed: RoutedExperts | None = None
 
 
 def fused_piece(pieces: dict[str, torch.Tensor], role: str) -> torch.Tensor:

python/freetoken/layers/quantization/moe/nvfp4.py

@@ -74,7 +74,10 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
         banks = (t["gate_up"], t["gate_up_scale"], t["gate_up_global"], t["down"], t["down_scale"], t["down_global"])
         alpha, limit = float(layer.alpha), limit_or_inf(layer)
         if is_prefill:
-            return fused_experts_nvfp4(x, *banks, topk_weights, topk_ids, view.n, layer.activation, layer.apply_router_weight_on_input, alpha, limit)
+            return fused_experts_nvfp4(
+                x, *banks, topk_weights, topk_ids, view.n, layer.activation, layer.apply_router_weight_on_input, alpha, limit,
+                routed_experts=None if view.routed is None else view.routed.count,
+            )
         return fused_experts_decode_nvfp4_marlin(x, *banks, topk_weights, topk_ids, layer.activation, layer.apply_router_weight_on_input, alpha, limit)
 
 
@@ -121,10 +124,13 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
         from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4_w4a4
 
         t = view.tensors
-        assert view.alphas is not None
+        routed = view.routed
+        input_scales = view.alphas if routed is None else routed.alphas
+        assert input_scales is not None
         return fused_experts_nvfp4_w4a4(
             x, t["gate_up"], t["gate_up_scale"], t["gate_up_global"], t["down"], t["down_scale"], t["down_global"],
-            view.alphas[0], view.alphas[1], topk_weights, topk_ids, view.n, layer.activation, float(layer.alpha), limit_or_inf(layer),
+            input_scales[0], input_scales[1], topk_weights, topk_ids, view.n, layer.activation, float(layer.alpha), limit_or_inf(layer),
+            routed_experts=None if routed is None else routed.count,
         )
 
 

python/freetoken/moe/fused.py

@@ -44,6 +44,11 @@ def fused_topk(
     return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded)
 
 
+# sgl_kernel's moe_align_block_size counts experts in one 1024-thread block and reads out of
+# bounds past it; a prefill over the offload slot cache aligns against the slot count.
+_SGL_ALIGN_MAX_EXPERTS = 1024
+
+
 def moe_align_block_size(
     topk_ids: torch.Tensor, block_size: int, num_experts: int
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -97,7 +102,7 @@ def moe_align_block_size(
 
             return moe_align_block_size_triton(topk_ids, block_size, num_experts)
 
-    if not is_sgl_kernel_installed():
+    if not is_sgl_kernel_installed() or num_experts + 1 > _SGL_ALIGN_MAX_EXPERTS:
         from freetoken.kernel.triton.moe_align import (
             moe_align_block_size as triton_moe_align_block_size,
         )

python/freetoken/moe/fused_nvfp4.py

@@ -123,16 +123,20 @@ def fused_experts_nvfp4(
     apply_router_weight_on_input: bool = False,
     act_alpha: float = 1.702,
     act_limit: float = 7.0,
+    routed_experts: int | None = None,
 ) -> torch.Tensor:
     """Prefill inline-NVFP4 MoE. ``topk_ids`` index rows of the bank tensors in
     ``[0, num_experts)``: full-layer banks with position == expert id (the
-    materialized ``[:E]`` slot view or the overlap double buffer), raw ids."""
+    materialized ``[:E]`` slot view or the overlap double buffer), raw ids; or the
+    cache slots a demand prefill chunk routes to, whose layer's ``routed_experts``
+    size the tiles as the stream's banks would."""
     M, H = hidden_states.shape
     top_k = topk_ids.shape[1]
     two_i = gate_up_packed.shape[1]
     inter = two_i // 2
     dev, dt = hidden_states.device, hidden_states.dtype
-    cfg = nvfp4_moe.prefill_config(M, top_k, num_experts, hidden_states.device.index)
+    tile_experts = num_experts if routed_experts is None else routed_experts
+    cfg = nvfp4_moe.prefill_config(M, top_k, tile_experts, hidden_states.device.index)
 
     sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, cfg["BLOCK_SIZE_M"], num_experts)
     tw = topk_weights.reshape(-1).contiguous()
@@ -173,22 +177,26 @@ def fused_experts_nvfp4_w4a4(
     activation: str = "silu",
     act_alpha: float = 1.702,
     act_limit: float = 7.0,
+    routed_experts: int | None = None,
 ) -> torch.Tensor:
     """Prefill NVFP4 MoE with NVFP4 activations (W4A4) where the GPU and the batch favour it, else :func:`fused_experts_nvfp4`.
 
-    ``*_input_scale`` are the per-expert static activation scales of the ``num_experts`` banks; each GEMM input is quantized under their max, as vLLM does.
+    ``*_input_scale`` are the per-expert static activation scales of the experts the routes are drawn from: the
+    ``num_experts`` banks, or a demand prefill chunk's layer of ``routed_experts`` where the banks are cache slots.
+    Each GEMM input is quantized under their max, as vLLM does.
     """
     M, H = hidden_states.shape
     top_k = topk_ids.shape[1]
     inter = gate_up_packed.shape[1] // 2
-    cfg = nvfp4_moe.w4a4_prefill_config(M, top_k, num_experts, H, inter, activation, hidden_states.dtype, hidden_states.device.index)
+    tile_experts = num_experts if routed_experts is None else routed_experts
+    cfg = nvfp4_moe.w4a4_prefill_config(M, top_k, tile_experts, H, inter, activation, hidden_states.dtype, hidden_states.device.index)
     if cfg is None:
         return fused_experts_nvfp4(
             hidden_states, gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global,
-            topk_weights, topk_ids, num_experts, activation, False, act_alpha, act_limit,
+            topk_weights, topk_ids, num_experts, activation, False, act_alpha, act_limit, routed_experts,
         )
     dev = hidden_states.device
-    g = torch.stack((gate_up_input_scale[:num_experts], down_input_scale[:num_experts])).float().amax(dim=1)
+    g = torch.stack((gate_up_input_scale[:tile_experts], down_input_scale[:tile_experts])).float().amax(dim=1)
     sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, cfg["BLOCK_SIZE_M"], num_experts)
     tw = topk_weights.reshape(-1).contiguous()
     num_valid = topk_ids.numel()

python/freetoken/moe/offload_cache.py

@@ -124,6 +124,9 @@ class OffloadMoeCache:
     cache_size: int
     device: torch.device
     cache_policy: str = "lru"
+    # A prefill chunk of at most this many tokens admits the experts it routes to into the slot
+    # cache, as decode does, instead of streaming every expert of each layer; 0 streams them all.
+    prefill_demand_tokens: int = 0
     prefill_overlap: bool = False
     # Prefill hit/miss split: experts already resident in the slot cache (slots
     # >= 2 * num_experts) are gathered device-side into the double buffer instead
@@ -172,6 +175,10 @@ def __post_init__(self) -> None:
         policy_ids = {"lru": 0}
         assert self.cache_policy in policy_ids
         assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target
+        if self.prefill_demand_tokens and self.decode_target != "gpu":
+            raise ValueError(
+                f"--moe-prefill-demand-tokens needs GPU decode; this cache decodes on {self.decode_target!r}"
+            )
         if self.layer_layouts is not None:
             assert self.layout is None, "pass layout or layer_layouts, not both"
             assert len(self.layer_layouts) == self.num_layers, (len(self.layer_layouts), self.num_layers)
@@ -1067,6 +1074,33 @@ def release_prefill_layer(self, layer_id: int) -> None:
             self._prefill_buffer_has_release_event[buffer_id] = True
         self._prefill_buffer_released[buffer_id] = True
 
+    def loads_prefill_on_demand(self, num_tokens: int) -> bool:
+        return 0 < num_tokens <= self.prefill_demand_tokens
+
+    def load_prefill_experts(self, layer_id: int, topk_ids: torch.Tensor) -> torch.Tensor:
+        """Make a demand prefill chunk's routed experts resident, copying in the missing ones,
+        and return ``topk_ids`` as slot ids. The chunk's distinct experts go to the slot cache
+        ``ENSURE_QUERY_WIDTH`` at a time, each query's misses copied before the next query
+        evicts: LRU evicts an earlier query's slots last while the chunk's distinct experts fit the
+        layer's pool, which one admission of the whole chunk needs as well. The admissions stay out
+        of the decode statistics."""
+        from freetoken.moe.offload_kernels import ENSURE_QUERY_WIDTH, ensure_experts
+
+        experts = torch.arange(self.num_experts, dtype=torch.int32, device=self.device)
+        routed = torch.zeros(self.num_experts, dtype=torch.bool, device=self.device)
+        routed.index_fill_(0, topk_ids.reshape(-1).long(), True)
+        # the routed experts in ascending order, then the first of them again, which the admission
+        # dedups; a chunk's routes reach at most as many experts as there are routes
+        ordered = torch.where(routed, experts, experts + self.num_experts).sort().values
+        ordered = ordered[: min(self.num_experts, topk_ids.numel())]
+        query = torch.where(ordered < self.num_experts, ordered, ordered[0])
+        self._pending_src_layer = layer_id
+        self._pending_whole_layer = False
+        for piece in query.split(ENSURE_QUERY_WIDTH):
+            ensure_experts(self, layer_id, piece, record_stats=False)
+            self.copy_missing()
+        return self.slot_for_id[layer_id][topk_ids.long()]
+
     def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None:
         from freetoken.moe.offload_kernels import ensure_experts
 

python/freetoken/moe/offload_kernels.py

@@ -27,8 +27,12 @@
 # Phase 2 of the hybrid ensure kernel scans a pool wider than this in tiles of this width.
 _HYBRID_EVICT_CHUNK = 8192
 
+# lru_ensure dedups its query in a [K, K] register block: flat in cost up to 128 ids, spilling past
+# that and failing to compile past 1024, so a wider admission goes to it in queries this wide.
+ENSURE_QUERY_WIDTH = 128
 
-def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None:
+
+def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor, *, record_stats: bool = True) -> None:
     """Make this layer's routed experts resident; rewrite ``expert_ids`` to slot ids.
 
     Delegates to flashlib's slot cache. ``id_base`` maps this layer's expert ids into the
@@ -48,7 +52,7 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None:
         cache.src_indices,
         cache.evict_slots,
         cache.num_indices,
-        stats=cache.lru_stats[layer_id] if cache.collect_stats else None,
+        stats=cache.lru_stats[layer_id] if cache.collect_stats and record_stats else None,
         id_base=layer_id * cache.num_experts,
     )
 

python/freetoken/server/args.py

@@ -864,6 +864,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         ),
     )
 
+    parser.add_argument(
+        "--moe-prefill-demand-tokens",
+        type=int,
+        default=ServerArgs.moe_prefill_demand_tokens,
+        help=(
+            "A prefill chunk of at most this many tokens loads only the experts it routes to, "
+            "through the expert cache as decode does, instead of every expert of each layer. "
+            "0 streams whole layers for every chunk."
+        ),
+    )
+
     parser.add_argument(
         "--enable-special-token-ckpt",
         action="store_true",