From 7bab776a1662308e4d1c52c98cf7047df605ff4b Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Tue, 27 Oct 2026 15:54:11 +0100 Subject: [PATCH] prain: tiered forward with CPU path --- prain/tiered.py | 152 +++++++++++++++++++++++++++++++++++++++++++ tests/test_tiered.py | 128 ++++++++++++++++++++++++++++++++++++ 2 files changed, 280 insertions(+) create mode 100644 prain/tiered.py create mode 100644 tests/test_tiered.py diff --git a/prain/tiered.py b/prain/tiered.py new file mode 100644 index 0000000..1925c8b --- /dev/null +++ b/prain/tiered.py @@ -0,0 +1,152 @@ +"""Run one MoE layer when only some experts fit in VRAM. Needs torch, not vLLM.""" + +import torch +import torch.nn.functional as F + +HEAT_INDEX = {"count": 0, "weight": 1, "recency": 2} + + +def route_stats(topk_weights, topk_ids, num_experts, recency_half_life): + """{expert: (tokens, router weight mass, recency-weighted mass)} for every routed expert, + with one host sync. Later tokens weigh more for recency. Route id -1 counts for nothing.""" + tokens = topk_ids.shape[0] + valid = (topk_ids >= 0).float() + ids = topk_ids.clamp_min(0).long().flatten() + position = torch.arange(tokens, device=topk_ids.device, dtype=torch.float32) + decay = torch.pow(0.5, (tokens - 1 - position) / recency_half_life)[:, None] + weight = topk_weights.float() * valid + stats = torch.stack([ + torch.bincount(ids, weights=valid.flatten(), minlength=num_experts), + torch.bincount(ids, weights=weight.flatten(), minlength=num_experts), + torch.bincount(ids, weights=(weight * decay).flatten(), minlength=num_experts), + ]).tolist() + return {e: (round(stats[0][e]), stats[1][e], stats[2][e]) + for e in range(num_experts) if stats[0][e] > 0} + + +def draft_rows(spans, rows, device): + """Bool [rows]: True for the token rows of draft requests (spans from prain.drafts).""" + mask = torch.zeros(rows, dtype=torch.bool) + for first, end in spans: + mask[first:min(end, rows)] = True + return mask.to(device) + + +def available_experts(topk_ids, draft, resident, num_experts): + """Bool [num_experts]: experts in VRAM, plus every expert a full (non-draft) token routes + to, because those are copied in this forward anyway.""" + ids = topk_ids.long() + full_routes = ((ids >= 0) & ~draft[:, None]).flatten().float() + available = torch.bincount(ids.clamp_min(0).flatten(), weights=full_routes, + minlength=num_experts) > 0 + if resident: + available[torch.tensor(sorted(resident), device=ids.device)] = True + return available + + +def degrade(topk_weights, topk_ids, draft, available, min_weight): + """Let draft tokens skip experts that are not available. + + A draft token keeps its heaviest unavailable routes only while less than min_weight of its + router weight would be computed; its kept routes are scaled up to its full router weight. + Dropped routes get id -1 and weight 0, which route_stats and remap ignore. + draft: bool [tokens]; available: bool [num_experts]. Full tokens are left unchanged.""" + ids = topk_ids.long() + valid = ids >= 0 + w = topk_weights.to(torch.promote_types(topk_weights.dtype, torch.float32)) * valid + missing = valid & draft[:, None] & ~available[ids.clamp_min(0)] + total = w.sum(1, keepdim=True) + short = (min_weight * total - (w * ~missing).sum(1, keepdim=True)).clamp_min(0) + order = (w * missing).argsort(1, descending=True) # heaviest missing route first + heavy = torch.gather(w * missing, 1, order) + keep_sorted = (heavy > 0) & (heavy.cumsum(1) - heavy < short) # still short before it + keep = torch.zeros_like(keep_sorted).scatter(1, order, keep_sorted) + drop = missing & ~keep + kept = (w * ~drop).sum(1, keepdim=True) + scale = torch.where(kept > 0, total / kept, torch.ones_like(kept)) + weights = torch.where(drop, torch.zeros_like(w), w * scale).to(topk_weights.dtype) + return weights, torch.where(drop, torch.full_like(topk_ids, -1), topk_ids) + + +def remap(topk_weights, topk_ids, slot_of, num_experts): + """Point routes at VRAM slots for one round. + + Routes to experts outside this round (and padding routes with id -1) get slot 0 + and weight 0, so any fused-MoE kernel can run them unchanged.""" + table = torch.full((num_experts,), -1, dtype=torch.int64) + if slot_of: + table[torch.tensor(list(slot_of), dtype=torch.int64)] = torch.tensor( + list(slot_of.values()), dtype=torch.int64 + ) + table = table.to(topk_ids.device) + ids = topk_ids.long() + slots = torch.where(ids >= 0, table[ids.clamp_min(0)], -1) + weights = torch.where(slots >= 0, topk_weights, torch.zeros_like(topk_weights)) + return weights, slots.clamp_min(0).to(topk_ids.dtype) + + +def cpu_experts(x, topk_weights, topk_ids, experts, weights_of): + """The routes to `experts`, computed on the CPU like vLLM's fused kernel with activation + "gelu_tanh": w13 = [gate; up], y = (gelu_tanh(x @ gate.T) * (x @ up.T)) @ w2.T times the + route weight. weights_of(e) -> (w13 [2I, H], w2 [H, I]). Returns [tokens, hidden] fp32.""" + out = torch.zeros(x.shape, dtype=torch.float32) + for e in sorted(experts): + rows, cols = (topk_ids == e).nonzero(as_tuple=True) + if rows.numel() == 0: + continue + w13, w2 = weights_of(e) + h = x[rows].to(w13.dtype) @ w13.t() + half = w13.shape[0] // 2 + act = F.gelu(h[:, :half].float(), approximate="tanh") * h[:, half:].float() + y = (act.to(w2.dtype) @ w2.t()).float() + out.index_add_(0, rows, y * topk_weights[rows, cols].float()[:, None]) + return out + + +def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight", + recency_half_life=32.0, miss=None, cpu_weights=None, on_stats=None, + after_enqueue=None, draft=None, min_weight=0.5): + """One MoE forward over VRAM slots. + + cache: SlotCache. load(expert, slot) copies an expert into a slot. kernel(x, weights, + slot_ids, used_slots) runs the stock fused-MoE kernel on the slot tensors. With a + [miss] config other than "copy", the experts it picks run on the CPU (cpu_weights(e) + gives their weights) while the GPU rounds run. on_stats(stats, tokens) sees the routing + statistics; after_enqueue(cpu, stats) runs once all GPU work of this forward is queued. + draft: bool [tokens] for the rows of draft requests, or None; see degrade(). + """ + raw = None + if draft is not None: + if on_stats is not None: # the trace records the routing, not what drafts skipped + raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life) + available = available_experts(topk_ids, draft, set(cache.pinned) | set(cache.slots), + cache.num_experts) + topk_weights, topk_ids = degrade(topk_weights, topk_ids, draft, available, min_weight) + stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life) # host sync + tokens = int(topk_ids.shape[0]) + if on_stats is not None: + on_stats(raw if raw is not None else stats, tokens) + counts = {e: s[0] for e, s in stats.items()} + cpu = cache.choose_cpu(counts, miss) if miss is not None else set() + if cpu: # copy the inputs now: later the copy would wait for the queued GPU rounds + host = x.cpu(), topk_weights.cpu(), topk_ids.cpu() + index = HEAT_INDEX[heat] + outs = [] + for slot_of, loads in cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu): + for expert, slot in loads: + load(expert, slot) + if slot_of: + weights, slots = remap(topk_weights, topk_ids, slot_of, cache.num_experts) + outs.append(kernel(x, weights, slots, list(slot_of.values()))) + if after_enqueue is not None: + after_enqueue(cpu, stats) + if cpu: # runs while the GPU works through its rounds + outs.append(cpu_experts(*host, cpu, cpu_weights).to(x.device)) + if not outs: + return torch.zeros_like(x) + if len(outs) == 1: + return outs[0].to(x.dtype) + # ponytail: every extra round re-runs all tokens with weight 0 for experts it does not + # hold; only big batches need many rounds. Upgrade: compact tokens per round. + acc = torch.promote_types(x.dtype, torch.float32) + return torch.stack([o.to(acc) for o in outs]).sum(0).to(x.dtype) diff --git a/tests/test_tiered.py b/tests/test_tiered.py new file mode 100644 index 0000000..cafb16b --- /dev/null +++ b/tests/test_tiered.py @@ -0,0 +1,128 @@ +import unittest + +try: + import torch +except ImportError: # the planning host has no torch; the vLLM container does + torch = None + +from prain.config import Miss, Policy +from prain.policy import SlotCache + +if torch is not None: + from prain.tiered import (available_experts, cpu_experts, degrade, remap, route_stats, + tiered_forward) + +HIDDEN, INTER, EXPERTS = 6, 4, 10 + + +def naive_moe(x, w13, w2, topk_weights, topk_ids): + """Gated-GELU experts, one route at a time; id -1 = no route.""" + out = torch.zeros(x.shape, dtype=torch.float64) + for t in range(x.shape[0]): + for j in range(topk_ids.shape[1]): + e = int(topk_ids[t, j]) + if e >= 0: + h = w13[e] @ x[t] + act = torch.nn.functional.gelu(h[:INTER], approximate="tanh") * h[INTER:] + out[t] += topk_weights[t, j] * (w2[e] @ act) + return out + + +@unittest.skipIf(torch is None, "needs torch") +class TieredForward(unittest.TestCase): + def setUp(self): + torch.manual_seed(0) + self.w13 = torch.randn(EXPERTS, 2 * INTER, HIDDEN, dtype=torch.float64) + self.w2 = torch.randn(EXPERTS, HIDDEN, INTER, dtype=torch.float64) + self.x = torch.randn(6, HIDDEN, dtype=torch.float64) + self.topk_ids = torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 0, 1], + [2, 3, -1], [4, 5, 6]], dtype=torch.int32) + self.topk_weights = torch.rand(6, 3, dtype=torch.float64) + self.expected = naive_moe(self.x, self.w13, self.w2, self.topk_weights, self.topk_ids) + + def run_tiered(self, pinned, cache_slots, miss=None, policy=Policy(), draft=None): + cache = SlotCache(pinned, EXPERTS, cache_slots, policy) + slots13 = torch.zeros(len(pinned) + cache_slots, 2 * INTER, HIDDEN, dtype=torch.float64) + slots2 = torch.zeros(len(pinned) + cache_slots, HIDDEN, INTER, dtype=torch.float64) + for slot, e in enumerate(pinned): + slots13[slot], slots2[slot] = self.w13[e], self.w2[e] + loads, seen = [], [] + + def load(expert, slot): + loads.append(expert) + slots13[slot], slots2[slot] = self.w13[expert], self.w2[expert] + + def kernel(x, w, ids, used): + return naive_moe(x, slots13, slots2, w, ids) + + out = tiered_forward(self.x, self.topk_weights, self.topk_ids, cache, load, kernel, + miss=miss, cpu_weights=lambda e: (self.w13[e], self.w2[e]), + on_stats=lambda stats, tokens: seen.append((tokens, sorted(stats))), + draft=draft, min_weight=0.5) + self.assertEqual(seen[0][0], 6) + self.seen_experts = seen[0][1] + return out, loads, cache + + def test_matches_dense_moe_when_everything_needs_rounds(self): + out, loads, _ = self.run_tiered(pinned=[0, 1], cache_slots=3) # 8 others -> 3 rounds + torch.testing.assert_close(out, self.expected) + self.assertEqual(sorted(loads), [2, 3, 4, 5, 6, 7, 8, 9]) + + def test_matches_dense_moe_in_one_round(self): + out, _, _ = self.run_tiered(pinned=[], cache_slots=10) + torch.testing.assert_close(out, self.expected) + + def test_cpu_path_matches_dense_moe(self): + for mode in ("cpu", "auto"): + out, loads, cache = self.run_tiered([0], 3, Miss(mode=mode, cpu_max_tokens=1)) + torch.testing.assert_close(out, self.expected, rtol=1e-5, atol=1e-5) + self.assertEqual(cache.cpu + cache.misses + cache.hits, 10) + self.assertEqual(cpu_experts(self.x, self.topk_weights, self.topk_ids, set(), + None).abs().sum().item(), 0.0) + + def test_route_stats(self): + stats = route_stats(torch.tensor([[0.5, 0.25], [1.0, 0.0]]), + torch.tensor([[2, 1], [2, -1]]), 4, recency_half_life=1.0) + self.assertEqual(sorted(stats), [1, 2]) + self.assertEqual(stats[2][0], 2) + self.assertAlmostEqual(stats[2][1], 1.5) + self.assertAlmostEqual(stats[2][2], 0.5 * 0.5 + 1.0) # the first token is one half-life old + self.assertAlmostEqual(stats[1][2], 0.125) + + def test_remap_zeroes_routes_outside_the_round(self): + weights, slots = remap(torch.ones(1, 3), torch.tensor([[2, 7, -1]]), {7: 4}, 10) + self.assertEqual(weights.tolist(), [[0.0, 1.0, 0.0]]) + self.assertEqual(slots.tolist(), [[0, 4, 0]]) + + def test_degrade_drops_unavailable_routes_of_draft_tokens_only(self): + weights = torch.tensor([[0.5, 0.3, 0.2], [0.5, 0.3, 0.2]]) + ids = torch.tensor([[0, 1, 2], [0, 1, 2]], dtype=torch.int32) + draft = torch.tensor([False, True]) + available = torch.tensor([False, True, False, False]) + w, i = degrade(weights, ids, draft, available, min_weight=0.5) + # the draft token computes 0.3 of 1.0: it keeps expert 0 (heaviest) to reach 0.5 + self.assertEqual(i.tolist(), [[0, 1, 2], [0, 1, -1]]) + torch.testing.assert_close(w, torch.tensor([[0.5, 0.3, 0.2], [0.625, 0.375, 0.0]])) + w, i = degrade(weights, ids, draft, available, min_weight=0.0) + self.assertEqual(i[1].tolist(), [-1, 1, -1]) + torch.testing.assert_close(w[1], torch.tensor([0.0, 1.0, 0.0])) + + def test_draft_tokens_use_vram_and_what_full_tokens_load(self): + draft = torch.tensor([False, True, False, True, True, False]) + out, loads, _ = self.run_tiered(pinned=[0], cache_slots=3, draft=draft) + # available: pinned expert 0 plus every expert of the full rows 0, 2 and 5 + available = torch.zeros(EXPERTS, dtype=torch.bool) + available[[0, 1, 2, 4, 5, 6, 7, 8]] = True + self.assertTrue(torch.equal( + available_experts(self.topk_ids, draft, {0}, EXPERTS), available)) + expected = naive_moe(self.x, self.w13, self.w2, + *degrade(self.topk_weights, self.topk_ids, draft, available, 0.5)) + torch.testing.assert_close(out, expected) + torch.testing.assert_close(out[[0, 2, 5]], self.expected[[0, 2, 5]]) # full rows exact + self.assertEqual(self.seen_experts, list(range(EXPERTS))) # trace sees the routing + + def test_drafts_are_exact_when_every_expert_is_in_vram(self): + out, loads, _ = self.run_tiered(pinned=list(range(EXPERTS)), cache_slots=0, + draft=torch.ones(6, dtype=torch.bool)) + torch.testing.assert_close(out, self.expected) + self.assertEqual(loads, [])