diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 4918a032f..e74b931fb 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -264,10 +264,10 @@ def _decode_routed( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor: - """On-demand load: ``ensure_experts`` rewrites ``topk_ids`` into cache slot - ids in place (loading missing experts), then the GEMM reads the full slot - cache. All device-side with fixed shapes, so the decode call is CUDA-graph - capturable. + """On-demand load: ``load_experts`` makes the routed experts resident and maps + ``topk_ids`` to their cache slots (in place when one admission takes the step), + then the GEMM reads the full slot cache. All device-side with fixed shapes, so + the decode call is CUDA-graph capturable. For ``decode_target == "cpu"`` the experts are instead computed on the CPU (high RAM bandwidth) straight from the host banks: ship hidden/routing to @@ -282,13 +282,12 @@ def _decode_routed( return executor.decode(self.layer_id, hidden_states, topk_weights, topk_ids) if cache.decode_target == "hybrid": return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids) - cache.ensure_experts(self.layer_id, topk_ids) - cache.copy_missing() + slots = cache.load_experts(self.layer_id, topk_ids) return self._expert_gemm( cache, hidden_states, topk_weights, - topk_ids, + slots, views=cache.bank_views(), n=None, alphas=cache.alphas_for_slots(self.layer_id), diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 35a316f10..fff0eab62 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -119,21 +119,20 @@ def _prefill_routed( # streaming buffers disown their borrowed slots on invalidation. cache = self.offload_cache assert cache is not None - # unpinned (LOCKED) layers must take the base materialize path: their copy_missing is the whole-layer pageable branch with position == expert id, which ensure_experts's LRU slot remap would contradict (the GEMM would gather other experts' weights) + # unpinned (LOCKED) layers must take the base materialize path: their copy_missing is the whole-layer pageable branch with position == expert id, which load_experts's LRU slot remap would contradict (the GEMM would gather other experts' weights) if ( hidden_states.shape[0] * self.top_k >= self.num_experts or cache.is_unpinned_layer(self.layer_id) ): return super()._prefill_routed(hidden_states, topk_weights, topk_ids) - cache.ensure_experts(self.layer_id, topk_ids) # in-place expert-id -> slot - cache.copy_missing() + slots = cache.load_experts(self.layer_id, topk_ids) if cache.collect_stats: cache.record_decode_stats(self.layer_id) return self._expert_gemm( cache, hidden_states, topk_weights, - topk_ids, + slots, views=cache.bank_views(), n=None, alphas=cache.alphas_for_slots(self.layer_id), diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 0969f7382..443d19945 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -224,8 +224,9 @@ def __post_init__(self) -> None: # is a one-off warm-up increment at capture time (<0.1% over a session). self.collect_stats = False # [num_layers, N_STATS] -- ensure_experts passes lru_stats[layer_id] straight to - # the kernel, which accumulates in the same launch. The stat_* tensors below stay - # for the hybrid path, whose kernel is still ours. + # the kernel, which accumulates in the same launch; load_experts adds a wider step's + # counts itself, once per step. The stat_* tensors below stay for the hybrid path, + # whose kernel is still ours. self.lru_stats = torch.zeros( (self.num_layers, N_STATS), dtype=torch.int64, device=self.device ) @@ -843,18 +844,57 @@ def release_prefill_layer(self, layer_id: int) -> None: self._prefill_buffer_has_release_event[buffer_id] = True self._prefill_buffer_released[buffer_id] = True + def load_experts(self, layer_id: int, topk_ids: torch.Tensor) -> torch.Tensor: + """Make the experts ``topk_ids`` routes to resident, copying in the missing ones, and return + ``topk_ids`` as slot ids. A step of at most ``ENSURE_ONE_CALL_WIDTH`` routes is one + admission that rewrites ``topk_ids`` in place. A wider one goes to the slot cache as its + distinct experts, ``ENSURE_QUERY_WIDTH`` at a time, each query's misses copied before the + next query evicts: LRU evicts an earlier query's slots last, and a layer never routes to + more experts than the cache holds. The statistics count either step once.""" + from freetoken.moe.offload_kernels import ENSURE_ONE_CALL_WIDTH, ENSURE_QUERY_WIDTH, ensure_experts + + if topk_ids.numel() <= ENSURE_ONE_CALL_WIDTH: + self.ensure_experts(layer_id, topk_ids) + self.copy_missing() + return topk_ids + self._count_routing(layer_id, topk_ids) + experts = torch.arange(self.num_experts, dtype=torch.int32, device=self.device) + routed = torch.zeros(self.num_experts, dtype=torch.bool, device=self.device) + routed.index_fill_(0, topk_ids.reshape(-1).long(), True) + # the routed experts in ascending order, then the first of them again, which the admission + # dedups; a step's routes reach at most as many experts as there are routes + ordered = torch.where(routed, experts, experts + self.num_experts).sort().values + ordered = ordered[: min(self.num_experts, topk_ids.numel())] + query = torch.where(ordered < self.num_experts, ordered, ordered[0]) + missed = torch.zeros((), dtype=torch.int64, device=self.device) + self._pending_src_layer = layer_id + self._pending_whole_layer = False + for piece in query.split(ENSURE_QUERY_WIDTH): + ensure_experts(self, layer_id, piece, record_stats=False) + missed += self.num_indices[0] + self.copy_missing() + if self.collect_stats: + stats = self.lru_stats[layer_id] + stats[Stat.ACTIVE] += routed.sum() + stats[Stat.MISS] += missed + stats[Stat.CALLS] += 1 + return self.slot_for_id[layer_id][topk_ids.long()] + def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: from freetoken.moe.offload_kernels import ensure_experts - if self.collect_decode_freq: - # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to - # slot ids in place), so snapshot the routing histogram before that happens. - ids = expert_ids.reshape(-1).long() - self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) + # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to slot ids in + # place), so snapshot the routing histogram before that happens. + self._count_routing(layer_id, expert_ids) self._pending_src_layer = layer_id self._pending_whole_layer = False ensure_experts(self, layer_id, expert_ids) + def _count_routing(self, layer_id: int, expert_ids: torch.Tensor) -> None: + if self.collect_decode_freq: + ids = expert_ids.reshape(-1).long() + self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) + def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None: """Capped-fetch LRU for the hybrid backend. @@ -905,7 +945,7 @@ def reset_stats(self) -> None: self.stat_steps_layer.zero_() def record_decode_stats(self, layer_id: int) -> None: - """No-op: ``ensure_experts`` accumulates into ``lru_stats`` inside its own launch. + """No-op: ``load_experts`` accumulates into ``lru_stats`` as it admits. Kept so the hybrid and non-hybrid call sites stay symmetric. The previous version was eight torch ops per layer per step, all captured into the decode graph. diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index cf513f52d..9c46f9a5a 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -15,8 +15,15 @@ os.getenv("FREETOKEN_HYBRID_FETCH", "recency").strip().lower() != "lowest_id" ) +# lru_ensure dedups its query in a [K, K] register block: flat in cost up to 128 ids, spilling past +# that and failing to compile past 1024, so a wider admission goes to it in queries this wide. +ENSURE_QUERY_WIDTH = 128 +# The widest decode step admitted in one call: up to 256 ids the spill costs less than the distinct experts' +# queries and their own launches, and past it far more (RTX 3090 Ti, CUDA-graph replay). +ENSURE_ONE_CALL_WIDTH = 256 -def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: + +def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor, *, record_stats: bool = True) -> None: """Make this layer's routed experts resident; rewrite ``expert_ids`` to slot ids. Delegates to flashlib's slot cache. ``id_base`` maps this layer's expert ids into the @@ -35,7 +42,7 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: cache.src_indices, cache.evict_slots, cache.num_indices, - stats=cache.lru_stats[layer_id] if cache.collect_stats else None, + stats=cache.lru_stats[layer_id] if cache.collect_stats and record_stats else None, id_base=layer_id * cache.num_experts, ) diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py index 0546fbcd6..763742cf9 100644 --- a/tests/moe/test_nvfp4_backends.py +++ b/tests/moe/test_nvfp4_backends.py @@ -439,6 +439,92 @@ def test_triton_decode_marlin_matches_baseline_kernel(): torch.testing.assert_close(marlin.float(), base.float(), rtol=2e-3, atol=2e-3) +def _streamed_decode(device, hidden, topk_weights, topk_ids): + """The decode GEMV over layer 1 materialized whole, where position == expert id.""" + from freetoken.moe.fused_nvfp4 import fused_experts_decode_nvfp4_marlin + + cache, _ = _triton_cache(device) + cache.materialize_layer(1) + cache.copy_missing() + return fused_experts_decode_nvfp4_marlin( + hidden, *cache.bank_views(E), topk_weights, topk_ids.clone(), "silu", False + ) + + +@cuda +@pytest.mark.parametrize(("cache_size", "query_width"), [(1100, None), (E, 3)]) +def test_triton_wide_decode_step_matches_the_streamed_layer_bit_for_bit(monkeypatch, cache_size, query_width): + """A decode step routing more often than one admission query compiles for loads its distinct + experts a query at a time: the GEMV over the slots must equal the one over the streamed layer + exactly, and the statistics must count the step once. The cache of one layer's experts, + filled three at a time, must keep every query's experts until the GEMV.""" + from freetoken.moe import offload_kernels + from freetoken.moe.fused_nvfp4 import fused_experts_decode_nvfp4_marlin + + if query_width is not None: + monkeypatch.setattr(offload_kernels, "ENSURE_QUERY_WIDTH", query_width) + device = torch.device("cuda") + torch.manual_seed(16) + M = 600 + hidden = torch.randn(M, H, dtype=torch.bfloat16, device=device) / 4 + topk_ids = torch.rand(M, E, device=device).argsort(dim=1)[:, :TOPK].to(torch.int32).contiguous() + topk_weights = torch.rand(M, TOPK, dtype=torch.float32, device=device) + streamed = _streamed_decode(device, hidden, topk_weights, topk_ids) + + cache, _ = _triton_cache(device, cache_size=cache_size) + cache.collect_stats = True + # earlier steps left the other layer's experts and one of this layer's resident, so slots are not expert ids + for layer_id, resident in ((0, [7, 6]), (1, [2])): + cache.load_experts(layer_id, torch.tensor([resident], dtype=torch.int32, device=device)) + stats = cache.lru_stats[1].clone() + slots = cache.load_experts(1, topk_ids) + decoded = fused_experts_decode_nvfp4_marlin( + hidden, *cache.bank_views(), topk_weights, slots, "silu", False + ) + + assert not torch.equal(slots, topk_ids), "the admission must have mapped expert ids to slots" + assert torch.equal(decoded, streamed) + # one step that met all E experts and missed every one but the resident expert 2 + assert (cache.lru_stats[1] - stats).tolist() == [E, E - 1, 1] + + +@cuda +def test_triton_wide_decode_step_replays_as_a_cuda_graph(monkeypatch): + """Decode runs as a CUDA graph, so a wide step's admission must capture: no host sync and fixed + shapes, and a replay on a cold cache admits and copies every expert itself.""" + from freetoken.moe import offload_kernels + from freetoken.moe.fused_nvfp4 import fused_experts_decode_nvfp4_marlin + + monkeypatch.setattr(offload_kernels, "ENSURE_ONE_CALL_WIDTH", 3) + monkeypatch.setattr(offload_kernels, "ENSURE_QUERY_WIDTH", 3) + device = torch.device("cuda") + torch.manual_seed(17) + M = 8 + hidden = torch.randn(M, H, dtype=torch.bfloat16, device=device) / 4 + topk_ids = torch.rand(M, E, device=device).argsort(dim=1)[:, :TOPK].to(torch.int32).contiguous() + topk_weights = torch.rand(M, TOPK, dtype=torch.float32, device=device) + streamed = _streamed_decode(device, hidden, topk_weights, topk_ids) + + cache, _ = _triton_cache(device, cache_size=2 * S) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + warm = cache.load_experts(1, topk_ids.clone()) + fused_experts_decode_nvfp4_marlin(hidden, *cache.bank_views(), topk_weights, warm, "silu", False) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + slots = cache.load_experts(1, topk_ids.clone()) + decoded = fused_experts_decode_nvfp4_marlin( + hidden, *cache.bank_views(), topk_weights, slots, "silu", False + ) + cache.reset() + graph.replay() + torch.cuda.synchronize() + + assert torch.equal(decoded, streamed) + + @cuda def test_b12x_decode_matches_dequant_reference(): """sm_120 + CUDA>=13 only: the flashinfer b12x W4A16 fused MoE over the slot cache diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index cd5f5dc77..575469f57 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 @@ -383,6 +385,101 @@ def fake_fused_decode( assert calls["topk_ids"].tolist() == [[5, 0]] +def test_offload_moe_layer_wide_decode_step_loads_its_experts_128_at_a_time(monkeypatch): + from freetoken.moe import offload_kernels + + layer, cache = _make_layer_and_cache(num_experts=160, cache_size=192) + cache.collect_decode_freq = True + cache.collect_stats = True + # 150 decode rows route 300 times, past one call's 256, to all 160 experts; the first 7 are resident already + routes = torch.arange(300, dtype=torch.int32).remainder(160).view(150, 2) + cache.slot_for_id[0, :7] = cache.cache_size - 1 - torch.arange(7, dtype=torch.int32) + events = [] + calls = {} + + monkeypatch.setattr( + "freetoken.layers.moe.fused_topk", + lambda *, hidden_states, gating_output, topk, renormalize: ( + torch.rand(hidden_states.shape[0], 2), routes[: hidden_states.shape[0]].clone() + ), + ) + monkeypatch.setattr(cache, "copy_missing", lambda: events.append(("copied",))) + + def admit(cache, layer_id, expert_ids, *, record_stats=True): + # the slot cache's contract: every queried expert resident, the query rewritten to its slots in place, + # and num_indices the count of experts it had to copy in + query = expert_ids.view(-1) + events.append(("admitted", query.tolist(), record_stats)) + cache.num_indices[0] = query[cache.slot_for_id[layer_id, query.long()] < 0].unique().numel() + cache.slot_for_id[layer_id, query.long()] = cache.cache_size - 1 - query + query.copy_(cache.cache_size - 1 - query) + + monkeypatch.setattr(offload_kernels, "ensure_experts", admit) + + def fake_fused_decode(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation, + apply_router_weight_on_input, act_alpha=1.0, act_limit=float("inf")): + calls["topk_ids"] = got_topk_ids.clone() + return hidden_states + + monkeypatch.setattr("freetoken.moe.fused.fused_experts_decode_impl", fake_fused_decode) + + layer.decode_forward(torch.randn(150, 8), torch.randn(150, 160)) + + admissions = [event for event in events if event[0] == "admitted"] + # lru_ensure takes at most 128 ids a query, never a step's routes + assert [len(query) for _, query, _ in admissions] == [128, 32] + # each query's misses are copied before the next query can evict + assert events == [event for admission in admissions for event in (admission, ("copied",))] + assert sorted(expert for _, query, _ in admissions for expert in query) == list(range(160)) + assert torch.equal(calls["topk_ids"], cache.cache_size - 1 - routes) + # counted as one step: 160 experts met, the 153 not resident missed, every route in the histogram once + assert not any(record_stats for *_, record_stats in admissions) + assert cache.lru_stats[0].tolist() == [160, 153, 1] + assert cache.decode_freq[0].tolist() == torch.bincount(routes.flatten().long(), minlength=160).tolist() + + events.clear() + narrow = torch.arange(8, dtype=torch.int32).view(4, 2) + assert cache.load_experts(0, narrow) is narrow + assert events == [("admitted", list(range(8)), True), ("copied",)] + + +def test_dsv4_offload_layer_wide_on_demand_chunk_loads_its_experts_128_at_a_time(monkeypatch): + from freetoken.layers.quantization import NoQuantConfig + from freetoken.models.deepseek_v4.args import DeepseekV4Args + from freetoken.models.deepseek_v4.moe import DSV4OffloadMoELayer + from freetoken.moe import offload_kernels + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + args = DeepseekV4Args(dim=8, moe_inter_dim=16, n_routed_experts=640, n_activated_experts=2, swiglu_limit=None) + layer = DSV4OffloadMoELayer(0, args, quant_config=NoQuantConfig(), prefix="model.layers.0.mlp.experts") + cache = OffloadMoeCache(num_layers=1, num_experts=640, cache_size=768, device=torch.device("cpu")) + cache.set_bank_sources({"gate_up": [torch.randn(640, 32, 8)], "down": [torch.randn(640, 8, 16)]}) + layer.offload_cache = cache + # 150 tokens route 300 times, past one call's 256 and under the 640 experts that would stream the layer + routes = torch.arange(300, dtype=torch.int32).view(150, 2) + queries = [] + gemm = {} + + def admit(cache, layer_id, expert_ids, *, record_stats=True): + query = expert_ids.view(-1) + queries.append(query.tolist()) + cache.slot_for_id[layer_id, query.long()] = cache.cache_size - 1 - query + query.copy_(cache.cache_size - 1 - query) + + monkeypatch.setattr(offload_kernels, "ensure_experts", admit) + monkeypatch.setattr(cache, "copy_missing", lambda: None) + monkeypatch.setattr( + layer, "_expert_gemm", + lambda cache, hidden_states, topk_weights, topk_ids, **_: gemm.setdefault("topk_ids", topk_ids.clone()), + ) + + layer._prefill_routed(torch.randn(150, 8), torch.rand(150, 2), routes.clone()) + + assert [len(query) for query in queries] == [128, 128, 44] + assert torch.equal(gemm["topk_ids"], cache.cache_size - 1 - routes) + + def test_lru_gpu_cache_assigns_unique_slots_for_large_miss_batch(): import pytest