From 13a48876a076478c36535310ac9c7ae3966219b3 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Tue, 13 Oct 2026 23:59:18 +0000 Subject: [PATCH 1/3] perf(moe): split hybrid decode routes inside the ensure kernel --- python/freetoken/layers/moe.py | 19 ++++------ python/freetoken/moe/offload_cache.py | 47 ++++++++----------------- python/freetoken/moe/offload_kernels.py | 45 ++++++++++++++++++++--- 3 files changed, 63 insertions(+), 48 deletions(-) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 4918a032fe..0fce3e8123 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -307,7 +307,7 @@ 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 @@ -315,14 +315,13 @@ def _decode_hybrid( """ 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() + # 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) + topk_weights = topk_weights.contiguous() + cpu_ids, gpu_w = torch.empty_like(topk_ids), torch.empty_like(topk_weights) + cache.ensure_experts_hybrid(self.layer_id, topk_ids, split=(topk_weights, cpu_ids, gpu_w, clamp)) gpu_slots = topk_ids - if cache.collect_stats: - cache.record_decode_stats_hybrid(self.layer_id) 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 +331,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, diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 0969f73829..53b3661c0b 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -224,24 +224,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 +854,7 @@ 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, split=None) -> None: """Capped-fetch LRU for the hybrid backend. Like :meth:`ensure_experts` but assigns slots to (and schedules copies for) at @@ -864,7 +863,8 @@ 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. + ``split``: see :func:`freetoken.moe.offload_kernels.ensure_experts_hybrid`.""" from freetoken.moe.offload_kernels import ensure_experts_hybrid if self.collect_decode_freq: @@ -873,7 +873,7 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None 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 + self, layer_id, expert_ids, self.hybrid_max_fetch, self.hybrid_fetch_fraction, split ) def materialize_layer(self, layer_id: int) -> None: @@ -911,23 +911,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()) diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index cf513f52d1..a36b66eeaa 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -41,7 +41,7 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: def ensure_experts_hybrid( - cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, fetch_fraction: float = 0.0 + cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, fetch_fraction: float = 0.0, split=None ) -> None: """Capped-fetch variant of ``ensure_experts`` (hybrid backend). @@ -52,12 +52,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).""" + pre-cap miss count (stats). + + ``split`` = (weights, cpu_ids, gpu_w, clamp), GPU only: the same launch also writes the CPU's + routes (expert id, -1 for a GPU route) and the GPU's weights (0 for a CPU route), and with + ``clamp`` rewrites CPU routes to slot 0 instead of -1.""" # Q16 fixed point so the GPU kernel and the CPU reference cap identically (no float). frac_q16 = min(1 << 16, max(0, round(fetch_fraction * (1 << 16)))) if not expert_ids.is_cuda: + assert split is None, "the route split is GPU-only" 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) + _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, max_fetch, frac_q16, split) def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None: @@ -95,7 +100,7 @@ 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, max_fetch: int, frac_q16: int, split=None ) -> None: block_e = triton.next_power_of_2(cache.num_experts) block_c = triton.next_power_of_2(cache.cache_size) @@ -112,6 +117,8 @@ def _ensure_experts_hybrid_gpu( cache.num_indices, cache.num_missing_full, cache.expert_recency, + *(split[:3] if split is not None else (expert_ids,) * 3), + *(cache.hybrid_stats if cache.collect_stats else (expert_ids,) * 2), layer_id, expert_ids.numel(), int(max_fetch), @@ -121,6 +128,9 @@ def _ensure_experts_hybrid_gpu( BLOCK_E=block_e, BLOCK_C=block_c, BY_RECENCY=_HYBRID_FETCH_BY_RECENCY, + SPLIT=split is not None, + CLAMP=split is not None and split[3], + STATS=cache.collect_stats, num_warps=num_warps, ) @@ -300,6 +310,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 +324,9 @@ def _ensure_experts_hybrid_kernel( BLOCK_E: tl.constexpr, BLOCK_C: tl.constexpr, BY_RECENCY: tl.constexpr, + SPLIT: tl.constexpr, + CLAMP: tl.constexpr, + STATS: tl.constexpr, ): """Capped-fetch timestamp-LRU (hybrid backend). @@ -401,8 +419,27 @@ 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) + if SPLIT: + 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: From 930cc184d5e54f19d1a15f9e6b2ea1686b942c5f Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Wed, 14 Oct 2026 01:48:26 +0000 Subject: [PATCH 2/3] refactor(moe): always split routes in the hybrid ensure kernel --- python/freetoken/layers/moe.py | 5 +- python/freetoken/moe/offload_cache.py | 15 ++- python/freetoken/moe/offload_kernels.py | 126 ++++++++---------------- tests/moe/test_hybrid_fetch.py | 47 +-------- 4 files changed, 54 insertions(+), 139 deletions(-) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 0fce3e8123..c7aa52cd6c 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -318,10 +318,7 @@ def _decode_hybrid( # 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) - topk_weights = topk_weights.contiguous() - cpu_ids, gpu_w = torch.empty_like(topk_ids), torch.empty_like(topk_weights) - cache.ensure_experts_hybrid(self.layer_id, topk_ids, split=(topk_weights, cpu_ids, gpu_w, clamp)) - gpu_slots = topk_ids + 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 diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 53b3661c0b..984b7124ad 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/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 @@ -854,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, split=None) -> 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,7 +869,7 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor, split=N 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. - ``split``: see :func:`freetoken.moe.offload_kernels.ensure_experts_hybrid`.""" + 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: @@ -872,8 +877,8 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor, split=N 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, split + 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: diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index a36b66eeaa..ea12b6fb17 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/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, split=None -) -> 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 @@ -54,15 +69,15 @@ def ensure_experts_hybrid( together. ``num_indices`` = capped fetch count (copy_missing); ``num_missing_full`` = pre-cap miss count (stats). - ``split`` = (weights, cpu_ids, gpu_w, clamp), GPU only: the same launch also writes the CPU's - routes (expert id, -1 for a GPU route) and the GPU's weights (0 for a CPU route), and with - ``clamp`` rewrites CPU routes to slot 0 instead of -1.""" - # Q16 fixed point so the GPU kernel and the CPU reference cap identically (no float). + 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: - assert split is None, "the route split is GPU-only" - 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, split) + 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: @@ -100,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, split=None + 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) @@ -117,7 +139,9 @@ def _ensure_experts_hybrid_gpu( cache.num_indices, cache.num_missing_full, cache.expert_recency, - *(split[:3] if split is not None else (expert_ids,) * 3), + weights, + routes.cpu_ids, + routes.gpu_w, *(cache.hybrid_stats if cache.collect_stats else (expert_ids,) * 2), layer_id, expert_ids.numel(), @@ -128,76 +152,12 @@ def _ensure_experts_hybrid_gpu( BLOCK_E=block_e, BLOCK_C=block_c, BY_RECENCY=_HYBRID_FETCH_BY_RECENCY, - SPLIT=split is not None, - CLAMP=split is not None and split[3], + 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,)]( @@ -324,7 +284,6 @@ def _ensure_experts_hybrid_kernel( BLOCK_E: tl.constexpr, BLOCK_C: tl.constexpr, BY_RECENCY: tl.constexpr, - SPLIT: tl.constexpr, CLAMP: tl.constexpr, STATS: tl.constexpr, ): @@ -419,12 +378,11 @@ 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) - if SPLIT: - 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) + 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] diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py index a26630b894..8cc8bfc590 100644 --- a/tests/moe/test_hybrid_fetch.py +++ b/tests/moe/test_hybrid_fetch.py @@ -2,17 +2,15 @@ Covers the two halves of --moe-hybrid-max-fetch auto: the profile reader that turns `ft bench bw` kernel bandwidths into a fetch fraction, and the ensure kernel's -per-step integer split (GPU kernel vs CPU reference mirror, and the balance rule). +per-step integer split rule. """ import json import os import pytest -import torch from freetoken.moe.bench_profile import PROFILE_VERSION, default_profile_path, load_backend_recommendation, load_hybrid_fetch_fraction -from freetoken.moe.offload_cache import OffloadMoeCache Q = 1 << 16 @@ -87,46 +85,3 @@ def write(path, name, verdict): # this card's own file wins over the legacy one write(default_profile_path(uuid), "FAKE GPU", "offload") assert load_backend_recommendation("bf16", gpu_name="FAKE GPU", gpu_uuid=uuid) == "offload" - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_hybrid_fraction_gpu_matches_cpu_reference(): - torch.manual_seed(0) - num_experts, cache_size, top_k, frac = 32, 40, 8, 0.415 - - def make(): - return OffloadMoeCache( - num_layers=2, num_experts=num_experts, cache_size=cache_size, - device=torch.device("cuda"), quant_format="bf16", decode_target="hybrid", - hybrid_max_fetch=num_experts, hybrid_fetch_fraction=frac, - ) - - gpu, ref = make(), make() - frac_q16 = round(frac * Q) - for step in range(64): - ids = torch.randperm(num_experts)[:top_k].to(torch.int32) - g, c = ids.clone().cuda(), ids.clone() # a CPU ids tensor drives the reference path - gpu.ensure_experts_hybrid(0, g) - ref.ensure_experts_hybrid(0, c) - missing = int(gpu.num_missing_full.item()) - fetched = int(gpu.num_indices.item()) - assert missing == int(ref.num_missing_full.item()) - assert fetched == int(ref.num_indices.item()) == _balanced_fetch(missing, frac_q16) - # slot rewrites (hit/fetched -> slot, overflow -> -1) and LRU state stay identical - assert torch.equal(g.cpu(), c) - assert torch.equal(gpu.slot_for_id.cpu(), ref.slot_for_id.cpu()) - assert torch.equal(gpu.id_of_slot.cpu(), ref.id_of_slot.cpu()) - assert (g >= 0).sum().item() == len(set(ids.tolist())) - (missing - fetched) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_hybrid_fixed_cap_unchanged(): - # fraction 0 (no profile / explicit --moe-hybrid-max-fetch) keeps the fixed cap. - cache = OffloadMoeCache( - num_layers=1, num_experts=32, cache_size=40, device=torch.device("cuda"), - quant_format="bf16", decode_target="hybrid", hybrid_max_fetch=1, - ) - ids = torch.arange(8, dtype=torch.int32).cuda() - cache.ensure_experts_hybrid(0, ids) - assert int(cache.num_missing_full.item()) == 8 - assert int(cache.num_indices.item()) == 1 From 47aa1a170767c07e07efb0e5524e26251f7128cc Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Wed, 14 Oct 2026 01:57:17 +0000 Subject: [PATCH 3/3] test(moe): drop the fetch split test left without its kernel check --- tests/moe/test_hybrid_fetch.py | 30 ++---------------------------- 1 file changed, 2 insertions(+), 28 deletions(-) diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py index 8cc8bfc590..fad4fa7843 100644 --- a/tests/moe/test_hybrid_fetch.py +++ b/tests/moe/test_hybrid_fetch.py @@ -1,8 +1,7 @@ """Hybrid decode's bandwidth-matched fetch split. -Covers the two halves of --moe-hybrid-max-fetch auto: the profile reader that turns -`ft bench bw` kernel bandwidths into a fetch fraction, and the ensure kernel's -per-step integer split rule. +Covers the profile reader of --moe-hybrid-max-fetch auto, which turns `ft bench bw` +kernel bandwidths into a fetch fraction. """ import json @@ -12,31 +11,6 @@ from freetoken.moe.bench_profile import PROFILE_VERSION, default_profile_path, load_backend_recommendation, load_hybrid_fetch_fraction -Q = 1 << 16 - - -def _balanced_fetch(num_missing: int, frac_q16: int) -> int: - """Reference split: F ~ frac * misses, rounded to whichever integer neighbor - minimizes the slower overlapped side (fetch ~ F*(1-frac), CPU ~ (M-F)*frac).""" - lo = (num_missing * frac_q16) >> 16 - cost = lambda f: max(f * (Q - frac_q16), (num_missing - f) * frac_q16) # noqa: E731 - return min(num_missing, lo if cost(lo) <= cost(lo + 1) else lo + 1) - - -def test_balanced_fetch_tracks_fraction(): - # The split follows fetched : cpu = pcie : (cpu - pcie) up to integer rounding, and - # never over/under-shoots by more than one expert. - for frac in (0.1, 0.415, 0.454, 0.7, 1.0): - q = round(frac * Q) - for m in range(0, 65): - f = _balanced_fetch(m, q) - assert 0 <= f <= m - assert abs(f - frac * m) <= 1.0 - # ceil would over-fetch here (the regression this rule fixed): 41.5% of 3 misses is - # 1.24 -> fetching 2 makes the PCIe side ~1.6x slower than balance; keep it at 1. - assert _balanced_fetch(3, round(0.415 * Q)) == 1 - assert _balanced_fetch(4, round(0.415 * Q)) == 2 - def test_load_hybrid_fetch_fraction(tmp_path): prof = {