perf(moe): split hybrid decode routes inside the ensure kernel

python/freetoken/layers/moe.py

@@ -307,22 +307,18 @@ def _decode_hybrid(
 
         The CPU pool is kicked off (``decode_submit``) before the GPU PCIe fetch + GEMM so
         the CPU overflow GEMV runs concurrently with the GPU work. Capture-safe: the
-        routing split is device-side elementwise and the CPU submit/sync are host nodes.
+        routing split runs inside the hybrid ensure kernel and the CPU submit/sync are host nodes.
         Each route is computed exactly once -- CPU-assigned routes reach the GPU as slot -1 with
         weight 0 (a kernel with ``supports_inactive_slots`` skips them without a read; any other reads
         slot 0, whose storage is a payload or the bank's zero fill) and the CPU ids are -1 for
         GPU-assigned routes (the C++ kernel skips id<0).
         """
         executor = cache.cpu_executor
         assert executor is not None, "CPU MoE executor was not initialized"
-        raw_ids = topk_ids.clone()
-        cache.ensure_experts_hybrid(self.layer_id, topk_ids)
-        on_gpu = topk_ids >= 0
-        cpu_ids = torch.where(on_gpu, -1, raw_ids).contiguous()
-        gpu_w = torch.where(on_gpu, topk_weights, 0.0).contiguous()
-        gpu_slots = topk_ids
-        if cache.collect_stats:
-            cache.record_decode_stats_hybrid(self.layer_id)
+        # the kernel reads every route's slot: give the inactive ones slot 0, which holds a
+        # complete payload or the bank's zero fill (never uninitialized storage), times weight 0
+        clamp = not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots)
+        gpu_slots, gpu_w, cpu_ids = cache.ensure_experts_hybrid(self.layer_id, topk_ids, topk_weights, clamp=clamp)
         pending = executor.decode_submit(self.layer_id, hidden_states, topk_weights, cpu_ids)
 
         # Measurement knob: FREETOKEN_HYBRID_OVERLAP=0 syncs the CPU pool *before* the
@@ -332,10 +328,6 @@ def _decode_hybrid(
         )
 
         cache.copy_missing()
-        if not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots):
-            # the kernel reads every route's slot: give the inactive ones slot 0, which holds a
-            # complete payload or the bank's zero fill (never uninitialized storage), times weight 0
-            gpu_slots = gpu_slots.clamp_min(0)
         gpu_routed = self._expert_gemm(
             cache,
             hidden_states,

python/freetoken/moe/offload_cache.py

@@ -3,7 +3,7 @@
 import math
 import os
 from dataclasses import dataclass
-from typing import Iterator
+from typing import TYPE_CHECKING, Iterator
 
 import torch
 from flashlib.kernels.slot_cache import N_STATS, Stat
@@ -28,6 +28,9 @@
 from freetoken.kernel.fast_index_copy import default_blocks_per_bank
 from freetoken.utils import init_logger
 
+if TYPE_CHECKING:
+    from freetoken.moe.offload_kernels import HybridRoutes
+
 logger = init_logger(__name__)
 
 # quant_format -> bank names, in registration order: the single place a format's bank
@@ -224,24 +227,23 @@ 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. The stat_* tensors below are
+        # the hybrid path's: views of one buffer its kernel accumulates into the same way.
         self.lru_stats = torch.zeros(
             (self.num_layers, N_STATS), dtype=torch.int64, device=self.device
         )
-        self.stat_missing = torch.zeros((), dtype=torch.int64, device=self.device)
-        self.stat_active = torch.zeros((), dtype=torch.int64, device=self.device)
-        self.stat_calls = torch.zeros((), dtype=torch.int64, device=self.device)
-        # hybrid only: experts actually fetched over PCIe (<= stat_missing). The CPU
-        # computes stat_missing - stat_fetched of them.
-        self.stat_fetched = torch.zeros((), dtype=torch.int64, device=self.device)
+        # (missing, fetched, active, calls), then the same per layer
+        stat_buf = torch.zeros(4 * (self.num_layers + 1), dtype=torch.int64, device=self.device)
+        self.hybrid_stats = (stat_buf[:4], stat_buf[4:])
+        # hybrid only: stat_fetched counts the experts actually fetched over PCIe (<= stat_missing).
+        # The CPU computes stat_missing - stat_fetched of them.
+        self.stat_missing, self.stat_fetched, self.stat_active, self.stat_calls = stat_buf[:4].unbind()
         # Per-layer counterparts of the scalars above (indexed by MoE-layer id). Same
         # device-side accumulation (graph-safe: layer_id is a static index per graph node),
         # so one req's per-layer miss rate is readable via decode_miss_stats_per_layer().
-        self.stat_missing_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device)
-        self.stat_active_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device)
-        self.stat_fetched_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device)
-        self.stat_steps_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device)
+        per_layer = stat_buf[4:].view(self.num_layers, 4)
+        self.stat_missing_layer, self.stat_fetched_layer = per_layer[:, 0], per_layer[:, 1]
+        self.stat_active_layer, self.stat_steps_layer = per_layer[:, 2], per_layer[:, 3]
         # Opt-in decode routing histogram (per layer, per expert) for cache-skew
         # analysis. Accumulated in ``ensure_experts`` from the raw expert ids before the
         # kernel rewrites them to slots. Only accurate with CUDA graphs disabled (the
@@ -855,7 +857,9 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None:
         self._pending_whole_layer = False
         ensure_experts(self, layer_id, expert_ids)
 
-    def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None:
+    def ensure_experts_hybrid(
+        self, layer_id: int, expert_ids: torch.Tensor, weights: torch.Tensor, clamp: bool = False
+    ) -> HybridRoutes:
         """Capped-fetch LRU for the hybrid backend.
 
         Like :meth:`ensure_experts` but assigns slots to (and schedules copies for) at
@@ -864,16 +868,17 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None
         left non-resident and ``expert_ids`` is rewritten to their cache slot (hit or
         freshly fetched) or ``-1`` (overflow -> compute on the CPU). ``num_indices`` holds
         the capped fetch count (for ``copy_missing``); ``num_missing_full`` the pre-cap
-        miss count (for stats). All device-side / fixed-shape, so it is CUDA-graph safe."""
+        miss count (for stats). All device-side / fixed-shape, so it is CUDA-graph safe.
+        Returns the routes split between GPU and CPU (:class:`freetoken.moe.offload_kernels.HybridRoutes`)."""
         from freetoken.moe.offload_kernels import ensure_experts_hybrid
 
         if self.collect_decode_freq:
             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_hybrid(
-            self, layer_id, expert_ids, self.hybrid_max_fetch, self.hybrid_fetch_fraction
+        return ensure_experts_hybrid(
+            self, layer_id, expert_ids, weights, self.hybrid_max_fetch, self.hybrid_fetch_fraction, clamp
         )
 
     def materialize_layer(self, layer_id: int) -> None:
@@ -911,23 +916,6 @@ def record_decode_stats(self, layer_id: int) -> None:
         was eight torch ops per layer per step, all captured into the decode graph.
         """
 
-    def record_decode_stats_hybrid(self, layer_id: int) -> None:
-        """Hybrid stats: full miss count (pre-cap), the PCIe-fetched count (capped), and
-        the active count. The CPU computes (missing - fetched) experts. Device-side;
-        accumulates both the scalar totals and the per-layer breakdown."""
-        assert 0 <= layer_id < self.num_layers, f"layer_id {layer_id} out of range [0, {self.num_layers})"
-        missing = self.num_missing_full.sum()
-        fetched = self.num_indices.sum()
-        active = self.active_mask.sum()
-        self.stat_missing += missing
-        self.stat_fetched += fetched
-        self.stat_active += active
-        self.stat_calls += 1
-        self.stat_missing_layer[layer_id] += missing
-        self.stat_fetched_layer[layer_id] += fetched
-        self.stat_active_layer[layer_id] += active
-        self.stat_steps_layer[layer_id] += 1
-
     def decode_miss_stats(self) -> dict:
         if self.decode_target == "hybrid":
             active = int(self.stat_active.item())

python/freetoken/moe/offload_kernels.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+from typing import NamedTuple
 
 import torch
 import triton
@@ -40,9 +41,23 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None:
     )
 
 
+class HybridRoutes(NamedTuple):
+    """One layer's routes split between the GPU and the CPU, each [bs, top_k]."""
+
+    gpu_slots: torch.Tensor  # cache slot of each route; a CPU route is -1, or 0 with clamp
+    gpu_w: torch.Tensor  # router weights, 0 for a CPU route
+    cpu_ids: torch.Tensor  # expert ids, -1 for a GPU route
+
+
 def ensure_experts_hybrid(
-    cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, fetch_fraction: float = 0.0
-) -> None:
+    cache,
+    layer_id: int,
+    expert_ids: torch.Tensor,
+    weights: torch.Tensor,
+    max_fetch: int,
+    fetch_fraction: float = 0.0,
+    clamp: bool = False,
+) -> HybridRoutes:
     """Capped-fetch variant of ``ensure_experts`` (hybrid backend).
 
     Identical LRU bookkeeping, but only the first ``max_fetch`` of this step's missing
@@ -52,12 +67,17 @@ def ensure_experts_hybrid(
     split (fraction = pcie_bw / cpu_bw): fetch ~fraction of the step's misses, rounded to
     the integer that makes the PCIe fetch and the CPU overflow compute finish closest to
     together. ``num_indices`` = capped fetch count (copy_missing); ``num_missing_full`` =
-    pre-cap miss count (stats)."""
-    # Q16 fixed point so the GPU kernel and the CPU reference cap identically (no float).
+    pre-cap miss count (stats).
+
+    The same launch splits the routes: ``expert_ids`` becomes the GPU slots in place, and
+    ``clamp`` gives CPU routes slot 0 instead of -1 for GEMM kernels that read every slot."""
+    assert expert_ids.is_cuda, "the hybrid ensure kernel runs on the GPU"
+    # Q16 fixed point: the kernel caps in integers.
     frac_q16 = min(1 << 16, max(0, round(fetch_fraction * (1 << 16))))
-    if not expert_ids.is_cuda:
-        return _ensure_experts_hybrid_cpu(cache, layer_id, expert_ids, max_fetch, frac_q16)
-    _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, max_fetch, frac_q16)
+    weights = weights.contiguous()
+    routes = HybridRoutes(expert_ids, torch.empty_like(weights), torch.empty_like(expert_ids))
+    _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, weights, routes, clamp, max_fetch, frac_q16)
+    return routes
 
 
 def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None:
@@ -95,7 +115,14 @@ def reset_cache(cache) -> None:
 
 
 def _ensure_experts_hybrid_gpu(
-    cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, frac_q16: int
+    cache,
+    layer_id: int,
+    expert_ids: torch.Tensor,
+    weights: torch.Tensor,
+    routes: HybridRoutes,
+    clamp: bool,
+    max_fetch: int,
+    frac_q16: int,
 ) -> None:
     block_e = triton.next_power_of_2(cache.num_experts)
     block_c = triton.next_power_of_2(cache.cache_size)
@@ -112,6 +139,10 @@ def _ensure_experts_hybrid_gpu(
         cache.num_indices,
         cache.num_missing_full,
         cache.expert_recency,
+        weights,
+        routes.cpu_ids,
+        routes.gpu_w,
+        *(cache.hybrid_stats if cache.collect_stats else (expert_ids,) * 2),
         layer_id,
         expert_ids.numel(),
         int(max_fetch),
@@ -121,73 +152,12 @@ def _ensure_experts_hybrid_gpu(
         BLOCK_E=block_e,
         BLOCK_C=block_c,
         BY_RECENCY=_HYBRID_FETCH_BY_RECENCY,
+        CLAMP=clamp,
+        STATS=cache.collect_stats,
         num_warps=num_warps,
     )
 
 
-def _ensure_experts_hybrid_cpu(
-    cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, frac_q16: int
-) -> None:
-    """CPU reference mirror of the hybrid kernel (eviction/fetch decisions bit-identical to
-    the GPU path; see tests/test_offload_lru_kernels.py). Fetches at most ``max_fetch`` (or
-    the bandwidth-matched ``~frac_q16/2^16 * misses`` when ``frac_q16`` > 0) of the missing
-    experts; overflow misses are rewritten to -1. With ``BY_RECENCY`` the fetch set is the
-    most-recently-active misses (ties -> lower id); else the lowest ids."""
-    seen = []
-    for expert in expert_ids.view(-1).tolist():
-        if expert not in seen:
-            seen.append(expert)
-
-    cache.active_mask.zero_()
-    step = int(cache.step.item()) + 1
-    cache.step.fill_(step)
-    for expert in seen:
-        cache.active_mask[expert] = 1
-
-    for expert in seen:
-        slot = int(cache.slot_for_id[layer_id, expert].item())
-        if slot != -1:
-            cache.usage[slot] = step
-
-    missing = [e for e in seen if int(cache.slot_for_id[layer_id, e].item()) == -1]
-    if _HYBRID_FETCH_BY_RECENCY:
-        rec = cache.expert_recency[layer_id].tolist()
-        missing.sort(key=lambda e: (-rec[e], e))
-    else:
-        missing.sort()
-    if frac_q16 > 0:
-        m, q = len(missing), 1 << 16
-        lo = (m * frac_q16) >> 16
-        cost = lambda f: max(f * (q - frac_q16), (m - f) * frac_q16)  # noqa: E731
-        max_fetch = lo if cost(lo) <= cost(lo + 1) else lo + 1
-    num_fetch = min(len(missing), int(max_fetch))
-    cache.num_missing_full.fill_(len(missing))
-    cache.num_indices.fill_(num_fetch)
-
-    usage = cache.usage.tolist()
-    for idx in range(num_fetch):
-        expert = missing[idx]
-        victim = min(range(cache.cache_size), key=lambda s: (usage[s], s))
-        old_id = int(cache.id_of_slot[victim].item())
-        if old_id >= 0:
-            cache.slot_for_id.view(-1)[old_id] = -1
-        cache.id_of_slot[victim] = layer_id * cache.num_experts + expert
-        cache.slot_for_id[layer_id, expert] = victim
-        cache.usage[victim] = step
-        usage[victim] = step
-        cache.evict_slots[idx] = victim
-        cache.src_indices[idx] = expert  # layer-local row
-
-    if _HYBRID_FETCH_BY_RECENCY:
-        for expert in seen:
-            cache.expert_recency[layer_id, expert] = step
-
-    # Overflow misses keep slot_for_id == -1, so the rewrite below yields -1 for them.
-    flat = expert_ids.view(-1)
-    for i in range(flat.numel()):
-        flat[i] = int(cache.slot_for_id[layer_id, int(flat[i].item())].item())
-
-
 def _materialize_layer_gpu(cache, layer_id: int) -> None:
     block = triton.next_power_of_2(max(cache.num_experts, cache.cache_size))
     _materialize_layer_kernel[(1,)](
@@ -300,6 +270,11 @@ def _ensure_experts_hybrid_kernel(
     num_indices_ptr,
     num_missing_full_ptr,
     expert_recency_ptr,
+    weights_ptr,
+    cpu_ids_ptr,
+    gpu_w_ptr,
+    stats_ptr,
+    stats_layer_ptr,
     layer_id,
     num_active,
     max_fetch,
@@ -309,6 +284,8 @@ def _ensure_experts_hybrid_kernel(
     BLOCK_E: tl.constexpr,
     BLOCK_C: tl.constexpr,
     BY_RECENCY: tl.constexpr,
+    CLAMP: tl.constexpr,
+    STATS: tl.constexpr,
 ):
     """Capped-fetch timestamp-LRU (hybrid backend).
 
@@ -401,8 +378,26 @@ def _ensure_experts_hybrid_kernel(
     for i in tl.range(num_active):
         e = tl.load(expert_ids_ptr + i)
         s = tl.load(slot_for_id_ptr + base + e)
+        on_gpu = s >= 0
+        tl.store(cpu_ids_ptr + i, tl.where(on_gpu, -1, e))
+        tl.store(gpu_w_ptr + i, tl.where(on_gpu, tl.load(weights_ptr + i), 0.0))
+        if CLAMP:
+            s = tl.maximum(s, 0)
         tl.store(expert_ids_ptr + i, s)
 
+    # stats_ptr: (missing, fetched, active, calls); stats_layer_ptr: the same per layer, [num_layers, 4]
+    if STATS:
+        n_active = tl.sum((is_active & e_mask).to(tl.int64))
+        tl.store(stats_ptr + 0, tl.load(stats_ptr + 0) + num_missing)
+        tl.store(stats_ptr + 1, tl.load(stats_ptr + 1) + num_fetch)
+        tl.store(stats_ptr + 2, tl.load(stats_ptr + 2) + n_active)
+        tl.store(stats_ptr + 3, tl.load(stats_ptr + 3) + 1)
+        row = stats_layer_ptr + layer_id * 4
+        tl.store(row + 0, tl.load(row + 0) + num_missing)
+        tl.store(row + 1, tl.load(row + 1) + num_fetch)
+        tl.store(row + 2, tl.load(row + 2) + n_active)
+        tl.store(row + 3, tl.load(row + 3) + 1)
+
     # Bump every active expert's recency to this step (LRU on the expert): an overflow miss
     # computed on the CPU now ranks high if it recurs, so it gets fetched next time.
     if BY_RECENCY: