fabolous005/prain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

prain: int8/int4 experts quantized once on the GPU, rounds through vLLM's expert_map, [miss] min_weight skips experts not in VRAM; scripts: tune_moe, quality (answer NLL on identical token ids), bench --context, measure --quant

Browse Source
This commit is contained in:
fabolous005
2026-10-28 22:07:25 +01:00
parent 3ae9bd73a7
commit 708e152549
15 changed files with 703 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_tiered.py
+37 -7
View File
@@ -9,7 +9,7 @@ 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,
from prain.tiered import (available_experts, cpu_experts, degrade, route_stats, slot_table,
tiered_forward)
HIDDEN, INTER, EXPERTS = 6, 4, 10
@@ -52,8 +52,9 @@ class TieredForward(unittest.TestCase):
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)
def kernel(x, w, ids, table, used):
self.assertTrue(bool((ids >= 0).all())) # padding id -1 arrives as num_experts
return naive_moe(x, slots13, slots2, w, table[ids.long()])
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]),
@@ -80,6 +81,33 @@ class TieredForward(unittest.TestCase):
self.assertEqual(cpu_experts(self.x, self.topk_weights, self.topk_ids, set(),
None).abs().sum().item(), 0.0)
def test_miss_min_weight_skips_missing_experts_for_every_token(self):
"""[miss] min_weight < 1: every token skips experts that are not in VRAM while it keeps
that share of its router weight; experts some token copies anyway serve all tokens."""
everyone = torch.ones(6, dtype=torch.bool)
resident = torch.zeros(EXPERTS, dtype=torch.bool)
resident[[0, 1, 2]] = True
_, first = degrade(self.topk_weights, self.topk_ids, everyone, resident, 0.5)
free = resident.clone()
free[first[first >= 0].long()] = True
w, i = degrade(self.topk_weights, self.topk_ids, everyone, free, 0.5)
out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=0.5))
torch.testing.assert_close(out, naive_moe(self.x, self.w13, self.w2, w, i))
self.assertEqual(sorted(loads), sorted(set(i[i >= 0].tolist()) - {0, 1, 2}))
self.assertLess(len(loads), 7) # computing everything copies all 7 others
self.assertEqual(self.seen_experts, list(range(EXPERTS))) # the trace sees the routing
out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=1.0))
torch.testing.assert_close(out, self.expected)
self.assertEqual(len(loads), 7)
def test_cpu_experts_take_matmul_functions(self):
"""Quantized experts hand the CPU path x -> x @ w.T functions instead of weights."""
def matmuls(e):
return (lambda v: v @ self.w13[e].t()), (lambda v: v @ self.w2[e].t())
out = cpu_experts(self.x, self.topk_weights, self.topk_ids, set(range(EXPERTS)), matmuls)
torch.testing.assert_close(out, self.expected.float())
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)
@@ -89,10 +117,10 @@ class TieredForward(unittest.TestCase):
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_slot_table_maps_the_round_and_skips_the_rest(self):
table = slot_table({7: 4, 2: 0}, 10, "cpu")
self.assertEqual(table.tolist(), [-1, -1, 0, -1, -1, -1, -1, 4, -1, -1, -1])
self.assertEqual(table.dtype, torch.int32) # one extra row: padding id 10 -> -1
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]])
@@ -115,6 +143,8 @@ class TieredForward(unittest.TestCase):
available[[0, 1, 2, 4, 5, 6, 7, 8]] = True
self.assertTrue(torch.equal(
available_experts(self.topk_ids, draft, {0}, EXPERTS), available))
_, kept = degrade(self.topk_weights, self.topk_ids, draft, available, 0.5)
available[kept[kept >= 0].long()] = True # experts some row copies anyway serve all rows
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)