5080deploy: opt-in on-demand experts for short prefills (FREETOKEN_MOE_SMALL_PREFILL_TOKENS, off)

python/freetoken/layers/moe.py

@@ -24,6 +24,12 @@
 # default. Set FREETOKEN_HYBRID_OVERLAP=0 to force the serial path (CPU sync before the
 # GPU work) -- a measurement-only escape hatch to A/B the overlap benefit.
 _HYBRID_OVERLAP = os.getenv("FREETOKEN_HYBRID_OVERLAP", "1") != "0"
+# Prefills of at most this many tokens load only their routed experts (the decode path's
+# on-demand LRU) instead of streaming every expert of every layer; 0 = off.
+_SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0"))
+# ids per LRU ensure in that path: the kernel's block is next_pow2(ids) x next_pow2(slots), and a
+# block much wider than decode's spills to local memory whose launch reservation can OOM
+_ENSURE_CHUNK_IDS = 32
 
 
 class MoELayer(BaseOP):
@@ -356,6 +362,22 @@ def _prefill_routed(
         pass through unmapped."""
         cache = self.offload_cache
         assert cache is not None
+        if 0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS and cache.decode_target in ("gpu", "hybrid"):
+            # A short extension (an agent turn over a cached prefix) touches a fraction of each
+            # layer's experts; fetching those beats streaming all of them (~68 GB here). The same
+            # token count holds in every layer, so the overlap double buffer is never begun.
+            slots = self._ensure_unique(cache, topk_ids)
+            if slots is not None:
+                return self._expert_gemm(
+                    cache,
+                    hidden_states,
+                    topk_weights,
+                    slots,
+                    views=cache.bank_views(),
+                    n=None,
+                    alphas=cache.alphas_for_slots(self.layer_id),
+                    is_prefill=False,
+                )
         if cache.prefill_overlap:
             views = self._wait_prefill_overlap(cache)
             out = self._expert_gemm(
@@ -383,6 +405,25 @@ def _prefill_routed(
             is_prefill=True,
         )
 
+    def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor | None:
+        """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots, or
+        None (take the whole-layer path) when they cannot all be resident at once.
+
+        The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a
+        prefill's tokens x top_k ids overflow Triton's element limit, and even a block under it
+        can fail to launch on a full GPU. It gets the unique ids (at most num_experts) in
+        decode-width chunks; each chunk's misses are copied before the next ensure reuses the plan."""
+        uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True)
+        n = uniq.numel()
+        if n > cache.cache_size // 2:  # leave room so a chunk never evicts the one before it
+            return None
+        uniq = uniq.to(torch.int32)
+        for start in range(0, n, _ENSURE_CHUNK_IDS):
+            part = uniq[start : start + _ENSURE_CHUNK_IDS]  # a contiguous view: ensure rewrites it in place
+            cache.ensure_experts(self.layer_id, part)
+            cache.copy_missing()
+        return uniq[inverse].view_as(topk_ids)
+
     def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, ...]:
         """Double-buffer choreography for this layer's overlap prefill: kick off the
         next layer's full-layer H2D copy, then return this layer's bank views (in

tests/moe/test_small_prefill_on_demand.py

@@ -0,0 +1,102 @@
+"""Short prefills load only their routed experts (FREETOKEN_MOE_SMALL_PREFILL_TOKENS) instead of
+streaming whole layers; above the threshold, or with it off, the whole-layer path is unchanged."""
+
+import pytest
+import torch
+
+import freetoken.layers.moe as moe_mod
+from freetoken.distributed import set_tp_info, try_get_tp_info
+
+
+def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None):
+    from freetoken.layers.moe import OffloadMoELayer
+    from freetoken.layers.quantization import NoQuantConfig
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(),
+                            prefix="model.layers.0.mlp.experts")
+    cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=cache_size or num_experts, device=device)
+    g = torch.Generator().manual_seed(0)
+    cache.set_bank_sources({
+        "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1],
+        "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1],
+    })
+    layer.offload_cache = cache
+    return layer, cache
+
+
+@pytest.mark.parametrize("tokens, threshold, on_demand", [
+    (3, 8, True), (8, 8, True), (9, 8, False), (3, 0, False),
+])
+def test_small_prefill_dispatch(monkeypatch, tokens, threshold, on_demand):
+    layer, cache = _layer_and_cache(torch.device("cpu"))
+    monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
+    calls = []
+    monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids: calls.append("ensure"))
+    monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append("materialize"))
+    monkeypatch.setattr(cache, "copy_missing", lambda: calls.append("copy"))
+    monkeypatch.setattr(cache, "begin_prefill", lambda: calls.append("begin_prefill"), raising=False)
+    got = {}
+
+    def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill):
+        got.update(n=n, is_prefill=is_prefill)
+        return hs
+
+    monkeypatch.setattr(layer, "_expert_gemm", fake_gemm)
+    hs = torch.randn(tokens, 16)
+    layer._prefill_routed(hs, torch.full((tokens, 2), 0.5), torch.zeros(tokens, 2, dtype=torch.int32))
+    if on_demand:
+        assert calls == ["ensure", "copy"] and got == {"n": None, "is_prefill": False}
+    else:
+        assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_small_prefill_matches_the_whole_layer_path(monkeypatch):
+    from freetoken.core import get_global_ctx  # noqa: F401  (layer reads the ctx only in forward())
+
+    dev = torch.device("cuda")
+    hs = torch.randn(5, 16, dtype=torch.bfloat16, device=dev) * 0.5
+    ids = torch.tensor([[0, 3], [1, 3], [7, 2], [3, 4], [5, 0]], dtype=torch.int32, device=dev)
+    w = torch.tensor([[0.6, 0.4], [0.5, 0.5], [0.9, 0.1], [0.3, 0.7], [0.2, 0.8]], device=dev)
+    outs = {}
+    for threshold in (0, 8):
+        monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
+        layer, cache = _layer_and_cache(dev, cache_size=32)
+        outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0)
+    torch.testing.assert_close(outs[8], outs[0], rtol=2e-2, atol=2e-2)
+
+
+def _run_and_check_path(layer, hs, w, ids, *, on_demand):
+    taken = []
+    real = getattr(layer, "_ensure_unique", None)
+    if real is not None:
+        layer._ensure_unique = lambda cache, t: taken.append(r := real(cache, t)) or r
+    out = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float()
+    assert (bool(taken) and taken[0] is not None) == on_demand
+    return out
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch):
+    # 200 tokens x top-10 = 2000 routed ids against a cache that rounds up to 2048 slots: the
+    # LRU kernel's block is next_pow2(ids) x next_pow2(slots), so the raw ids overflow Triton's
+    # 1M-element limit, and a 512-id block failed to launch on a full GPU (CUDA OOM). The path
+    # must hand it the unique ids in decode-width chunks.
+    dev = torch.device("cuda")
+    tokens, experts, top_k = 200, 512, 10
+    g = torch.Generator(device=dev).manual_seed(1)
+    hs = torch.randn(tokens, 16, dtype=torch.bfloat16, device=dev, generator=g) * 0.5
+    ids = torch.randint(0, experts, (tokens, top_k), dtype=torch.int32, device=dev, generator=g)
+    w = torch.softmax(torch.randn(tokens, top_k, device=dev, generator=g), dim=-1)
+    outs = {}
+    for threshold in (0, 1024):
+        monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold)
+        layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100)
+        widths, ensure = [], cache.ensure_experts
+        monkeypatch.setattr(cache, "ensure_experts", lambda lid, q: (widths.append(q.numel()), ensure(lid, q))[1])
+        outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0)
+    assert widths and max(widths) <= 32, f"LRU ensure widths {sorted(set(widths))}: wider than decode's"
+    torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2)