From ea496e1bc19653030a5cd79e3ff58d91a498d145 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Tue, 20 Oct 2026 09:18:52 +0000 Subject: [PATCH] perf(moe): scan the hybrid eviction pool in chunks and add fetch rounding modes --- python/freetoken/moe/offload_kernels.py | 93 ++++++++++++++++++------- tests/moe/test_offload.py | 92 ++++++++++++++++++++++++ 2 files changed, 160 insertions(+), 25 deletions(-) diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index 5b28c1be8..e1fcbcda1 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -15,6 +15,17 @@ _HYBRID_FETCH_BY_RECENCY = ( os.getenv("FREETOKEN_HYBRID_FETCH", "recency").strip().lower() != "lowest_id" ) +# How fetch_fraction * misses becomes a fetch count: "min-wait" (default) picks the integer +# that makes the PCIe fetch and the CPU compute finish closest together, "stochastic" rounds +# up with probability equal to the fractional part, "prefer-fetch" always rounds up. +_HYBRID_FETCH_ROUNDING = {"min-wait": 0, "stochastic": 1, "prefer-fetch": 2} +_hybrid_fetch_round = (os.getenv("FREETOKEN_HYBRID_FETCH_ROUND") or "min-wait").strip().lower() +if _hybrid_fetch_round not in _HYBRID_FETCH_ROUNDING: + raise ValueError(f"FREETOKEN_HYBRID_FETCH_ROUND={_hybrid_fetch_round!r}: expected one of {sorted(_HYBRID_FETCH_ROUNDING)}") +_HYBRID_FETCH_ROUND = _HYBRID_FETCH_ROUNDING[_hybrid_fetch_round] +_HYBRID_FETCH_SEED = 0 +# Phase 2 of the hybrid ensure kernel scans a pool wider than this in tiles of this width. +_HYBRID_EVICT_CHUNK = 8192 def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: @@ -66,9 +77,9 @@ def ensure_experts_hybrid( experts are given a slot and scheduled for copy; the overflow misses stay non-resident and their ``expert_ids`` positions are rewritten to ``-1`` (compute on the CPU). ``fetch_fraction`` > 0 replaces the fixed cap with the bandwidth-matched - 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`` = + split (fraction = pcie_bw / cpu_bw): fetch ~fraction of the step's misses, rounded as + ``FREETOKEN_HYBRID_FETCH_ROUND`` says (default: the integer that makes the PCIe fetch and the CPU overflow compute finish closest together). + ``num_indices`` = capped fetch count (copy_missing); ``num_missing_full`` = pre-cap miss count (stats). The same launch splits the routes: ``expert_ids`` becomes the GPU slots in place, and @@ -140,8 +151,9 @@ def _ensure_experts_hybrid_gpu( id_of_slot, usage = cache.pool_state(layer_id) num_slots = id_of_slot.numel() block_e = triton.next_power_of_2(cache.num_experts) - block_c = triton.next_power_of_2(num_slots) - num_warps = 8 if block_c >= 2048 else 4 + chunk = min(_HYBRID_EVICT_CHUNK, triton.next_power_of_2(num_slots)) + # with 4 warps a 4096+ chunk spills on sm80-sm89 + num_warps = 8 if chunk >= 4096 else 4 _ensure_experts_hybrid_kernel[(1,)]( expert_ids, cache.slot_for_id, @@ -165,9 +177,12 @@ def _ensure_experts_hybrid_gpu( int(fetch_cap), cache.num_experts, num_slots, + _HYBRID_FETCH_SEED, BLOCK_E=block_e, - BLOCK_C=block_c, + CHUNK=chunk, + SMALL=num_slots <= chunk, BY_RECENCY=_HYBRID_FETCH_BY_RECENCY, + ROUND=_HYBRID_FETCH_ROUND, CLAMP=clamp, STATS=cache.collect_stats, num_warps=num_warps, @@ -298,9 +313,12 @@ def _ensure_experts_hybrid_kernel( fetch_cap, num_experts: tl.constexpr, cache_size, + seed, BLOCK_E: tl.constexpr, - BLOCK_C: tl.constexpr, + CHUNK: tl.constexpr, + SMALL: tl.constexpr, BY_RECENCY: tl.constexpr, + ROUND: tl.constexpr, CLAMP: tl.constexpr, STATS: tl.constexpr, ): @@ -319,7 +337,12 @@ def _ensure_experts_hybrid_kernel( most-recently active before this step (LRU on the expert, via ``expert_recency``), breaking ties toward the lower expert id -- this prioritizes *recurring* misses for caching, lowering the steady miss rate. Otherwise the lowest expert ids are fetched - (``missing_rank``), the original routing-blind heuristic.""" + (``missing_rank``), the original routing-blind heuristic. + + Victims are the least-used slots, ties to the lowest slot. Every usage is at most the + step before this launch, so the hit slots stamped with ``step`` in Phase 1 are never + picked. A pool of at most CHUNK slots stays in registers (SMALL); a larger one is scanned + from memory in CHUNK tiles per victim.""" step = tl.load(step_ptr) + 1 tl.store(step_ptr, step) base = layer_id * num_experts @@ -341,11 +364,19 @@ def _ensure_experts_hybrid_kernel( # F * (1 - frac), CPU time with (M - F) * frac; they balance at F = frac * M. Pick # the integer neighbor that minimizes the slower (max) side of the overlap. lo = (num_missing * fetch_frac_q16) >> 16 - cost_lo = tl.maximum(lo * ((1 << 16) - fetch_frac_q16), (num_missing - lo) * fetch_frac_q16) - cost_hi = tl.maximum( - (lo + 1) * ((1 << 16) - fetch_frac_q16), (num_missing - lo - 1) * fetch_frac_q16 - ) - max_fetch = tl.where(cost_lo <= cost_hi, lo, lo + 1) + if ROUND == 1: + # step is unique per launch, so each (layer, step) draws its own number + rem = (num_missing * fetch_frac_q16) & 0xFFFF + u = tl.rand(seed, (step & 0x7FFFFFFF).to(tl.int32)) + max_fetch = lo + ((u * 65536.0) < rem.to(tl.float32)).to(tl.int32) + elif ROUND == 2: + max_fetch = (num_missing * fetch_frac_q16 + 0xFFFF) >> 16 + else: + cost_lo = tl.maximum(lo * ((1 << 16) - fetch_frac_q16), (num_missing - lo) * fetch_frac_q16) + cost_hi = tl.maximum( + (lo + 1) * ((1 << 16) - fetch_frac_q16), (num_missing - lo - 1) * fetch_frac_q16 + ) + max_fetch = tl.where(cost_lo <= cost_hi, lo, lo + 1) is_hit = is_active & (slot >= 0) # a pool smaller than num_experts may hold fewer free slots than misses num_fetch = tl.minimum(tl.minimum(num_missing, tl.minimum(max_fetch, fetch_cap)), cache_size - tl.sum(is_hit.to(tl.int32))) @@ -366,18 +397,27 @@ def _ensure_experts_hybrid_kernel( # ---- Phase 2: evict victims by argmin(usage), only for the capped fetches ---- if num_fetch > 0: - off_c = tl.arange(0, BLOCK_C) - c_mask = off_c < cache_size - oid = tl.load(id_of_slot_ptr + off_c, mask=c_mask, other=-1) - u = tl.load(usage_ptr + off_c, mask=c_mask, other=9223372036854775807).to(tl.int64) - owner_active = c_mask & False - for i in tl.range(num_active): - ei = tl.load(expert_ids_ptr + i) - owner_active = owner_active | (oid == base + ei) - u = tl.where(owner_active | (~c_mask), 9223372036854775807, u) + off_k = tl.arange(0, CHUNK) + tl.debug_barrier() # the scan below must see the hit stamps + if SMALL: + u_reg = tl.load(usage_ptr + off_k, mask=off_k < cache_size, other=9223372036854775807).to(tl.int64) for i in tl.range(num_fetch): - victim = tl.argmin(u, axis=0).to(tl.int32) - old_id = tl.sum(tl.where(off_c == victim, oid, 0)) + if SMALL: + victim = tl.argmin(u_reg, axis=0).to(tl.int32) + u_reg = tl.where(off_k == victim, step, u_reg) + else: + best_u = tl.full((), 9223372036854775807, tl.int64) + best_s = tl.full((), 0, tl.int32) + for c0 in tl.range(0, cache_size, CHUNK): + offs = c0 + off_k + uu = tl.load(usage_ptr + offs, mask=offs < cache_size, other=9223372036854775807).to(tl.int64) + cmin = tl.min(uu, axis=0) + carg = (tl.argmin(uu, axis=0) + c0).to(tl.int32) + take = cmin < best_u + best_s = tl.where(take, carg, best_s) + best_u = tl.where(take, cmin, best_u) + victim = best_s + old_id = tl.load(id_of_slot_ptr + victim) if old_id >= 0: tl.store(slot_for_id_ptr + old_id, -1) if BY_RECENCY: @@ -390,12 +430,15 @@ def _ensure_experts_hybrid_kernel( tl.store(usage_ptr + victim, step) tl.store(evict_slots_ptr + i, victim) tl.store(src_indices_ptr + i, e) # layer-local row - u = tl.where(off_c == victim, 9223372036854775807, u) + tl.debug_barrier() # the next scan must see this victim's stamp + tl.debug_barrier() # Phase 3 reads slot_for_id entries written above # ---- Phase 3: rewrite expert_ids -> slot id (hit/fetched) or -1 (overflow -> CPU) ---- for i in tl.range(num_active): e = tl.load(expert_ids_ptr + i) s = tl.load(slot_for_id_ptr + base + e) + # every warp loads e, one thread stores below; a warp that loads after that store indexes with a slot id + tl.debug_barrier() 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)) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 8833cdf98..90bb5d3a4 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -920,6 +920,98 @@ def test_slot_pools_keep_each_layer_in_its_own_slots(target): assert [v[0].shape[0] for v in map(cache.pool_bank_views, range(L))] == [8, 8, 24] +def _hybrid_ensure_reference(st, layer, ids, weights, num_experts, pool, frac, mode, kernel_fetch): + """One hybrid ensure call on CPU copies of the cache state, written from the LRU contract.""" + base, size = pool + step = int(st["step"]) + 1 + slot_for_id, recency = st["slot_for_id"], st["expert_recency"] + p_ids, p_usage = st["id_of_slot"][base : base + size], st["usage"][base : base + size] + active = sorted(set(ids.tolist())) + hits = [e for e in active if slot_for_id[layer, e] >= 0] + missing = [e for e in active if slot_for_id[layer, e] < 0] + m, fq = len(missing), min(1 << 16, max(0, round(frac * (1 << 16)))) + lo, room = (m * fq) >> 16, size - len(hits) + if mode == "min-wait": + cost_lo = max(lo * ((1 << 16) - fq), (m - lo) * fq) + cost_hi = max((lo + 1) * ((1 << 16) - fq), (m - lo - 1) * fq) + num_fetch = min(m, lo if cost_lo <= cost_hi else lo + 1, room) + elif mode == "prefer-fetch": + num_fetch = min(m, (m * fq + 0xFFFF) >> 16, room) + else: + allowed = {min(m, lo, room)} if (m * fq) & 0xFFFF == 0 else {min(m, lo, room), min(m, lo + 1, room)} + assert kernel_fetch in allowed + num_fetch = kernel_fetch + taken = torch.zeros(size, dtype=torch.bool) + for e in hits: + p_usage[slot_for_id[layer, e]] = step + taken[slot_for_id[layer, e]] = True + plan = [] + for e in sorted(missing, key=lambda e: (-int(recency[layer, e]), e))[:num_fetch]: + victim = int(torch.argmin(p_usage.masked_fill(taken, torch.iinfo(torch.int64).max))) # first minimum + old = int(p_ids[victim]) + if old >= 0: + slot_for_id.view(-1)[old] = -1 + p_ids[victim] = layer * num_experts + e + slot_for_id[layer, e] = victim + p_usage[victim] = step + taken[victim] = True + plan.append((victim, e)) + recency[layer, active] = step + st["step"] = torch.tensor(step) + slots = slot_for_id[layer, ids.long()] + on_gpu = slots >= 0 + return st, plan, slots, torch.where(on_gpu, -1, ids), torch.where(on_gpu, weights, 0.0), m, active + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("mode", ["min-wait", "stochastic", "prefer-fetch"]) +@pytest.mark.parametrize("clamp", [False, True]) +@pytest.mark.parametrize("layout", ["global", "global-tiled", "per-layer"]) +def test_hybrid_ensure_matches_lru_reference(monkeypatch, layout, clamp, mode): + from freetoken.moe import offload_kernels + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + monkeypatch.setattr(offload_kernels, "_HYBRID_FETCH_ROUND", offload_kernels._HYBRID_FETCH_ROUNDING[mode]) + torch.manual_seed(0) + # global-tiled: more slots than one Phase 2 tile, so victims come from the tiled scan + L, E, top_k = (20, 512, 10) if layout == "global-tiled" else (4, 128, 8) + size = {"global": 300, "global-tiled": offload_kernels._HYBRID_EVICT_CHUNK + 808, "per-layer": 4 * 96}[layout] + pools = {"pool_of_layer": list(range(L)), "pool_sizes": [96] * L} if layout == "per-layer" else {} + cache = OffloadMoeCache( + num_layers=L, num_experts=E, cache_size=size, device=torch.device("cuda"), + decode_target="hybrid", hybrid_max_fetch=1, **pools, + ) + popularity = 1.0 / torch.arange(1, E + 1, dtype=torch.float64) ** 0.6 + names = ["slot_for_id", "id_of_slot", "usage", "step", "expert_recency"] + + # fetch every miss until every pool is full, then the benched-style fraction + for frac, calls in ((1.0, 4000), (0.413, 60 * L)): + for call in range(calls): + layer = call % L + if frac == 1.0 and bool((cache.id_of_slot >= 0).all()) and layer == 0: + break + cache.hybrid_fetch_fraction = [frac] * L + ids = torch.cat([torch.multinomial(popularity, top_k) for _ in range(2)]).to(torch.int32) + weights = torch.rand(ids.numel()) + st = {n: getattr(cache, n).cpu().clone() for n in names} + p = cache.pool_of_layer[layer] + routes = cache.ensure_experts_hybrid(layer, ids.cuda(), weights.cuda(), clamp=clamp) + n = int(cache.num_indices) + st, plan, slots, cpu_ids, gpu_w, m, active = _hybrid_ensure_reference( + st, layer, ids, weights, E, (cache.pool_bases[p], cache.pool_sizes[p]), frac, mode, n + ) + assert int(cache.num_missing_full) == m and n == len(plan) + assert cache.evict_slots[:n].cpu().tolist() == [v for v, _ in plan] + assert cache.src_indices[:n].cpu().tolist() == [e for _, e in plan] + assert torch.equal(routes.gpu_slots.cpu(), slots.clamp(min=0) if clamp else slots) + assert torch.equal(routes.cpu_ids.cpu(), cpu_ids) and torch.equal(routes.gpu_w.cpu(), gpu_w) + assert cache.active_mask.cpu().nonzero().flatten().tolist() == active + for name in names: + assert torch.equal(getattr(cache, name).cpu(), st[name]), name + assert bool((cache.id_of_slot >= 0).all()) + + def _make_split_cache(num_layers=2, locked=(1,), prefill_overlap=False, device="cpu"): """A [gate_up, down] bf16 cache with the given layers LOCKED (rest pinned).""" from freetoken.moe.host_banks import HostResidency