diff --git a/docs/cli.md b/docs/cli.md index 138683b011..2ce63b768a 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 23afd35735..e8479084db 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 7982095ac4..8ea48787fe 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 598bd1193f..ce0a6584c3 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 @@ -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]): + slots = cache.load_prefill_experts(self.layer_id, topk_ids) + return self._expert_gemm( + cache, + hidden_states, + topk_weights, + slots, + 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: views = self._wait_prefill_overlap(cache) out = self._expert_gemm( @@ -410,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 905b80ef5d..6b65344d60 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 1230f36f08..5bd0a28ec1 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 b17a91359b..6f76cf4504 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 2a5e8bacd0..0e31b61f37 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.py b/python/freetoken/moe/fused.py index 5c34fb33ca..65b39a5957 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/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py index 6bedd57b12..7afc04a9e9 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/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 4a6a297df9..5fd779c4de 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,6 +1074,33 @@ 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 loads_prefill_on_demand(self, num_tokens: int) -> bool: + 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) -> None: from freetoken.moe.offload_kernels import ensure_experts diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index e1fcbcda1d..1380cf1587 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -27,8 +27,12 @@ # 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) -> 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 +52,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 efe3b8c0bd..b5f49a72b6 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", diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py index da34dba20d..85ba08054c 100644 --- a/tests/moe/test_fused_moe.py +++ b/tests/moe/test_fused_moe.py @@ -418,3 +418,21 @@ 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) + + +@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]]) diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py index 58a55bfa64..8fb9bfd8d7 100644 --- a/tests/moe/test_nvfp4_backends.py +++ b/tests/moe/test_nvfp4_backends.py @@ -328,6 +328,126 @@ def test_triton_overlap_prefill_matches_dequant_reference(): cache.release_prefill_layer(layer_id) +@cuda +@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. 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 = 600 + hidden = torch.randn(M, H, dtype=torch.bfloat16, device=device) / 4 + topk_ids = torch.rand(M, E, device=device).argsort(dim=1)[:, :TOPK].to(torch.int32).contiguous() + topk_weights = torch.rand(M, TOPK, dtype=torch.float32, device=device) + + streamed_cache, _ = _triton_cache(device) + streamed_cache.materialize_layer(1) + streamed_cache.copy_missing() + streamed = fused_experts_nvfp4( + hidden, *streamed_cache.stage_views(1, 0), topk_weights, topk_ids.clone(), E, "silu", False + ) + + demand_cache, _ = _triton_cache(device, cache_size=cache_size) + demand_cache.collect_stats = True + # 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.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" + assert torch.equal(demanded, streamed) + 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 + 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 9c74169149..97eea62965 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 @@ -125,6 +127,83 @@ 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): + from freetoken.layers.quantization import RoutedExperts + 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: ( + torch.rand(hidden_states.shape[0], 2), routes[: hidden_states.shape[0]].clone() + ), + ) + monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: events.append(("materialized", layer_id))) + monkeypatch.setattr(cache, "copy_missing", lambda: events.append(("copied",))) + + def admit(cache, layer_id, expert_ids, *, record_stats=True): + # the slot cache's contract: every queried expert resident, the query rewritten to its slots in place + 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(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")): + calls["w1"] = w1 + calls["topk_ids"] = got_topk_ids.clone() + 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)) + + 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 + # 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)) + 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 freetoken.moe.offload_cache import OffloadMoeCache @@ -240,6 +319,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