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

python/freetoken/layers/moe.py

@@ -264,10 +264,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
@@ -282,13 +282,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.bank_views(),
             n=None,
             alphas=cache.alphas_for_slots(self.layer_id),

python/freetoken/models/deepseek_v4/moe.py

@@ -119,21 +119,20 @@ 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)
         if (
             hidden_states.shape[0] * self.top_k >= self.num_experts
             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.bank_views(),
             n=None,
             alphas=cache.alphas_for_slots(self.layer_id),

python/freetoken/moe/offload_cache.py

@@ -224,8 +224,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 stay
-        # for the hybrid path, whose kernel is still ours.
+        # the kernel, which accumulates in the same launch; load_experts adds a wider step's
+        # counts itself, once per step. The stat_* tensors below stay for the hybrid path,
+        # whose kernel is still ours.
         self.lru_stats = torch.zeros(
             (self.num_layers, N_STATS), dtype=torch.int64, device=self.device
         )
@@ -843,18 +844,57 @@ 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, and a layer never routes to
+        more experts than the cache holds. 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) -> None:
         """Capped-fetch LRU for the hybrid backend.
 
@@ -905,7 +945,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

@@ -15,8 +15,15 @@
     os.getenv("FREETOKEN_HYBRID_FETCH", "recency").strip().lower() != "lowest_id"
 )
 
+# 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
@@ -35,7 +42,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

@@ -439,6 +439,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.bank_views(E), 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.bank_views(), 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.bank_views(), 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.bank_views(), 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

tests/moe/test_offload.py

@@ -23,18 +23,20 @@ def _bf16_offload_layer(layer_id: int, num_experts: int, top_k: int, hidden_size
     )
 
 
-def _make_layer_and_cache():
+def _make_layer_and_cache(num_experts: int = 4, cache_size: int = 6):
     from freetoken.moe.offload_cache import OffloadMoeCache
 
     _init_tp()
-    layer = _bf16_offload_layer(0, 4, 2, 8, 16)
+    layer = _bf16_offload_layer(0, num_experts, 2, 8, 16)
     cache = OffloadMoeCache(
         num_layers=1,
-        num_experts=4,
-        cache_size=6,
+        num_experts=num_experts,
+        cache_size=cache_size,
         device=torch.device("cpu"),
     )
-    cache.set_bank_sources({"gate_up": [torch.randn(4, 32, 8)], "down": [torch.randn(4, 8, 16)]})
+    cache.set_bank_sources(
+        {"gate_up": [torch.randn(num_experts, 32, 8)], "down": [torch.randn(num_experts, 8, 16)]}
+    )
     layer.offload_cache = cache
     return layer, cache
 
@@ -383,6 +385,101 @@ def fake_fused_decode(
     assert calls["topk_ids"].tolist() == [[5, 0]]
 
 
+def test_offload_moe_layer_wide_decode_step_loads_its_experts_128_at_a_time(monkeypatch):
+    from freetoken.moe import offload_kernels
+
+    layer, cache = _make_layer_and_cache(num_experts=160, cache_size=192)
+    cache.collect_decode_freq = True
+    cache.collect_stats = True
+    # 150 decode rows route 300 times, past one call's 256, to all 160 experts; the first 7 are resident already
+    routes = torch.arange(300, dtype=torch.int32).remainder(160).view(150, 2)
+    cache.slot_for_id[0, :7] = cache.cache_size - 1 - torch.arange(7, dtype=torch.int32)
+    events = []
+    calls = {}
+
+    monkeypatch.setattr(
+        "freetoken.layers.moe.fused_topk",
+        lambda *, hidden_states, gating_output, topk, renormalize: (
+            torch.rand(hidden_states.shape[0], 2), routes[: hidden_states.shape[0]].clone()
+        ),
+    )
+    monkeypatch.setattr(cache, "copy_missing", lambda: events.append(("copied",)))
+
+    def admit(cache, layer_id, expert_ids, *, record_stats=True):
+        # the slot cache's contract: every queried expert resident, the query rewritten to its slots in place,
+        # and num_indices the count of experts it had to copy in
+        query = expert_ids.view(-1)
+        events.append(("admitted", query.tolist(), record_stats))
+        cache.num_indices[0] = query[cache.slot_for_id[layer_id, query.long()] < 0].unique().numel()
+        cache.slot_for_id[layer_id, query.long()] = cache.cache_size - 1 - query
+        query.copy_(cache.cache_size - 1 - query)
+
+    monkeypatch.setattr(offload_kernels, "ensure_experts", admit)
+
+    def fake_fused_decode(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation,
+                          apply_router_weight_on_input, act_alpha=1.0, act_limit=float("inf")):
+        calls["topk_ids"] = got_topk_ids.clone()
+        return hidden_states
+
+    monkeypatch.setattr("freetoken.moe.fused.fused_experts_decode_impl", fake_fused_decode)
+
+    layer.decode_forward(torch.randn(150, 8), torch.randn(150, 160))
+
+    admissions = [event for event in events if event[0] == "admitted"]
+    # lru_ensure takes at most 128 ids a query, never a step's routes
+    assert [len(query) for _, query, _ in admissions] == [128, 32]
+    # each query's misses are copied before the next query can evict
+    assert events == [event for admission in admissions for event in (admission, ("copied",))]
+    assert sorted(expert for _, query, _ in admissions for expert in query) == list(range(160))
+    assert torch.equal(calls["topk_ids"], cache.cache_size - 1 - routes)
+    # counted as one step: 160 experts met, the 153 not resident missed, every route in the histogram once
+    assert not any(record_stats for *_, record_stats in admissions)
+    assert cache.lru_stats[0].tolist() == [160, 153, 1]
+    assert cache.decode_freq[0].tolist() == torch.bincount(routes.flatten().long(), minlength=160).tolist()
+
+    events.clear()
+    narrow = torch.arange(8, dtype=torch.int32).view(4, 2)
+    assert cache.load_experts(0, narrow) is narrow
+    assert events == [("admitted", list(range(8)), True), ("copied",)]
+
+
+def test_dsv4_offload_layer_wide_on_demand_chunk_loads_its_experts_128_at_a_time(monkeypatch):
+    from freetoken.layers.quantization import NoQuantConfig
+    from freetoken.models.deepseek_v4.args import DeepseekV4Args
+    from freetoken.models.deepseek_v4.moe import DSV4OffloadMoELayer
+    from freetoken.moe import offload_kernels
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    _init_tp()
+    args = DeepseekV4Args(dim=8, moe_inter_dim=16, n_routed_experts=640, n_activated_experts=2, swiglu_limit=None)
+    layer = DSV4OffloadMoELayer(0, args, quant_config=NoQuantConfig(), prefix="model.layers.0.mlp.experts")
+    cache = OffloadMoeCache(num_layers=1, num_experts=640, cache_size=768, device=torch.device("cpu"))
+    cache.set_bank_sources({"gate_up": [torch.randn(640, 32, 8)], "down": [torch.randn(640, 8, 16)]})
+    layer.offload_cache = cache
+    # 150 tokens route 300 times, past one call's 256 and under the 640 experts that would stream the layer
+    routes = torch.arange(300, dtype=torch.int32).view(150, 2)
+    queries = []
+    gemm = {}
+
+    def admit(cache, layer_id, expert_ids, *, record_stats=True):
+        query = expert_ids.view(-1)
+        queries.append(query.tolist())
+        cache.slot_for_id[layer_id, query.long()] = cache.cache_size - 1 - query
+        query.copy_(cache.cache_size - 1 - query)
+
+    monkeypatch.setattr(offload_kernels, "ensure_experts", admit)
+    monkeypatch.setattr(cache, "copy_missing", lambda: None)
+    monkeypatch.setattr(
+        layer, "_expert_gemm",
+        lambda cache, hidden_states, topk_weights, topk_ids, **_: gemm.setdefault("topk_ids", topk_ids.clone()),
+    )
+
+    layer._prefill_routed(torch.randn(150, 8), torch.rand(150, 2), routes.clone())
+
+    assert [len(query) for query in queries] == [128, 128, 44]
+    assert torch.equal(gemm["topk_ids"], cache.cache_size - 1 - routes)
+
+
 
 def test_lru_gpu_cache_assigns_unique_slots_for_large_miss_batch():
     import pytest