316 lines
14 KiB
Python
316 lines
14 KiB
Python
"""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)),
|
|
# One dimension: cpu_max_tokens only matters for auto, and coordinate descent would never
|
|
# reach auto with a high threshold once auto <= 1 lost to cpu.
|
|
("miss", ("mode", "cpu_max_tokens"),
|
|
(("copy", 1), ("cpu", 1)) + tuple(("auto", n) for n in (1, 2, 4, 8, 16))),
|
|
("miss", "admit", (0, 1, 2)),
|
|
)
|
|
_PART = {"policy": 0, "prefetch": 1, "miss": 2}
|
|
|
|
|
|
def _with(variant, table, name, value):
|
|
names, values = (name, value) if isinstance(name, tuple) else ((name,), (value,))
|
|
parts = list(variant)
|
|
parts[_PART[table]] = replace(parts[_PART[table]], **dict(zip(names, values)))
|
|
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]
|