prain: vLLM tiered experts layer and plugin

2026-10-27 15:55:57 +01:00
parent 7bab776a16
commit 036f4611cc
4 changed files with 450 additions and 0 deletions
@@ -0,0 +1,33 @@
+import importlib.util
+import unittest
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+GPU = torch is not None and torch.cuda.is_available() and importlib.util.find_spec("vllm") is not None
+
+
+@unittest.skipUnless(GPU, "needs vLLM and a GPU")
+class KernelMatch(unittest.TestCase):
+    def test_cpu_experts_compute_what_the_vllm_kernel_computes(self):
+        from vllm.config import VllmConfig, set_current_vllm_config
+        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+
+        from prain.tiered import cpu_experts
+
+        torch.manual_seed(0)
+        tokens, hidden, inter, experts = 8, 256, 128, 4
+        x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
+        w13 = torch.randn(experts, 2 * inter, hidden, dtype=torch.bfloat16) * 0.1
+        w2 = torch.randn(experts, hidden, inter, dtype=torch.bfloat16) * 0.1
+        topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
+        topk_weights = torch.rand(tokens, 2)
+        with set_current_vllm_config(VllmConfig()):
+            gpu = fused_experts(x.cuda(), w13.cuda(), w2.cuda(), topk_weights.cuda(),
+                                topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
+                                global_num_experts=experts).float().cpu()
+        cpu = cpu_experts(x, topk_weights, topk_ids, set(range(experts)), lambda e: (w13[e], w2[e]))
+        torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)