From 168251e81c34ef99322e22a49bc43960d0b2eebe Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Tue, 27 Oct 2026 13:52:09 +0100 Subject: [PATCH] prain: traces, simulator and tuner --- prain/sim.py | 311 ++++++++++++++++++++++++++++++++++++++++++++++ tests/test_sim.py | 83 +++++++++++++ 2 files changed, 394 insertions(+) create mode 100644 prain/sim.py create mode 100644 tests/test_sim.py diff --git a/prain/sim.py b/prain/sim.py new file mode 100644 index 0000000..eb14cf8 --- /dev/null +++ b/prain/sim.py @@ -0,0 +1,311 @@ +"""Routing traces, a synthetic trace generator, a simulator with a simple cost model, and a +tuner that searches policy variants and per-layer slot counts. Pure Python, no torch. + +A trace is gzip JSON lines, one per MoE layer per forward, in execution order: +{"l": layer, "t": tokens, "e": [experts], "c": [tokens per expert], "w": [router weight +mass], "r": [recency-weighted mass], "p": {"1": [[expert, prob], ...], "2": [...]}}. +"p" holds predictions for layers l+1 and l+2 (pregate) and is optional. Routing does not +depend on placement, so one trace can evaluate any placement and policy. +""" + +import gzip +import json +import random +from dataclasses import dataclass, field, replace + +from prain.config import GIB, LayerPlacement, Miss, Policy, Prefetch +from prain.policy import COUNTERS, Belady, SlotCache + +HEAT_FIELD = {"count": "c", "weight": "w", "recency": "r"} + + +@dataclass(frozen=True) +class Cost: + """Rough RX 6600 XT numbers. scripts/measure.py writes measured ones as JSON.""" + expert_mib: float = 11.34 # one expert + h2d_gbps: float = 12.0 # pinned RAM -> VRAM, PCIe 4.0 x8 + disk_gbps: float = 3.0 # pack file -> RAM + round_ms: float = 0.12 # one fused-MoE kernel launch + route_ms: float = 0.0012 # GPU work per (token, expert) route + other_ms: float = 0.6 # attention + dense MLP of one layer per forward + sync_ms: float = 0.05 # host sync per layer + cpu_ms: float = 0.5 # CPU work per (token, expert) route + + def copy_ms(self, disk): + mb = self.expert_mib * 1.048576 + return mb / self.h2d_gbps + (mb / self.disk_gbps if disk else 0.0) + + +def load_cost(path, expert_mib): + values = {"expert_mib": expert_mib} + if path: + with open(path) as f: + values |= json.load(f) + values["expert_mib"] = expert_mib + return Cost(**values) + + +def write_trace(path, records, mode="wt"): + with gzip.open(path, mode) as f: + for r in records: + f.write(json.dumps(r, separators=(",", ":")) + "\n") + + +def read_forwards(path): + """[{layer: record}] per forward, in execution order.""" + forwards, last = [], None + with gzip.open(path, "rt") as f: + for line in f: + if not line.strip(): + continue + r = json.loads(line) + if last is None or r["l"] <= last: + forwards.append({}) + forwards[-1][r["l"]] = r + last = r["l"] + return forwards + + +def synthetic_trace(layers=30, experts=128, top_k=8, forwards=300, prompt_every=100, + prompt_tokens=128, zipf=1.1, stay=0.5, accuracy=0.7, seed=0): + """Zipf-popular experts per layer, decode tokens that keep some of the previous token's + experts, a prompt every prompt_every forwards, and pregate predictions that are right + with probability accuracy (accuracy**2 two layers ahead).""" + rng = random.Random(seed) + rank = [rng.sample(range(experts), experts) for _ in range(layers)] + weights = [1 / (i + 1) ** zipf for i in range(experts)] + previous = [[] for _ in range(layers)] + records = [] + for f in range(forwards): + tokens = prompt_tokens if f % prompt_every == 0 else 1 + routed = [] + for layer in range(layers): + count, recency = {}, {} + for t in range(tokens): + chosen = {e for e in previous[layer] if rng.random() < stay} if tokens == 1 else set() + while len(chosen) < top_k: + chosen.add(rank[layer][rng.choices(range(experts), weights)[0]]) + decay = 0.5 ** ((tokens - 1 - t) / 32) + for e in chosen: + count[e] = count.get(e, 0) + 1 + recency[e] = recency.get(e, 0.0) + decay / top_k + if tokens == 1: + previous[layer] = sorted(chosen) + routed.append((count, recency)) + for layer, (count, recency) in enumerate(routed): + ids = sorted(count) + record = {"l": layer, "t": tokens, "e": ids, "c": [count[e] for e in ids], + "w": [round(count[e] / top_k, 4) for e in ids], + "r": [round(recency[e], 4) for e in ids]} + predictions = {} + for d in (1, 2): + if layer + d < layers: + nxt = routed[layer + d][0] + guess = [] + for e in sorted(nxt, key=lambda e: (-nxt[e], e))[:top_k]: + g = e if rng.random() < accuracy ** d else rng.randrange(experts) + if g not in guess: + guess.append(g) + predictions[str(d)] = [[g, round(1 / top_k, 4)] for g in guess] + if predictions: + record["p"] = predictions + records.append(record) + return records + + +@dataclass +class Result: + total_ms: float = 0.0 + decode_ms: float = 0.0 # summed over forwards with fewer than scan_tokens tokens + decode_forwards: int = 0 + prefill_ms: float = 0.0 + prefill_tokens: int = 0 + counters: dict = field(default_factory=dict) + layer_ms: dict = field(default_factory=dict) + + @property + def hit_rate(self): + c = self.counters + return c["hits"] / max(1, c["hits"] + c["misses"] + c["cpu"]) + + @property + def ms_per_decode(self): + return self.decode_ms / max(1, self.decode_forwards) + + +def _future(forwards, layer): + future, n = {}, 0 + for fwd in forwards: + if layer in fwd: + n += 1 + for e in fwd[layer]["e"]: + future.setdefault(e, []).append(n) + return future + + +def simulate(forwards, placements, policy, prefetch, miss, cost, belady=False): + """Replay a trace. Model: one PCIe link serves copies in issue order; demand copies of + a round wait for the previous round's kernel; async copies (prefetch, admit-behind) are + issued after the layer's demand copies; the CPU path runs alongside the GPU.""" + if prefetch.mode == "pregate" and not any("p" in r for fwd in forwards for r in fwd.values()): + raise ValueError("prefetch.mode = pregate needs a trace with predictions ('p')") + caches, disk = {}, {} + for layer, p in placements.items(): + n = len(p.vram) + len(p.ram) + len(p.disk) + disk[layer] = set(p.disk) + cost_of = [policy.disk_cost if e in disk[layer] else 1.0 for e in range(n)] + evictor = Belady(_future(forwards, layer)) if belady else None + caches[layer] = SlotCache(p.vram, n, p.cache_slots, policy, cost_of, evictor) + result = Result(layer_ms={layer: 0.0 for layer in caches}) + key, t, link, ready = HEAT_FIELD[policy.heat], 0.0, 0.0, {} + for fwd in forwards: + start, tokens = t, 0 + for layer in sorted(fwd): + if layer not in caches: + continue + rec, cache = fwd[layer], caches[layer] + tokens = rec["t"] + counts = dict(zip(rec["e"], rec["c"])) + heat = dict(zip(rec["e"], rec[key])) + cpu = cache.choose_cpu(counts, miss) + gpu = t + cost.sync_ms + for slot_of, loads in cache.plan(heat, tokens, cpu): + if loads: + link = max(link, gpu) + sum(cost.copy_ms(e in disk[layer]) for e, _ in loads) + gpu = link + land = max((ready.pop((layer, e), 0.0) for e in slot_of), default=0.0) + gpu = max(gpu, land) + cost.round_ms + cost.route_ms * sum(counts[e] for e in slot_of) + cpu_end = t + cost.sync_ms + cost.cpu_ms * sum(counts[e] for e in cpu) + later = [] + if cpu and miss.admit: + hot = sorted((e for e in cpu if e not in disk[layer]), key=lambda e: (-heat[e], e)) + later += [(layer, e) for e, _ in cache.prefetch([(e, 1.0) for e in hot], miss.admit)] + target = layer + prefetch.distance + if prefetch.mode == "pregate" and target in caches: + guesses = [(e, p) for e, p in rec.get("p", {}).get(str(prefetch.distance), []) + if p >= prefetch.min_prob and e not in disk[target]] + later += [(target, e) for e, _ in caches[target].prefetch(guesses, prefetch.budget)] + for where in later: + link = max(link, t + cost.sync_ms) + cost.copy_ms(False) + ready[where] = link + end = max(gpu, cpu_end) + cost.other_ms + result.layer_ms[layer] += end - t + t = end + if tokens < policy.scan_tokens: + result.decode_ms += t - start + result.decode_forwards += 1 + else: + result.prefill_ms += t - start + result.prefill_tokens += tokens + result.total_ms = t + result.counters = {k: sum(getattr(c, k) for c in caches.values()) for k in COUNTERS} + return result + + +def describe(result, bound, cost): + c = result.counters + copied = (c["misses"] + c["prefetches"]) * cost.expert_mib / 1024 + rate = 1000 / result.ms_per_decode if result.decode_forwards else 0.0 + return [ + f"forwards {result.decode_forwards} decode + prefill of {result.prefill_tokens} tokens", + f"VRAM hit rate {100 * result.hit_rate:.1f}% " + f"(eviction alone at best: {100 * bound.hit_rate:.1f}%)", + f"copies {c['misses']} + prefetched {c['prefetches']} (used {c['prefetch_hits']}, " + f"wasted {c['prefetch_wasted']}) = {copied:.1f} GiB, CPU-run {c['cpu']}", + f"decode {result.ms_per_decode:.1f} ms/forward ({rate:.1f} tok/s), " + f"prefill {result.prefill_ms / max(1, result.prefill_tokens):.2f} ms/token", + f"simulated total {result.total_ms / 1000:.2f} s", + ] + + +DIMENSIONS = ( + ("policy", "eviction", ("lru", "lfu", "gdsf")), + ("policy", "heat", ("count", "weight", "recency")), + ("policy", "scan_keep", (0, 1, 2)), + ("policy", "promote", (0, 2, 4)), + ("prefetch", "mode", ("none", "pregate")), + ("prefetch", "distance", (1, 2)), + ("prefetch", "budget", (1, 2, 4)), + ("miss", "mode", ("copy", "auto", "cpu")), + ("miss", "admit", (0, 1, 2)), +) +_PART = {"policy": 0, "prefetch": 1, "miss": 2} + + +def _with(variant, table, name, value): + parts = list(variant) + parts[_PART[table]] = replace(parts[_PART[table]], **{name: value}) + return tuple(parts) + + +def ram_tier(forwards, layers, num_experts, rows): + """The `rows` most-used (layer, expert) pairs of the trace -> {layer: set of experts}.""" + uses = {(layer, e): 0 for layer in layers for e in range(num_experts)} + for fwd in forwards: + for layer, rec in fwd.items(): + for e in rec["e"]: + uses[(layer, e)] += 1 + ranked = sorted(uses, key=lambda le: (-uses[le], le))[:rows] + ram = {layer: set() for layer in layers} + for layer, e in ranked: + ram[layer].add(e) + return ram + + +def tune(forwards, num_experts, expert_vram_gib, ram_gib, cost, base=(Policy(), Prefetch(), Miss()), + passes=2, window=3): + """Coordinate descent over DIMENSIONS with uniform slots, then per-layer slot counts by + greedy marginal gain. Returns (variant, ram, slots_per_layer, result, bound, rows).""" + layers = sorted({layer for fwd in forwards for layer in fwd}) + expert_size = cost.expert_mib * 2**20 + slots = int(expert_vram_gib * GIB // expert_size) // len(layers) + if slots < 1: + raise ValueError("expert VRAM budget is less than one slot per layer") + ram = ram_tier(forwards, layers, num_experts, int(ram_gib * GIB // expert_size)) + + def placements(per_layer): + return {layer: LayerPlacement((), tuple(sorted(ram[layer])), + tuple(e for e in range(num_experts) if e not in ram[layer]), + per_layer[layer]) for layer in layers} + + uniform = placements({layer: slots for layer in layers}) + predicted = any("p" in r for fwd in forwards for r in fwd.values()) + results = {} + + def run(variant): + if variant not in results: + ok = variant[0].scan_keep < slots + results[variant] = simulate(forwards, uniform, *variant, cost) if ok else None + return results[variant] + + best = base + if best[1].mode == "pregate" and not predicted: + best = _with(best, "prefetch", "mode", "none") + for _ in range(passes): + for table, name, values in DIMENSIONS: + if table == "prefetch" and (not predicted or (name != "mode" and best[1].mode == "none")): + continue + if (table, name) == ("miss", "admit") and best[2].mode == "copy": + continue + scored = [(r.total_ms, i, v) for i, v in enumerate(_with(best, table, name, x) for x in values) + if (r := run(v)) is not None] + best = min(scored)[2] + lo, hi = max(1, slots - window, best[0].scan_keep + 1), slots + window + curves = {} + for s in range(lo, hi + 1): + r = simulate(forwards, placements({layer: s for layer in layers}), *best, cost) + curves |= {(layer, s): ms for layer, ms in r.layer_ms.items()} + per = {layer: lo for layer in layers} + for _ in range((slots - lo) * len(layers)): + gains = [(curves[(layer, per[layer])] - curves[(layer, per[layer] + 1)], -layer) + for layer in layers if per[layer] < hi] + if not gains: + break + per[-max(gains)[1]] += 1 + final = simulate(forwards, placements(per), *best, cost) + if final.total_ms > results[best].total_ms: # greedy on approximate curves can lose + per, final = {layer: slots for layer in layers}, results[best] + bound = simulate(forwards, placements(per), Policy(scan_tokens=best[0].scan_tokens), + Prefetch(), Miss(), cost, belady=True) + rows = sorted((r.total_ms, str(v), v, r) for v, r in results.items() if r is not None) + return best, ram, per, final, bound, [(v, r) for _, _, v, r in rows] diff --git a/tests/test_sim.py b/tests/test_sim.py new file mode 100644 index 0000000..2874f1c --- /dev/null +++ b/tests/test_sim.py @@ -0,0 +1,83 @@ +import os +import unittest + +from prain.config import LayerPlacement, Miss, Policy, Prefetch +from prain.sim import (Cost, read_forwards, simulate, synthetic_trace, tune, write_trace) +from tests.helpers import TmpCase + +SMALL = dict(layers=4, experts=32, top_k=4, forwards=120, prompt_every=60, prompt_tokens=32) + + +def forwards_of(records): + forwards, last = [], None + for r in records: + if last is None or r["l"] <= last: + forwards.append({}) + forwards[-1][r["l"]] = r + last = r["l"] + return forwards + + +def all_ram(slots, layers=4, experts=32): + return {layer: LayerPlacement((), tuple(range(experts)), (), slots) for layer in range(layers)} + + +class Traces(TmpCase): + def test_round_trip_and_determinism(self): + records = synthetic_trace(**SMALL) + self.assertEqual(records, synthetic_trace(**SMALL)) + path = os.path.join(self.dir, "t.jsonl.gz") + write_trace(path, records[:8]) + write_trace(path, records[8:], mode="at") # runs append to the same trace + forwards = read_forwards(path) + self.assertEqual(len(forwards), 120) + self.assertEqual(forwards[0][0], records[0]) + self.assertEqual(forwards[0][0]["t"], 32) + self.assertEqual(sorted(forwards[1][1]["p"]), ["1", "2"]) + self.assertNotIn("p", forwards[1][3]) # nothing after the last layer + + def test_copy_cost(self): + self.assertAlmostEqual(Cost().copy_ms(False), 0.991, places=3) + self.assertAlmostEqual(Cost().copy_ms(True), 4.955, places=3) + + +class Simulate(unittest.TestCase): + def setUp(self): + self.forwards = forwards_of(synthetic_trace(**SMALL)) + + def run_sim(self, policy=Policy(), prefetch=Prefetch(), miss=Miss(), slots=3, **kw): + return simulate(self.forwards, all_ram(slots), policy, prefetch, miss, Cost(), **kw) + + def test_belady_bounds_every_policy(self): + bound = self.run_sim(belady=True).hit_rate + for eviction in ("lru", "lfu", "gdsf"): + self.assertLessEqual(self.run_sim(Policy(eviction=eviction)).hit_rate, bound) + + def test_cpu_path_beats_copies_when_the_cpu_is_faster_than_pcie(self): + copy = self.run_sim().ms_per_decode + self.assertLess(self.run_sim(miss=Miss(mode="auto")).ms_per_decode, copy) + + def test_prefetch_helps_when_predictions_are_right(self): + none = self.run_sim().hit_rate + good = self.run_sim(prefetch=Prefetch(mode="pregate", budget=4)) + self.assertGreater(good.hit_rate, none) + self.assertGreater(good.counters["prefetch_hits"], good.counters["prefetch_wasted"]) + + def test_pregate_needs_predictions(self): + stripped = [{layer: {k: v for k, v in r.items() if k != "p"} for layer, r in f.items()} + for f in self.forwards] + with self.assertRaisesRegex(ValueError, "needs a trace with predictions"): + simulate(stripped, all_ram(3), Policy(), Prefetch(mode="pregate"), Miss(), Cost()) + + +class Tune(unittest.TestCase): + def test_tune_beats_the_defaults_within_the_budget(self): + forwards = forwards_of(synthetic_trace(**SMALL)) + cost = Cost(expert_mib=64.0) # 1 GiB = 16 slots = 4 per layer + best, ram, per, final, bound, rows = tune(forwards, 32, 1.0, 1.0, cost) + self.assertLessEqual(sum(per.values()), 16) + self.assertEqual(sum(len(r) for r in ram.values()), 16) # 1 GiB of RAM tier + default = next(r for v, r in rows if v == (Policy(), Prefetch(), Miss())) + self.assertLessEqual(final.total_ms, default.total_ms) + self.assertEqual(rows[0][1].total_ms, min(r.total_ms for _, r in rows)) + self.assertGreaterEqual(bound.hit_rate, 0.0)