From d313c0d9a229de4c2c24c48403f43fbd544a9b46 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/9] 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 58a55bfa6..bb7754206 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -328,6 +328,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 9c7416914..c7bb38068 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 dee1c00335912920e5d401a7698334c7fc4cc567 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/9] 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 138683b01..2ce63b768 100644
--- a/docs/cli.md
+++ b/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 |
 
diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index 23afd3573..e8479084d 100644
--- a/python/freetoken/engine/config.py
+++ b/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
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 7982095ac..8ea48787f 100644
--- a/python/freetoken/engine/engine.py
+++ b/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",
diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 598bd1193..9442b75fa 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -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]):
+            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.pool_bank_views(self.layer_id),
+                n=cache.pool_slots(self.layer_id),
+                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 4a6a297df..132c43b4f 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/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,17 +1074,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, weights: torch.Tensor, clamp: bool = False
diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py
index e1fcbcda1..746943e21 100644
--- a/python/freetoken/moe/offload_kernels.py
+++ b/python/freetoken/moe/offload_kernels.py
@@ -28,7 +28,7 @@
 _HYBRID_EVICT_CHUNK = 8192
 
 
-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 +48,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 efe3b8c0b..b5f49a72b 100644
--- a/python/freetoken/server/args.py
+++ b/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",

From 98379dc847fcd690942f7da275cdcba78b1d34b5 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/9] 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 c7bb38068..09a1ce546 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 2645b8cc5077ee3a61767a39bfeff3e937fba3b7 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/9] 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 09a1ce546..ca4403f88 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 84c2d0489d581807f2a52c99c865748e193637cd 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/9] 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   | 32 +++++++++--
 python/freetoken/moe/offload_kernels.py |  4 ++
 tests/moe/test_nvfp4_backends.py        | 31 ++++++----
 tests/moe/test_offload.py               | 76 +++++++++++++++++--------
 5 files changed, 105 insertions(+), 43 deletions(-)

diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 9442b75fa..385adff59 100644
--- a/python/freetoken/layers/moe.py
+++ b/python/freetoken/layers/moe.py
@@ -357,13 +357,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.pool_bank_views(self.layer_id),
                 n=cache.pool_slots(self.layer_id),
                 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 132c43b4f..5fd779c4d 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -1075,19 +1075,43 @@ 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 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, *, 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, weights: torch.Tensor, clamp: bool = False
diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py
index 746943e21..1380cf158 100644
--- a/python/freetoken/moe/offload_kernels.py
+++ b/python/freetoken/moe/offload_kernels.py
@@ -27,6 +27,10 @@
 # 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, *, 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 bb7754206..bbb75228c 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -329,15 +329,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,21 +352,24 @@ def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit():
     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
+        hidden, *streamed_cache.stage_views(1, 0), 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
+        hidden, *demand_cache.pool_bank_views(1), topk_weights, slots, demand_cache.pool_slots(1), "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 ca4403f88..0726b1a4b 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])
 
-    calls.clear()
-    layer.prefill_forward(torch.randn(2, 8), torch.randn(2, 4))
-    assert calls["materialized"] == 0 and "ensured" not in calls
+    events.clear()
+    layer.prefill_forward(torch.randn(97, 8), torch.randn(97, 160))
+    assert events == [("materialized", 0), ("copied",)]
+
+
+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 1998f4ec7dbc48489c84ad3be17dba03e65e6c95 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/9] 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 5c34fb33c..65b39a595 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]:
@@ -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,
         )
diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py
index da34dba20..792a9b863 100644
--- a/tests/moe/test_fused_moe.py
+++ b/tests/moe/test_fused_moe.py
@@ -418,3 +418,47 @@ def test_fused_topk_sigmoid_matches_torch_reference(num_tokens, num_experts, top
     assert weights.dtype == torch.float32 and ids.dtype == torch.int32
     assert torch.equal(ids, ref_i)
     torch.testing.assert_close(weights, ref_w, rtol=1e-5, atol=1e-6)
+
+
+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]])

From 430dd2a28b8058a1f733ca211071e17185bcabef Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Thu, 22 Oct 2026 23:51:38 +0300
Subject: [PATCH 7/9] test(moe): a demand prefill chunk's nvfp4 gemms take
 their tiles and w4a4 scales from its layer

main's nvfp4 prefill sizes its tiles by tokens per expert, so over a 1100-slot pool the demand path
ran other tiles than the stream and its output stopped matching it; w4a4 took its activation scale
over every slot's input scale, other layers' included.
---
 tests/moe/test_nvfp4_backends.py | 52 ++++++++++++++++++++++++++++++++
 tests/moe/test_offload.py        | 10 ++++++
 2 files changed, 62 insertions(+)

diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py
index bbb75228c..367c823a2 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -372,6 +372,58 @@ def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit(monkeypatc
     assert torch.equal(demand_cache.lru_stats, decode_stats)
 
 
+def test_w4a4_prefill_over_cache_slots_takes_the_routed_layers_scales_and_tiles(monkeypatch):
+    """A demand prefill chunk's W4A4 GEMMs read cache slots, but the stream quantizes its activations under the max of
+    the layer's own input scales and sizes its tiles by the layer's experts: other layers' slots carry scales this
+    layer never routes to."""
+    from types import SimpleNamespace
+
+    from freetoken.layers.quantization import ExpertView, RoutedExperts
+    from freetoken.layers.quantization.moe.nvfp4 import W4a4Nvfp4MoEKernel
+    from freetoken.moe import fused_nvfp4
+
+    seen = {}
+    monkeypatch.setattr(fused_nvfp4, "fused_experts_nvfp4_w4a4", lambda *args, **kwargs: seen.update(args=args, kwargs=kwargs))
+    slots = 1100
+    layer_scales = (torch.full((E,), 2.0), torch.full((E,), 3.0))
+    slot_scales = (torch.full((slots,), 9.0), torch.full((slots,), 9.0))
+    banks = {role: torch.zeros(slots, 1) for role in ("gate_up", "gate_up_scale", "gate_up_global", "down", "down_scale", "down_global")}
+    view = ExpertView(banks, n=slots, alphas=slot_scales, routed=RoutedExperts(E, layer_scales))
+    W4a4Nvfp4MoEKernel().apply(
+        SimpleNamespace(activation="silu", alpha=1.702, limit=7.0),
+        torch.zeros(4, H), torch.ones(4, TOPK), torch.zeros(4, TOPK, dtype=torch.int32), view, is_prefill=True,
+    )
+
+    assert seen["args"][7] is layer_scales[0] and seen["args"][8] is layer_scales[1]
+    assert seen["args"][11] == slots
+    assert seen["kwargs"]["routed_experts"] == E
+
+
+def test_w4a4_prefill_sizes_its_tiles_by_the_routed_experts_and_hands_them_to_its_fallback(monkeypatch):
+    """Where W4A4 does not pay, the Triton W4A16 prefill takes over: it aligns over the slots and sizes its tiles
+    by the layer's experts, as W4A4 itself would have."""
+    from freetoken.kernel import nvfp4_moe
+    from freetoken.moe import fused_nvfp4
+
+    seen = {}
+
+    def without_w4a4(M, top_k, num_experts, *rest):
+        seen["tile_experts"] = num_experts
+        return None
+
+    monkeypatch.setattr(nvfp4_moe, "w4a4_prefill_config", without_w4a4)
+    monkeypatch.setattr(fused_nvfp4, "fused_experts_nvfp4", lambda *args: seen.update(fallback=args))
+    slots = 1100
+    banks = [torch.zeros(slots, 2 * I)] * 6
+    fused_nvfp4.fused_experts_nvfp4_w4a4(
+        torch.zeros(4, H), *banks, torch.ones(E), torch.ones(E), torch.ones(4, TOPK), torch.zeros(4, TOPK, dtype=torch.int32),
+        slots, routed_experts=E,
+    )
+
+    assert seen["tile_experts"] == E
+    assert seen["fallback"][9] == slots and seen["fallback"][-1] == E
+
+
 @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 0726b1a4b..99a4cf8ab 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -128,6 +128,7 @@ def fake_fused(
 
 
 def test_offload_moe_layer_short_prefill_loads_its_experts_through_the_slot_cache(monkeypatch):
+    from freetoken.layers.quantization import RoutedExperts
     from freetoken.moe import offload_kernels
 
     layer, cache = _make_layer_and_cache(num_experts=160, cache_size=192)
@@ -163,6 +164,13 @@ def fake_fused(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation
         return hidden_states
 
     monkeypatch.setattr("freetoken.moe.fused.fused_experts_impl", fake_fused)
+    apply = layer.quant_method.apply
+
+    def apply_and_keep_the_view(*args, **kwargs):
+        calls["view"] = args[3]
+        return apply(*args, **kwargs)
+
+    monkeypatch.setattr(layer.quant_method, "apply", apply_and_keep_the_view)
 
     hidden_states = torch.randn(96, 8)
     out = layer.prefill_forward(hidden_states, torch.randn(96, 160))
@@ -179,6 +187,8 @@ def fake_fused(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation
     # the grouped GEMM reads the whole slot cache at the slots the admission wrote
     assert calls["w1"].shape[0] == cache.cache_size
     assert torch.equal(calls["topk_ids"], cache.cache_size - 1 - routes[:96])
+    # and sizes itself by the layer's experts, which bf16 keeps no scales for
+    assert calls["view"].routed == RoutedExperts(160, None)
 
     events.clear()
     layer.prefill_forward(torch.randn(97, 8), torch.randn(97, 160))

From cf39859e579cf0c2cd2218f03a2196570619f033 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Thu, 22 Oct 2026 23:53:26 +0300
Subject: [PATCH 8/9] fix(moe): size a demand prefill chunk's nvfp4 tiles and
 w4a4 scales by its layer's experts

The expert view of a chunk loaded on demand carries the layer's expert count and input scales beside
its slot range: the triton and w4a4 prefill gemms align over the slots, size their tiles by the layer's
experts and quantize under the layer's scales, as the stream does.
---
 python/freetoken/layers/moe.py                |  6 +++--
 .../freetoken/layers/quantization/__init__.py |  3 ++-
 .../layers/quantization/moe/__init__.py       |  4 +--
 .../freetoken/layers/quantization/moe/base.py | 10 +++++++
 .../layers/quantization/moe/nvfp4.py          | 12 ++++++---
 python/freetoken/moe/fused_nvfp4.py           | 20 +++++++++-----
 tests/moe/test_nvfp4_backends.py              | 26 ++++++++++++++++++-
 7 files changed, 66 insertions(+), 15 deletions(-)

diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py
index 385adff59..ce0a6584c 100644
--- a/python/freetoken/layers/moe.py
+++ b/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
@@ -366,6 +366,7 @@ def _prefill_routed(
                 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:
@@ -424,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
diff --git a/python/freetoken/layers/quantization/__init__.py b/python/freetoken/layers/quantization/__init__.py
index 905b80ef5..6b65344d6 100644
--- a/python/freetoken/layers/quantization/__init__.py
+++ b/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",
diff --git a/python/freetoken/layers/quantization/moe/__init__.py b/python/freetoken/layers/quantization/moe/__init__.py
index 1230f36f0..5bd0a28ec 100644
--- a/python/freetoken/layers/quantization/moe/__init__.py
+++ b/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",
 ]
diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py
index b17a91359..6f76cf450 100644
--- a/python/freetoken/layers/quantization/moe/base.py
+++ b/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:
diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py
index 2a5e8bacd..0e31b61f3 100644
--- a/python/freetoken/layers/quantization/moe/nvfp4.py
+++ b/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,
         )
 
 
diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py
index 6bedd57b1..7afc04a9e 100644
--- a/python/freetoken/moe/fused_nvfp4.py
+++ b/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()
diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py
index 367c823a2..8fb9bfd8d 100644
--- a/tests/moe/test_nvfp4_backends.py
+++ b/tests/moe/test_nvfp4_backends.py
@@ -364,7 +364,8 @@ def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit(monkeypatc
     decode_stats = demand_cache.lru_stats.clone()
     slots = demand_cache.load_prefill_experts(1, topk_ids)
     demanded = fused_experts_nvfp4(
-        hidden, *demand_cache.pool_bank_views(1), topk_weights, slots, demand_cache.pool_slots(1), "silu", False
+        hidden, *demand_cache.pool_bank_views(1), topk_weights, slots, demand_cache.pool_slots(1), "silu", False,
+        routed_experts=E,
     )
 
     assert not torch.equal(slots, topk_ids), "the admission must have mapped expert ids to slots"
@@ -372,6 +373,29 @@ def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit(monkeypatc
     assert torch.equal(demand_cache.lru_stats, decode_stats)
 
 
+def test_triton_prefill_over_cache_slots_sizes_its_tiles_by_the_routed_layer(monkeypatch):
+    """The Triton prefill over a demand chunk's slots aligns over the slots and hands the GEMM the layer's experts to
+    size its tiles by, as the stream's banks would."""
+    from types import SimpleNamespace
+
+    from freetoken.layers.quantization import ExpertView, RoutedExperts
+    from freetoken.layers.quantization.moe.nvfp4 import TritonNvfp4MoEKernel
+    from freetoken.moe import fused_nvfp4
+
+    seen = {}
+    monkeypatch.setattr(fused_nvfp4, "fused_experts_nvfp4", lambda *args, **kwargs: seen.update(args=args, kwargs=kwargs))
+    slots = 1100
+    banks = {role: torch.zeros(slots, 1) for role in ("gate_up", "gate_up_scale", "gate_up_global", "down", "down_scale", "down_global")}
+    view = ExpertView(banks, n=slots, routed=RoutedExperts(E))
+    TritonNvfp4MoEKernel().apply(
+        SimpleNamespace(activation="silu", alpha=1.702, limit=7.0, apply_router_weight_on_input=False),
+        torch.zeros(4, H), torch.ones(4, TOPK), torch.zeros(4, TOPK, dtype=torch.int32), view, is_prefill=True,
+    )
+
+    assert seen["args"][9] == slots
+    assert seen["kwargs"]["routed_experts"] == E
+
+
 def test_w4a4_prefill_over_cache_slots_takes_the_routed_layers_scales_and_tiles(monkeypatch):
     """A demand prefill chunk's W4A4 GEMMs read cache slots, but the stream quantizes its activations under the max of
     the layer's own input scales and sizes its tiles by the layer's experts: other layers' slots carry scales this

From 5b2f0a08befb1044853a925fae920b4a20297121 Mon Sep 17 00:00:00 2001
From: Yevhenii <evgeny.kotyrlo@gmail.com>
Date: Fri, 23 Oct 2026 02:20:07 +0300
Subject: [PATCH 9/9] test(moe): drop two demand-prefill tests the tests README
 no longer admits

The dense-model test echoed _DENSE_MOE_SETTINGS back, and the align test's only
check was which kernel it called; the GPU test sorts real routes through
sgl_kernel at 2100 experts, which is where the out-of-bounds read shows.
---
 tests/moe/test_fused_moe.py | 26 --------------------------
 tests/moe/test_offload.py   | 31 -------------------------------
 2 files changed, 57 deletions(-)

diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py
index 792a9b863..85ba08054 100644
--- a/tests/moe/test_fused_moe.py
+++ b/tests/moe/test_fused_moe.py
@@ -420,32 +420,6 @@ def test_fused_topk_sigmoid_matches_torch_reference(num_tokens, num_experts, top
     torch.testing.assert_close(weights, ref_w, rtol=1e-5, atol=1e-6)
 
 
-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):
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index 99a4cf8ab..97eea6296 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -550,37 +550,6 @@ 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