prain: int4 MoE on RDNA2: GEMV-style decode kernel, fp16 rounds, staged expert copies

Measured on the live 16k agent setup (A/B with runtime flags, greedy answers vs reference):
decode 3.49 -> 5.24 tok/s, first token 4.3 -> 3.3 s, 2.3k-token prompt 27.5 -> 20.6 s,
same answers. Dropped after measuring: route stats on the CPU, a Triton router GEMV,
pregate prefetch from disk (all within noise or slower).

The GEMV keeps one config instead of autotuning: Triton benchmarks with a 256 MiB buffer
that torch keeps cached, which left 8 MiB of VRAM and made lazy kernel loads fail
(hipErrorNoBinaryForGpu). prain also returns torch cached blocks after materializing and
whenever the driver has < 128 MiB free. quality.py score --greedy-only skips prompt logprobs,
which OOM the 16k config.

tests/test_kernels.py

@@ -0,0 +1,111 @@
+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 Kernels(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(0)
+
+    def test_unpack_writes_what_to_gpu_int4_computes(self):
+        from prain.kernels import unpack_int4
+        from prain.quant import to_gpu_int4
+
+        for n, half in ((256, 128), (1408, 1408), (2816, 352)):
+            cpu_layout = torch.randint(0, 256, (n, half), dtype=torch.uint8, device="cuda")
+            dst = torch.zeros_like(cpu_layout)
+            unpack_int4(cpu_layout.flatten(), dst)
+            self.assertTrue(torch.equal(dst, to_gpu_int4(cpu_layout)), (n, half))
+
+    def test_decode_kernel_computes_what_fused_experts_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.kernels import moe_int4_decode
+        from prain.quant import dequantize, moe_quant_config, quantize
+        from prain.tiered import cpu_experts, slot_table
+
+        hidden, inter, experts, topk = 256, 128, 6, 4
+        w13 = torch.randn(experts, 2 * inter, hidden) * 0.1
+        w2 = torch.randn(experts, hidden, inter) * 0.1
+        parts = [quantize(w, "int4") for w in (w13, w2)]  # gpu layout
+        held = [1, 3, 4]  # the round's experts, in slots 0..2
+        slots = [[q[held].cuda(), s[held].cuda()] for q, s in parts]
+        table = slot_table({e: i for i, e in enumerate(held)}, experts, "cuda")
+        for tokens in (1, 2, 3, 4):
+            x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
+            ids = torch.stack([torch.randperm(experts)[:topk] for _ in range(tokens)]).int()
+            ids[0, -1] = experts  # padding route: the table's last row
+            w = torch.rand(tokens, topk)
+            got = moe_int4_decode(x.cuda(), w.cuda(), ids.cuda(), table, *slots[0], *slots[1])
+            with set_current_vllm_config(VllmConfig()):
+                vllm = fused_experts(x.cuda(), slots[0][0], slots[1][0], w.cuda(), ids.cuda(),
+                                     activation=MoEActivation.GELU_TANH,
+                                     global_num_experts=table.numel(), expert_map=table,
+                                     quant_config=moe_quant_config("int4", *[s[1] for s in slots]))
+            exact = cpu_experts(x, w, torch.where(ids < experts, ids, -1), set(held), lambda e: tuple(
+                dequantize(q, s, "int4")[e] for q, s in parts))
+            self.assertEqual((got.shape, got.dtype), ((tokens, hidden), torch.bfloat16))
+            torch.testing.assert_close(got.float().cpu(), exact, rtol=3e-2, atol=3e-2)
+            torch.testing.assert_close(got.float(), vllm.float(), rtol=3e-2, atol=3e-2)
+
+    def test_staged_copy_fills_a_slot_like_the_repack_path(self):
+        from prain.experts import _Staging, stage_int4
+        from prain.quant import quantize, to_cpu_int4, to_gpu_int4
+
+        w13, w2 = torch.randn(2, 256, 256), torch.randn(2, 256, 128)
+        pack = [[to_cpu_int4(q), s] for q, s in (quantize(w, "int4") for w in (w13, w2))]
+        slots = [[torch.zeros((3, *t.shape[1:]), dtype=t.dtype, device="cuda") for t in mats]
+                 for mats in pack]
+        ring = _Staging(sum(t[0].nbytes for mats in pack for t in mats), "cuda")
+        stream = torch.cuda.current_stream()
+        for expert, slot in ((1, 2), (0, 0), (1, 1), (0, 2), (1, 0)):  # wraps the ring
+            stage_int4([[t[expert] for t in mats] for mats in pack], slots, slot, ring, stream)
+            torch.cuda.synchronize()
+            for mats, dsts in zip(pack, slots):
+                self.assertTrue(torch.equal(dsts[0][slot], to_gpu_int4(mats[0][expert].cuda())))
+                self.assertTrue(torch.equal(dsts[1][slot], mats[1][expert].cuda()))
+
+
+    def test_int4_rounds_on_rdna2_match_the_bf16_kernel_even_when_fp16_overflows(self):
+        from types import SimpleNamespace
+
+        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.experts import TieredRoutedExperts
+        from prain.quant import moe_quant_config, quantize
+        from prain.tiered import slot_table
+
+        hidden, inter, experts = 256, 128, 4
+        parts = [quantize(w, "int4") for w in (torch.randn(experts, 2 * inter, hidden) * 0.1,
+                                                 torch.randn(experts, hidden, inter) * 0.1)]
+        slots = [[q.cuda(), s.cuda()] for q, s in parts]
+        config = moe_quant_config("int4", slots[0][1], slots[1][1])
+        layer = SimpleNamespace(_slots=slots, activation=MoEActivation.GELU_TANH, _qconfig=config)
+        table = slot_table({e: e for e in range(experts)}, experts, "cuda")
+        with set_current_vllm_config(VllmConfig()):
+            for tokens, scale in ((2, 1.0), (64, 1.0), (64, 3000.0)):  # decode kernel, fp16, overflow
+                x = (torch.randn(tokens, hidden) * scale).bfloat16().cuda()
+                ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int().cuda()
+                w = torch.rand(tokens, 2).cuda()
+                got = TieredRoutedExperts._int4_gfx10(layer, x, w, ids, table)
+                want = fused_experts(x, slots[0][0], slots[1][0], w, ids, activation=layer.activation,
+                                     global_num_experts=table.numel(), expert_map=table,
+                                     quant_config=config)
+                self.assertTrue(torch.isfinite(got).all(), (tokens, scale))
+                torch.testing.assert_close(got.float(), want.float(), rtol=3e-2, atol=3e-2 * scale,
+                                           msg=f"{tokens} tokens, x * {scale}")
+
+
+if __name__ == "__main__":
+    unittest.main()