From aae638d877f23f3e1f954cc8c9b7b2703bf85478 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Wed, 30 Sep 2026 14:07:09 +0300
Subject: [PATCH 1/6] test(moe): a short prefill chunk loads its experts
 through the slot cache, bit-identical to the stream

---
 tests/moe/test_nvfp4_backends.py | 35 ++++++++++++++++++++++++++
 tests/moe/test_offload.py        | 43 ++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py
index 0546fbcd6f..91abeee1a8 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -325,6 +325,41 @@ def test_triton_overlap_prefill_matches_dequant_reference():
         cache.release_prefill_layer(layer_id)
 
 
+@cuda
+def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit():
+    """A short prefill chunk loads only its routed experts, through the slot cache, and runs the
+    same grouped GEMM over the slots: its output must equal the whole-layer stream's exactly, and
+    its admissions must not count as decode misses."""
+    from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4
+
+    device = torch.device("cuda")
+    torch.manual_seed(15)
+    M = 12
+    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_cache, _ = _triton_cache(device)
+    streamed_cache.materialize_layer(1)
+    streamed_cache.copy_missing()
+    streamed = fused_experts_nvfp4(
+        hidden, *streamed_cache.bank_views(E), topk_weights, topk_ids.clone(), E, "silu", False
+    )
+
+    demand_cache, _ = _triton_cache(device, cache_size=2 * S)
+    demand_cache.collect_stats = True
+    slots = topk_ids.clone()
+    demand_cache.ensure_experts(1, slots, counts_as_decode=False)
+    demand_cache.copy_missing()
+    demanded = fused_experts_nvfp4(
+        hidden, *demand_cache.bank_views(), topk_weights, slots, demand_cache.cache_size, "silu", False
+    )
+
+    assert not torch.equal(slots, topk_ids), "the admission must have rewritten expert ids to slots"
+    assert torch.equal(demanded, streamed)
+    assert int(demand_cache.lru_stats.sum()) == 0
+
+
 @cuda
 def test_triton_swigluoai_matches_dequant_reference():
     """MiniMax-M3's swigluoai routed experts through the Triton prefill grouped GEMM
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index cd5f5dc773..5e3a20c8fb 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -125,6 +125,49 @@ def fake_fused(
     assert calls["topk_ids"].tolist() == [[2, 1]]
 
 
+def test_offload_moe_layer_short_prefill_loads_its_experts_through_the_slot_cache(monkeypatch):
+    layer, cache = _make_layer_and_cache()
+    cache.prefill_demand_tokens = 1
+    topk_weights = torch.tensor([[0.7, 0.3]], dtype=torch.float32)
+    hidden_states = torch.randn(1, 8)
+    calls = {}
+
+    monkeypatch.setattr(
+        "freetoken.layers.moe.fused_topk",
+        lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, torch.tensor([[2, 1]], dtype=torch.int32)),
+    )
+    monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.setdefault("materialized", layer_id))
+
+    def ensure(layer_id, expert_ids, *, counts_as_decode=True):
+        calls["ensured"] = (layer_id, expert_ids.tolist(), counts_as_decode)
+        expert_ids.copy_(torch.tensor([[5, 3]], dtype=torch.int32))
+
+    monkeypatch.setattr(cache, "ensure_experts", ensure)
+    monkeypatch.setattr(cache, "copy_missing", lambda: calls.setdefault("copied", True))
+
+    def fake_fused(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["w1"] = w1
+        calls["topk_ids"] = got_topk_ids.clone()
+        return hidden_states
+
+    monkeypatch.setattr("freetoken.moe.fused.fused_experts_impl", fake_fused)
+
+    out = layer.prefill_forward(hidden_states, torch.randn(1, 4))
+
+    assert out is hidden_states
+    assert "materialized" not in calls
+    assert calls["ensured"] == (0, [[2, 1]], False)
+    assert calls["copied"] is True
+    # the grouped GEMM reads the whole slot cache at the slots the admission wrote
+    assert calls["w1"].shape[0] == cache.cache_size
+    assert calls["topk_ids"].tolist() == [[5, 3]]
+
+    calls.clear()
+    layer.prefill_forward(torch.randn(2, 8), torch.randn(2, 4))
+    assert calls["materialized"] == 0 and "ensured" not in calls
+
+
 def test_offload_moe_layer_prefill_overlap_prefetches_layers_into_two_buffers(monkeypatch):
     from freetoken.moe.offload_cache import OffloadMoeCache
 

From 689f3d137e23aefbb1eecb961e59f4afd72e2a0c Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Wed, 30 Sep 2026 14:08:19 +0300
Subject: [PATCH 2/6] feat(moe): load a short prefill chunk's experts on demand
 through the slot cache

---
 docs/cli.md                             |  1 +
 python/freetoken/engine/config.py       |  1 +
 python/freetoken/engine/engine.py       |  2 ++
 python/freetoken/layers/moe.py          | 27 +++++++++++++++++++------
 python/freetoken/moe/offload_cache.py   | 16 ++++++++++++---
 python/freetoken/moe/offload_kernels.py |  4 ++--
 python/freetoken/server/args.py         | 11 ++++++++++
 7 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index 2166e9e86e..b4c318e424 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -89,6 +89,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does.
 | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor |
 | `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4) |
 | `--moe-hybrid-max-fetch` | auto | With `hybrid`: max experts fetched over PCIe per layer per step; rest computed on CPU |
+| `--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 |
 
diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index b6fa8b8949..4b30fe7fda 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -47,6 +47,7 @@ class EngineConfig:
     moe_cache_auto: bool = False
     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 >= 12.8
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 3988caf602..5e02fb2ed8 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -753,6 +753,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,
@@ -1459,6 +1460,7 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int, *, reserved: int
     "moe_cpu_layers": None,
     "moe_cpu_threads": 0,
     "moe_hybrid_max_fetch": -1,
+    "moe_prefill_demand_tokens": 0,
     "moe_prefill_overlap": True,
     "moe_prefill_hit_d2d": False,
     "expert_load": "auto",
diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 4918a032fe..710a44c79a 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -253,8 +253,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``.
     # ------------------------------------------------------------------
 
@@ -355,12 +356,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]):
+            cache.ensure_experts(self.layer_id, topk_ids, counts_as_decode=False)
+            cache.copy_missing()
+            return self._expert_gemm(
+                cache,
+                hidden_states,
+                topk_weights,
+                topk_ids,
+                views=cache.bank_views(),
+                n=cache.cache_size,
+                alphas=cache.alphas_for_slots(self.layer_id),
+                is_prefill=True,
+            )
         if cache.prefill_overlap:
             views = self._wait_prefill_overlap(cache)
             out = self._expert_gemm(
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 0969f73829..2afbb1a79f 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -108,6 +108,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
@@ -150,6 +153,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.layout is None:
             assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}"
         # Attached by the engine for decode_target == "cpu" (CpuMoeExecutor); None
@@ -843,17 +850,20 @@ 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 ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None:
+    def loads_prefill_on_demand(self, num_tokens: int) -> bool:
+        return num_tokens <= self.prefill_demand_tokens
+
+    def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor, *, counts_as_decode: bool = True) -> None:
         from freetoken.moe.offload_kernels import ensure_experts
 
-        if self.collect_decode_freq:
+        if self.collect_decode_freq and counts_as_decode:
             # ``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))
         self._pending_src_layer = layer_id
         self._pending_whole_layer = False
-        ensure_experts(self, layer_id, expert_ids)
+        ensure_experts(self, layer_id, expert_ids, record_stats=counts_as_decode)
 
     def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None:
         """Capped-fetch LRU for the hybrid backend.
diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py
index cf513f52d1..5df2ffec68 100644
--- a/python/freetoken/moe/offload_kernels.py
+++ b/python/freetoken/moe/offload_kernels.py
@@ -16,7 +16,7 @@
 )
 
 
-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 +35,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,
     )
 
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 7c0ed9b047..519447081b 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -781,6 +781,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",

From 96c731085aca69b74f5fae066bcdc806df540948 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Wed, 30 Sep 2026 17:38:27 +0300
Subject: [PATCH 3/6] test(moe): a dense model drops the demand-prefill
 threshold

---
 tests/moe/test_offload.py | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index 5e3a20c8fb..665cd7876e 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -498,6 +498,37 @@ def test_adjust_config_converts_moe_cache_rate_to_cache_size(monkeypatch):
     assert is_offload_moe_strategy(config.moe_strategy)
 
 
+def test_adjust_config_drops_the_demand_prefill_threshold_for_a_dense_model():
+    from types import SimpleNamespace
+
+    from freetoken.distributed import DistributedInfo
+    from freetoken.engine.config import EngineConfig
+    from freetoken.engine.engine import _adjust_config
+
+    config = EngineConfig(
+        model_path="/tmp/freetoken-test-model",
+        tp_info=DistributedInfo(rank=0, size=1),
+        dtype=torch.float16,
+        attention_backend="triton",
+        moe_prefill_demand_tokens=512,
+    )
+    object.__setattr__(
+        config,
+        "model_config",
+        SimpleNamespace(
+            has_swa_attention=False,
+            has_linear_attention=False,
+            is_moe=False,
+            num_layers=10,
+            expert_quant="none",
+        ),
+    )
+
+    _adjust_config(config)
+
+    assert config.moe_prefill_demand_tokens == 0
+
+
 def test_graph_capture_reuses_warm_offload_cache_before_capture(monkeypatch):
     import freetoken.core as core
     from freetoken.core import Context, Req, get_global_ctx

From a8f24b1764ccc54a7a1509be7f9a923ca23f1a24 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Thu, 1 Oct 2026 06:27:12 +0300
Subject: [PATCH 4/6] test(moe): the demand-prefill threshold is refused off
 GPU decode

OffloadMoeCache raises when --moe-prefill-demand-tokens is set for cpu or hybrid decode, whose experts are not read from the slot cache a demand-loaded chunk fills; the test drives both targets.
---
 tests/moe/test_offload.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index 665cd7876e..2cbd29eb29 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -283,6 +283,22 @@ def test_offload_moe_cache_marlin_rejects_slot_count_beyond_kernel_limit():
         )
 
 
+@pytest.mark.parametrize("decode_target", ["cpu", "hybrid"])
+def test_offload_moe_cache_rejects_demand_prefill_without_gpu_decode(decode_target):
+    # a demand-loaded chunk admits its experts into GPU decode's slot cache, which cpu and hybrid decode do not read
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    with pytest.raises(ValueError, match="needs GPU decode"):
+        OffloadMoeCache(
+            num_layers=1,
+            num_experts=4,
+            cache_size=6,
+            device=torch.device("cpu"),
+            decode_target=decode_target,
+            prefill_demand_tokens=64,
+        )
+
+
 def test_prefill_overlap_prefetch_invalidates_borrowed_unified_cache_slots():
     from freetoken.moe.offload_cache import OffloadMoeCache
 

From 4bea268f2dd4413759892b52332c4051c62e8fe4 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Tue, 13 Oct 2026 22:05:27 +0300
Subject: [PATCH 5/6] fix(moe): admit a demand prefill chunk's distinct experts
 128 at a time

lru_ensure dedups its query in a [K, K] block, so a chunk's tokens x top_k ids spilled past
128 and failed to compile past 1024. The chunk's distinct experts now go to it 128 at a time,
the width its dedup stays flat to, each query's misses copied before the next query evicts,
and a threshold of 0 never takes a chunk on demand.

Co-authored-by: Gabriel A. Devenyi <gdevenyi@gmail.com>
---
 python/freetoken/layers/moe.py          |  5 +-
 python/freetoken/moe/offload_cache.py   | 31 ++++++++--
 python/freetoken/moe/offload_kernels.py |  4 ++
 tests/moe/test_nvfp4_backends.py        | 27 ++++++---
 tests/moe/test_offload.py               | 76 +++++++++++++++++--------
 5 files changed, 102 insertions(+), 41 deletions(-)

diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 710a44c79a..061f280582 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -364,13 +364,12 @@ def _prefill_routed(
         cache = self.offload_cache
         assert cache is not None
         if cache.loads_prefill_on_demand(hidden_states.shape[0]):
-            cache.ensure_experts(self.layer_id, topk_ids, counts_as_decode=False)
-            cache.copy_missing()
+            slots = cache.load_prefill_experts(self.layer_id, topk_ids)
             return self._expert_gemm(
                 cache,
                 hidden_states,
                 topk_weights,
-                topk_ids,
+                slots,
                 views=cache.bank_views(),
                 n=cache.cache_size,
                 alphas=cache.alphas_for_slots(self.layer_id),
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 2afbb1a79f..3468e3d514 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -851,19 +851,42 @@ def release_prefill_layer(self, layer_id: int) -> None:
         self._prefill_buffer_released[buffer_id] = True
 
     def loads_prefill_on_demand(self, num_tokens: int) -> bool:
-        return num_tokens <= self.prefill_demand_tokens
+        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, and a layer never routes to more
+        experts than the cache holds. 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, *, counts_as_decode: bool = True) -> None:
+    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 and counts_as_decode:
+        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))
         self._pending_src_layer = layer_id
         self._pending_whole_layer = False
-        ensure_experts(self, layer_id, expert_ids, record_stats=counts_as_decode)
+        ensure_experts(self, layer_id, expert_ids)
 
     def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None:
         """Capped-fetch LRU for the hybrid backend.
diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py
index 5df2ffec68..8c1514ed92 100644
--- a/python/freetoken/moe/offload_kernels.py
+++ b/python/freetoken/moe/offload_kernels.py
@@ -15,6 +15,10 @@
     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
+
 
 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.
diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py
index 91abeee1a8..f02806f222 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -326,15 +326,21 @@ def test_triton_overlap_prefill_matches_dequant_reference():
 
 
 @cuda
-def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit():
+@pytest.mark.parametrize(("cache_size", "query_width"), [(1100, None), (E, 3)])
+def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit(monkeypatch, cache_size, query_width):
     """A short prefill chunk loads only its routed experts, through the slot cache, and runs the
     same grouped GEMM over the slots: its output must equal the whole-layer stream's exactly, and
-    its admissions must not count as decode misses."""
+    its admissions must not count as decode misses. The chunk routes more often than one
+    admission query compiles for, and the wider cache aligns past 1024 experts; the cache of one
+    layer's experts, filled three at a time, must keep every query's experts until the GEMM."""
+    from freetoken.moe import offload_kernels
     from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4
 
+    if query_width is not None:
+        monkeypatch.setattr(offload_kernels, "ENSURE_QUERY_WIDTH", query_width)
     device = torch.device("cuda")
     torch.manual_seed(15)
-    M = 12
+    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)
@@ -346,18 +352,21 @@ def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit():
         hidden, *streamed_cache.bank_views(E), topk_weights, topk_ids.clone(), E, "silu", False
     )
 
-    demand_cache, _ = _triton_cache(device, cache_size=2 * S)
+    demand_cache, _ = _triton_cache(device, cache_size=cache_size)
     demand_cache.collect_stats = True
-    slots = topk_ids.clone()
-    demand_cache.ensure_experts(1, slots, counts_as_decode=False)
-    demand_cache.copy_missing()
+    # decode 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])):
+        demand_cache.ensure_experts(layer_id, torch.tensor([resident], dtype=torch.int32, device=device))
+        demand_cache.copy_missing()
+    decode_stats = demand_cache.lru_stats.clone()
+    slots = demand_cache.load_prefill_experts(1, topk_ids)
     demanded = fused_experts_nvfp4(
         hidden, *demand_cache.bank_views(), topk_weights, slots, demand_cache.cache_size, "silu", False
     )
 
-    assert not torch.equal(slots, topk_ids), "the admission must have rewritten expert ids to slots"
+    assert not torch.equal(slots, topk_ids), "the admission must have mapped expert ids to slots"
     assert torch.equal(demanded, streamed)
-    assert int(demand_cache.lru_stats.sum()) == 0
+    assert torch.equal(demand_cache.lru_stats, decode_stats)
 
 
 @cuda
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index 2cbd29eb29..c29692c9fa 100644
--- a/tests/moe/test_offload.py
+++ b/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
 
@@ -126,24 +128,33 @@ def fake_fused(
 
 
 def test_offload_moe_layer_short_prefill_loads_its_experts_through_the_slot_cache(monkeypatch):
-    layer, cache = _make_layer_and_cache()
-    cache.prefill_demand_tokens = 1
-    topk_weights = torch.tensor([[0.7, 0.3]], dtype=torch.float32)
-    hidden_states = torch.randn(1, 8)
+    from freetoken.moe import offload_kernels
+
+    layer, cache = _make_layer_and_cache(num_experts=160, cache_size=192)
+    cache.prefill_demand_tokens = 96
+    cache.collect_decode_freq = True
+    # the first 96 tokens route 192 times, to all 160 experts
+    routes = torch.arange(194, dtype=torch.int32).remainder(160).view(97, 2)
+    events = []
     calls = {}
 
     monkeypatch.setattr(
         "freetoken.layers.moe.fused_topk",
-        lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, torch.tensor([[2, 1]], dtype=torch.int32)),
+        lambda *, hidden_states, gating_output, topk, renormalize: (
+            torch.rand(hidden_states.shape[0], 2), routes[: hidden_states.shape[0]].clone()
+        ),
     )
-    monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.setdefault("materialized", layer_id))
+    monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: events.append(("materialized", layer_id)))
+    monkeypatch.setattr(cache, "copy_missing", lambda: events.append(("copied",)))
 
-    def ensure(layer_id, expert_ids, *, counts_as_decode=True):
-        calls["ensured"] = (layer_id, expert_ids.tolist(), counts_as_decode)
-        expert_ids.copy_(torch.tensor([[5, 3]], dtype=torch.int32))
+    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
+        query = expert_ids.view(-1)
+        events.append(("admitted", query.tolist(), record_stats))
+        cache.slot_for_id[layer_id, query.long()] = cache.cache_size - 1 - query
+        query.copy_(cache.cache_size - 1 - query)
 
-    monkeypatch.setattr(cache, "ensure_experts", ensure)
-    monkeypatch.setattr(cache, "copy_missing", lambda: calls.setdefault("copied", True))
+    monkeypatch.setattr(offload_kernels, "ensure_experts", admit)
 
     def fake_fused(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation, apply_router_weight_on_input,
                    act_alpha=1.0, act_limit=float("inf")):
@@ -153,19 +164,34 @@ def fake_fused(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation
 
     monkeypatch.setattr("freetoken.moe.fused.fused_experts_impl", fake_fused)
 
-    out = layer.prefill_forward(hidden_states, torch.randn(1, 4))
-
+    hidden_states = torch.randn(96, 8)
+    out = layer.prefill_forward(hidden_states, torch.randn(96, 160))
+
+    admissions = [event for event in events if event[0] == "admitted"]
+    # lru_ensure takes at most 128 ids a query, never a chunk'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 not any(record_stats for *_, record_stats in admissions)
+    assert int(cache.decode_freq.sum()) == 0
     assert out is hidden_states
-    assert "materialized" not in calls
-    assert calls["ensured"] == (0, [[2, 1]], False)
-    assert calls["copied"] is True
     # the grouped GEMM reads the whole slot cache at the slots the admission wrote
     assert calls["w1"].shape[0] == cache.cache_size
-    assert calls["topk_ids"].tolist() == [[5, 3]]
+    assert torch.equal(calls["topk_ids"], cache.cache_size - 1 - routes[:96])
+
+    events.clear()
+    layer.prefill_forward(torch.randn(97, 8), torch.randn(97, 160))
+    assert events == [("materialized", 0), ("copied",)]
 
-    calls.clear()
-    layer.prefill_forward(torch.randn(2, 8), torch.randn(2, 4))
-    assert calls["materialized"] == 0 and "ensured" not in calls
+
+def test_offload_moe_cache_takes_a_prefill_chunk_on_demand_only_within_its_threshold():
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    cache = OffloadMoeCache(num_layers=1, num_experts=4, cache_size=6, device=torch.device("cpu"))
+    assert not cache.loads_prefill_on_demand(0)
+    cache.prefill_demand_tokens = 2
+    assert [cache.loads_prefill_on_demand(tokens) for tokens in range(4)] == [False, True, True, False]
 
 
 def test_offload_moe_layer_prefill_overlap_prefetches_layers_into_two_buffers(monkeypatch):

From 9ed4e79d5660bf0643205b791a85c04b12bb28b7 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Tue, 13 Oct 2026 22:05:28 +0300
Subject: [PATCH 6/6] fix(moe): align on triton past sgl_kernel's 1024-expert
 block

sgl_kernel's moe_align_block_size counts experts in one 1024-thread block and reads out of
bounds past it. A demand prefill chunk aligns against the slot count, so with sgl_kernel
installed the first one poisoned the CUDA context.

Co-authored-by: Gabriel A. Devenyi <gdevenyi@gmail.com>
---
 python/freetoken/moe/fused.py |  7 +++++-
 tests/moe/test_fused_moe.py   | 44 +++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py
index a90aa4c792..2f559f988e 100644
--- a/python/freetoken/moe/fused.py
+++ b/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]:
@@ -86,7 +91,7 @@ def moe_align_block_size(
     """
     from freetoken.kernel.backend import is_sgl_kernel_installed
 
-    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,
         )
diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py
index cafe778449..887121e9ff 100644
--- a/tests/moe/test_fused_moe.py
+++ b/tests/moe/test_fused_moe.py
@@ -369,3 +369,47 @@ def test_fused_topk_softmax_is_cuda_graph_capturable():
     torch.cuda.synchronize()
     assert (ids[:3] != -1).all()
     assert (ids[3:] == -1).all()
+
+
+def test_moe_align_leaves_sgl_kernel_past_its_one_block_of_experts(monkeypatch):
+    """sgl_kernel's align counts experts in one 1024-thread block, and every call hands it a
+    sentinel expert past the real ones: 1023 experts fit, 1024 align on Triton."""
+    import sys
+    import types
+
+    from freetoken.kernel import backend
+    from freetoken.kernel.triton import moe_align
+    from freetoken.moe.fused import moe_align_block_size
+
+    calls = []
+    sgl_kernel = types.ModuleType("sgl_kernel")
+    sgl_kernel.moe_align_block_size = lambda topk_ids, num_experts, *_: calls.append(("sgl_kernel", num_experts))
+    monkeypatch.setitem(sys.modules, "sgl_kernel", sgl_kernel)
+    monkeypatch.setattr(backend, "is_sgl_kernel_installed", lambda: True)
+    monkeypatch.setattr(
+        moe_align, "moe_align_block_size",
+        lambda topk_ids, block_size, num_experts: calls.append(("triton", num_experts)),
+    )
+
+    for num_experts in (1023, 1024, 2100):
+        moe_align_block_size(torch.zeros((4, 2), dtype=torch.int32), 16, num_experts)
+
+    assert calls == [("sgl_kernel", 1024), ("triton", 1024), ("triton", 2100)]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+@pytest.mark.parametrize("num_experts", [1023, 2100])
+def test_moe_align_sorts_every_route_into_its_expert_blocks_past_1024_experts(num_experts):
+    """A prefill over the offload slot cache aligns against its slot count."""
+    from freetoken.moe.fused import moe_align_block_size
+
+    torch.manual_seed(0)
+    ids = torch.randint(0, num_experts, (300, 10), device="cuda", dtype=torch.int32)
+    sorted_ids, expert_ids, padded = moe_align_block_size(ids, 16, num_experts)
+    n = int(padded)
+    flat = ids.flatten()
+    rows = sorted_ids[:n] < flat.numel()
+    real = sorted_ids[:n][rows]
+    assert sorted(real.tolist()) == list(range(flat.numel()))
+    block_of = torch.arange(n, device="cuda") // 16
+    assert torch.equal(flat[real.long()], expert_ids[block_of[rows]])