fix(moe): admit a wide offloaded decode step's experts 128 at a time

python/freetoken/layers/moe.py

@@ -263,10 +263,10 @@ def _decode_routed(
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
     ) -> torch.Tensor:
-        """On-demand load: ``ensure_experts`` rewrites ``topk_ids`` into cache slot
-        ids in place (loading missing experts), then the GEMM reads the full slot
-        cache. All device-side with fixed shapes, so the decode call is CUDA-graph
-        capturable.
+        """On-demand load: ``load_experts`` makes the routed experts resident and maps
+        ``topk_ids`` to their cache slots (in place when one admission takes the step),
+        then the GEMM reads the full slot cache. All device-side with fixed shapes, so
+        the decode call is CUDA-graph capturable.
 
         For ``decode_target == "cpu"`` the experts are instead computed on the CPU
         (high RAM bandwidth) straight from the host banks: ship hidden/routing to
@@ -281,13 +281,12 @@ def _decode_routed(
             return executor.decode(self.layer_id, hidden_states, topk_weights, topk_ids)
         if cache.decode_target == "hybrid":
             return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids)
-        cache.ensure_experts(self.layer_id, topk_ids)
-        cache.copy_missing()
+        slots = cache.load_experts(self.layer_id, topk_ids)
         return self._expert_gemm(
             cache,
             hidden_states,
             topk_weights,
-            topk_ids,
+            slots,
             views=cache.pool_bank_views(self.layer_id),
             n=None,
             alphas=cache.alphas_for_slots(self.layer_id),

python/freetoken/models/deepseek_v4/moe.py

@@ -121,22 +121,21 @@ def _prefill_routed(
         # streaming buffers disown their borrowed slots on invalidation.
         cache = self.offload_cache
         assert cache is not None
-        # unpinned (LOCKED) layers must take the base materialize path: their copy_missing is the whole-layer pageable branch with position == expert id, which ensure_experts's LRU slot remap would contradict (the GEMM would gather other experts' weights)
+        # unpinned (LOCKED) layers must take the base materialize path: their copy_missing is the whole-layer pageable branch with position == expert id, which load_experts's LRU slot remap would contradict (the GEMM would gather other experts' weights)
         # one threshold for all layers (the smallest pool): with overlap prefill, the layers of a chunk must all stream or all take the slot path
         if (
             hidden_states.shape[0] * self.top_k >= min(self.num_experts, *cache.pool_sizes)
             or cache.is_unpinned_layer(self.layer_id)
         ):
             return super()._prefill_routed(hidden_states, topk_weights, topk_ids)
-        cache.ensure_experts(self.layer_id, topk_ids)  # in-place expert-id -> slot
-        cache.copy_missing()
+        slots = cache.load_experts(self.layer_id, topk_ids)
         if cache.collect_stats:
             cache.record_decode_stats(self.layer_id)
         return self._expert_gemm(
             cache,
             hidden_states,
             topk_weights,
-            topk_ids,
+            slots,
             views=cache.pool_bank_views(self.layer_id),
             n=None,
             alphas=cache.alphas_for_slots(self.layer_id),

python/freetoken/moe/offload_cache.py

@@ -273,8 +273,9 @@ def __post_init__(self) -> None:
         # is a one-off warm-up increment at capture time (<0.1% over a session).
         self.collect_stats = False
         # [num_layers, N_STATS] -- ensure_experts passes lru_stats[layer_id] straight to
-        # the kernel, which accumulates in the same launch. The stat_* tensors below are
-        # the hybrid path's: views of one buffer its kernel accumulates into the same way.
+        # the kernel, which accumulates in the same launch; load_experts adds a wider step's
+        # counts itself, once per step. The stat_* tensors below are the hybrid path's: views
+        # of one buffer its kernel accumulates into the same way.
         self.lru_stats = torch.zeros(
             (self.num_layers, N_STATS), dtype=torch.int64, device=self.device
         )
@@ -1067,18 +1068,58 @@ 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 load_experts(self, layer_id: int, topk_ids: torch.Tensor) -> torch.Tensor:
+        """Make the experts ``topk_ids`` routes to resident, copying in the missing ones, and return
+        ``topk_ids`` as slot ids. A step of at most ``ENSURE_ONE_CALL_WIDTH`` routes is one
+        admission that rewrites ``topk_ids`` in place. A wider one goes to the slot cache as its
+        distinct experts, ``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 step's distinct
+        experts fit the layer's pool, which one admission of the whole step needs as well. The
+        statistics count either step once."""
+        from freetoken.moe.offload_kernels import ENSURE_ONE_CALL_WIDTH, ENSURE_QUERY_WIDTH, ensure_experts
+
+        if topk_ids.numel() <= ENSURE_ONE_CALL_WIDTH:
+            self.ensure_experts(layer_id, topk_ids)
+            self.copy_missing()
+            return topk_ids
+        self._count_routing(layer_id, topk_ids)
+        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 step'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])
+        missed = torch.zeros((), dtype=torch.int64, device=self.device)
+        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)
+            missed += self.num_indices[0]
+            self.copy_missing()
+        if self.collect_stats:
+            stats = self.lru_stats[layer_id]
+            stats[Stat.ACTIVE] += routed.sum()
+            stats[Stat.MISS] += missed
+            stats[Stat.CALLS] += 1
+        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
 
-        if self.collect_decode_freq:
-            # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to
-            # slot ids in place), so snapshot the routing histogram before that happens.
-            ids = expert_ids.reshape(-1).long()
-            self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids))
+        # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to slot ids in
+        # place), so snapshot the routing histogram before that happens.
+        self._count_routing(layer_id, expert_ids)
         self._pending_src_layer = layer_id
         self._pending_whole_layer = False
         ensure_experts(self, layer_id, expert_ids)
 
+    def _count_routing(self, layer_id: int, expert_ids: torch.Tensor) -> None:
+        if self.collect_decode_freq:
+            ids = expert_ids.reshape(-1).long()
+            self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids))
+
     def ensure_experts_hybrid(
         self, layer_id: int, expert_ids: torch.Tensor, weights: torch.Tensor, clamp: bool = False
     ) -> HybridRoutes:
@@ -1136,7 +1177,7 @@ def reset_stats(self) -> None:
         self.stat_steps_layer.zero_()
 
     def record_decode_stats(self, layer_id: int) -> None:
-        """No-op: ``ensure_experts`` accumulates into ``lru_stats`` inside its own launch.
+        """No-op: ``load_experts`` accumulates into ``lru_stats`` as it admits.
 
         Kept so the hybrid and non-hybrid call sites stay symmetric. The previous version
         was eight torch ops per layer per step, all captured into the decode graph.

python/freetoken/moe/offload_kernels.py

@@ -27,8 +27,15 @@
 # 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
+# The widest decode step admitted in one call: up to 256 ids the spill costs less than the distinct experts'
+# queries and their own launches, and past it far more (RTX 3090 Ti, CUDA-graph replay).
+ENSURE_ONE_CALL_WIDTH = 256
 
-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 +55,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,
     )
 

tests/moe/test_nvfp4_backends.py

@@ -582,6 +582,92 @@ def test_triton_decode_marlin_matches_baseline_kernel():
     torch.testing.assert_close(marlin.float(), base.float(), rtol=2e-3, atol=2e-3)
 
 
+def _streamed_decode(device, hidden, topk_weights, topk_ids):
+    """The decode GEMV over layer 1 materialized whole, where position == expert id."""
+    from freetoken.moe.fused_nvfp4 import fused_experts_decode_nvfp4_marlin
+
+    cache, _ = _triton_cache(device)
+    cache.materialize_layer(1)
+    cache.copy_missing()
+    return fused_experts_decode_nvfp4_marlin(
+        hidden, *cache.stage_views(1, 0), topk_weights, topk_ids.clone(), "silu", False
+    )
+
+
+@cuda
+@pytest.mark.parametrize(("cache_size", "query_width"), [(1100, None), (E, 3)])
+def test_triton_wide_decode_step_matches_the_streamed_layer_bit_for_bit(monkeypatch, cache_size, query_width):
+    """A decode step routing more often than one admission query compiles for loads its distinct
+    experts a query at a time: the GEMV over the slots must equal the one over the streamed layer
+    exactly, and the statistics must count the step once. The cache of one layer's experts,
+    filled three at a time, must keep every query's experts until the GEMV."""
+    from freetoken.moe import offload_kernels
+    from freetoken.moe.fused_nvfp4 import fused_experts_decode_nvfp4_marlin
+
+    if query_width is not None:
+        monkeypatch.setattr(offload_kernels, "ENSURE_QUERY_WIDTH", query_width)
+    device = torch.device("cuda")
+    torch.manual_seed(16)
+    M = 600
+    hidden = torch.randn(M, H, dtype=torch.bfloat16, device=device) / 4
+    topk_ids = torch.rand(M, E, device=device).argsort(dim=1)[:, :TOPK].to(torch.int32).contiguous()
+    topk_weights = torch.rand(M, TOPK, dtype=torch.float32, device=device)
+    streamed = _streamed_decode(device, hidden, topk_weights, topk_ids)
+
+    cache, _ = _triton_cache(device, cache_size=cache_size)
+    cache.collect_stats = True
+    # earlier steps left the other layer's experts and one of this layer's resident, so slots are not expert ids
+    for layer_id, resident in ((0, [7, 6]), (1, [2])):
+        cache.load_experts(layer_id, torch.tensor([resident], dtype=torch.int32, device=device))
+    stats = cache.lru_stats[1].clone()
+    slots = cache.load_experts(1, topk_ids)
+    decoded = fused_experts_decode_nvfp4_marlin(
+        hidden, *cache.pool_bank_views(1), topk_weights, slots, "silu", False
+    )
+
+    assert not torch.equal(slots, topk_ids), "the admission must have mapped expert ids to slots"
+    assert torch.equal(decoded, streamed)
+    # one step that met all E experts and missed every one but the resident expert 2
+    assert (cache.lru_stats[1] - stats).tolist() == [E, E - 1, 1]
+
+
+@cuda
+def test_triton_wide_decode_step_replays_as_a_cuda_graph(monkeypatch):
+    """Decode runs as a CUDA graph, so a wide step's admission must capture: no host sync and fixed
+    shapes, and a replay on a cold cache admits and copies every expert itself."""
+    from freetoken.moe import offload_kernels
+    from freetoken.moe.fused_nvfp4 import fused_experts_decode_nvfp4_marlin
+
+    monkeypatch.setattr(offload_kernels, "ENSURE_ONE_CALL_WIDTH", 3)
+    monkeypatch.setattr(offload_kernels, "ENSURE_QUERY_WIDTH", 3)
+    device = torch.device("cuda")
+    torch.manual_seed(17)
+    M = 8
+    hidden = torch.randn(M, H, dtype=torch.bfloat16, device=device) / 4
+    topk_ids = torch.rand(M, E, device=device).argsort(dim=1)[:, :TOPK].to(torch.int32).contiguous()
+    topk_weights = torch.rand(M, TOPK, dtype=torch.float32, device=device)
+    streamed = _streamed_decode(device, hidden, topk_weights, topk_ids)
+
+    cache, _ = _triton_cache(device, cache_size=2 * S)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        warm = cache.load_experts(1, topk_ids.clone())
+        fused_experts_decode_nvfp4_marlin(hidden, *cache.pool_bank_views(1), topk_weights, warm, "silu", False)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        slots = cache.load_experts(1, topk_ids.clone())
+        decoded = fused_experts_decode_nvfp4_marlin(
+            hidden, *cache.pool_bank_views(1), topk_weights, slots, "silu", False
+        )
+    cache.reset()
+    graph.replay()
+    torch.cuda.synchronize()
+
+    assert torch.equal(decoded, streamed)
+
+
 @cuda
 def test_b12x_decode_matches_dequant_reference():
     """sm_120 + CUDA>=13 only: the flashinfer b12x W4A16 fused MoE over the slot cache