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, [])