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.

prain/experts.py

@@ -34,9 +34,9 @@ from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
 
 from prain import drafts, profile
 from prain.config import layer_index, load_config
+from prain.kernels import DECODE_ROWS, gfx10, moe_int4_decode, unpack_int4
 from prain.policy import SlotCache, write_stats
-from prain.quant import (GROUP, cpu_matmul, layout, moe_quant_config, quantize, to_cpu_int4,
-                         to_gpu_int4)
+from prain.quant import (GROUP, cpu_matmul, layout, moe_quant_config, quantize, to_cpu_int4)
 from prain.sim import write_trace
 from prain.tiered import draft_rows, tiered_forward
 
@@ -46,6 +46,68 @@ _COPY_STREAM = []  # one side stream shared by all layers: async copies stay in
 _TRACE = []  # trace records not written yet
 _DRAFT_MASK = [None, None]  # (spans, their bool row mask on the GPU) of the current forward
 _DRAFT_TOKENS = [0]  # draft token rows seen, for the stats log
+_STAGING = {}  # stream -> _Staging
+LOW_VRAM = 128 * 2**20  # driver-free bytes below which torch's cache goes back (_release_cache)
+
+
+def _release_cache():
+    """Hand torch's cached free blocks (the freed one-row bf16 expert weights, ~11 MiB per layer,
+    and load-time temporaries) back to the driver. HIP loads kernel code objects lazily into
+    VRAM; with a full card and nothing returned, a kernel's first launch fails with
+    hipErrorNoBinaryForGpu ("no kernel image")."""
+    before = torch.cuda.mem_get_info()[0]
+    torch.cuda.empty_cache()
+    free, total = torch.cuda.mem_get_info()
+    logger.info("prain: %.0f MiB of VRAM free for kernels and activations (%.0f MiB before "
+                "releasing torch's cache, %.2f of %.2f GiB reserved)", free / 2**20,
+                before / 2**20, torch.cuda.memory_reserved() / 2**30, total / 2**30)
+
+
+class _Staging:
+    """A ring of pinned host buffers and GPU buffers for staged int4 copies: the CPU copies an
+    expert's tensors from the page cache into one pinned buffer, one DMA moves it, and
+    unpack_int4 writes the slot. The CPU fills the next buffer while the last DMA runs."""
+    SLOTS = 4
+
+    def __init__(self, nbytes, device):
+        self.host = [torch.empty(nbytes, dtype=torch.uint8).pin_memory() for _ in range(self.SLOTS)]
+        self.dev = [torch.empty(nbytes, dtype=torch.uint8, device=device)
+                    for _ in range(self.SLOTS)]
+        self.done = [None] * self.SLOTS  # event: the DMA that last read this pinned buffer
+        self.i = 0
+
+    def take(self):
+        self.i = (self.i + 1) % self.SLOTS
+        if self.done[self.i] is not None:
+            self.done[self.i].synchronize()
+        return self.host[self.i], self.dev[self.i], self.i
+
+
+def stage_int4(sources, slots, slot, ring, stream):
+    """Copy one int4 expert: sources [[w13 q, w13 s], [w2 q, w2 s]] on the CPU (q in the CPU
+    layout) -> slots[...][slot] on the GPU (q in the GPU layout), through `ring`."""
+    host, dev, i = ring.take()
+    parts, offset = [], 0
+    with profile.span("memcpy"):
+        for mats in sources:
+            for src in mats:
+                n = src.nbytes
+                host[offset:offset + n].copy_(src.reshape(-1).view(torch.uint8))
+                parts.append((offset, n))
+                offset += n
+    with profile.span("h2d"):
+        dev[:offset].copy_(host[:offset], non_blocking=True)
+        ring.done[i] = torch.cuda.Event()
+        ring.done[i].record(stream)
+    with profile.span("unpack"):
+        parts = iter(parts)
+        for dsts in slots:
+            for t, dst in enumerate(dsts):
+                o, n = next(parts)
+                if t == 0:  # the q tensor: the pack keeps the CPU kernel's layout
+                    unpack_int4(dev[o:o + n], dst[slot])
+                else:
+                    dst[slot].view(-1).view(torch.uint8).copy_(dev[o:o + n])
 
 
 def _install_draft_hook():
@@ -272,6 +334,7 @@ class TieredRoutedExperts(RoutedExperts):
                 self._seal(self.quant)
             self.pack_complete, self._writing, self._bf16 = True, False, None
         self._pack = self._pack_of(self.quant)  # [[w13 tensors], [w2 tensors]] per expert row
+        self._expert_bytes = sum(t[0].nbytes for mats in self._pack for t in mats)
         if self._qslots is None:
             self._slots, self._qconfig = [[self.w13_weight.data], [self.w2_weight.data]], None
         else:
@@ -307,16 +370,21 @@ class TieredRoutedExperts(RoutedExperts):
         return self._ram[row]
 
     def _copy(self, expert, slot):
+        if self.quant == "int4":  # the pack keeps the CPU kernel's layout: stage, then unpack
+            return self._copy_staged(expert, slot)
         for dsts, srcs in zip(self._slots, self._source(expert)):
-            for i, (dst, src) in enumerate(zip(dsts, srcs)):
-                if i == 0 and self.quant == "int4":  # the pack keeps the CPU kernel's order
-                    with profile.span("h2d"):
-                        src = src.to(dst.device, non_blocking=True)
-                    with profile.span("repack"):
-                        src = to_gpu_int4(src)
+            for dst, src in zip(dsts, srcs):
                 with profile.span("h2d" if src.device.type == "cpu" else "d2d"):
                     dst[slot].copy_(src, non_blocking=True)
 
+    def _copy_staged(self, expert, slot):
+        """int4 on the current stream: one pinned buffer, one DMA, unpack into the slot."""
+        stream = torch.cuda.current_stream()
+        ring = _STAGING.get(stream)
+        if ring is None:
+            ring = _STAGING[stream] = _Staging(self._expert_bytes, stream.device)
+        stage_int4(self._source(expert), self._slots, slot, ring, stream)
+
     def _load(self, expert, slot):
         """Demand copy on the compute stream."""
         pending = self._ready_events.pop(slot, None)
@@ -325,7 +393,8 @@ class TieredRoutedExperts(RoutedExperts):
         self._copy(expert, slot)
 
     def _load_async(self, expert, slot):
-        """Copy on the side stream; only RAM-tier (pinned) experts are truly asynchronous."""
+        """Copy on the side stream; truly asynchronous for int4 (staged) and pinned RAM-tier
+        experts."""
         stream = _copy_stream()
         last = self._last_use.get(slot)
         if last is not None:  # a queued kernel may still read the slot's old expert
@@ -337,16 +406,19 @@ class TieredRoutedExperts(RoutedExperts):
         self._ready_events[slot] = event
 
     def _run_slots(self, x, weights, ids, table, used_slots):
-        """vLLM's Triton fused-MoE kernel on the slot tensors; expert_map=table skips the
-        routes of experts outside this round."""
+        """The fused-MoE kernel on the slot tensors (vLLM's, or prain's for int4 on RDNA2);
+        expert_map=table skips the routes of experts outside this round."""
         current = torch.cuda.current_stream()
         for slot in used_slots:
             pending = self._ready_events.pop(slot, None)
             if pending is not None:
                 current.wait_event(pending)
-        out = fused_experts(x, self._slots[0][0], self._slots[1][0], weights, ids,
-                            activation=self.activation, global_num_experts=table.numel(),
-                            expert_map=table, quant_config=self._qconfig)
+        if self.quant == "int4" and gfx10(x.device.index or 0):
+            out = self._int4_gfx10(x, weights, ids, table)
+        else:
+            out = fused_experts(x, self._slots[0][0], self._slots[1][0], weights, ids,
+                                activation=self.activation, global_num_experts=table.numel(),
+                                expert_map=table, quant_config=self._qconfig)
         if self._async:
             event = torch.cuda.Event()
             event.record(current)
@@ -354,6 +426,22 @@ class TieredRoutedExperts(RoutedExperts):
                 self._last_use[slot] = event
         return out
 
+    def _int4_gfx10(self, x, weights, ids, table):
+        """RDNA2 has no bf16 dot instructions, which vLLM's int4 kernel uses: a GEMV-style
+        kernel for a few tokens (10x faster at one), else vLLM's kernel in fp16 (2x faster),
+        redone in bf16 if fp16 overflowed."""
+        if (x.shape[0] <= DECODE_ROWS
+                and getattr(self.activation, "value", self.activation) == "gelu_tanh"):
+            return moe_int4_decode(x, weights, ids, table, *self._slots[0], *self._slots[1])
+        run = lambda x, config: fused_experts(  # noqa: E731
+            x, self._slots[0][0], self._slots[1][0], weights, ids, activation=self.activation,
+            global_num_experts=table.numel(), expert_map=table, quant_config=config)
+        out = run(x.half(), moe_quant_config("int4", self._slots[0][1].half(),
+                                             self._slots[1][1].half()))
+        if not torch.isfinite(out).all():  # host sync, once per round of a prefill
+            out = run(x, self._qconfig)
+        return out.to(x.dtype)
+
     def _predict(self, module, args):
         """Pre-hook of this layer's router: guess the experts of later layers from the same
         residual stream (their routers normalize their input, so the guess is cheap and
@@ -383,8 +471,9 @@ class TieredRoutedExperts(RoutedExperts):
                 self._load_async(e, slot)
         target = LAYERS.get(self.prain_layer + prefetch.distance)
         if prefetch.mode == "pregate" and target is not None and target._ready:
+            staged = target.quant == "int4"  # staged copies are async from every tier
             guesses = [(e, p) for e, p in self._predictions.get(prefetch.distance, ())
-                       if p >= prefetch.min_prob and e in target._ram_row]
+                       if p >= prefetch.min_prob and (staged or e in target._ram_row)]
             for e, slot in target.cache.prefetch(guesses, prefetch.budget):
                 target._load_async(e, slot)
 
@@ -405,6 +494,8 @@ class TieredRoutedExperts(RoutedExperts):
         if not self._ready:
             self._materialize()
             self._ready = True
+            if self.prain_layer == max(LAYERS):  # every layer has its slots now
+                _release_cache()
         policy = self.prain.policy
         out = tiered_forward(
             x, topk_weights, topk_ids, self.cache, self._load, self._run_slots,
@@ -422,6 +513,8 @@ class TieredRoutedExperts(RoutedExperts):
         self._steps += 1
         if self._steps % self.prain.stats_every or self.prain_layer != min(LAYERS):
             return
+        if torch.cuda.mem_get_info()[0] < LOW_VRAM:  # other processes or activations took it
+            _release_cache()
         if self.prain.trace_path:
             _flush_trace()
         if not self.prain.stats_path:

prain/gemv.py

@@ -15,13 +15,15 @@ import torch.nn.functional as F
 import triton
 import triton.language as tl
 
+from prain.kernels import gfx10
+
 MAX_ROWS = 4
+# ponytail: one config, the one autotuning picked for every Gemma 4 dense shape on gfx1032.
+# Triton's autotuner benchmarks with a 256 MiB buffer that torch keeps cached, and an 8 GB card
+# has no 256 MiB to spare (lazy kernel loads then fail). Re-measure if the GPU or model changes.
+BLOCK_N, BLOCK_K, NUM_WARPS = 4, 512, 8
 
 
-@triton.autotune(
-    configs=[triton.Config({"BLOCK_N": n, "BLOCK_K": k}, num_warps=w)
-             for n, k, w in ((4, 512, 8), (8, 512, 8), (16, 256, 4), (32, 256, 2))],
-    key=["M", "N", "K"])
 @triton.jit
 def _gemv(w_ptr, x_ptr, y_ptr, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
           BLOCK_K: tl.constexpr):
@@ -42,8 +44,7 @@ def _gemv(w_ptr, x_ptr, y_ptr, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.const
 @functools.cache
 def fast_device(index):
     """gfx10 (RDNA2) GPUs, where vLLM has no skinny-GEMM kernels."""
-    name = getattr(torch.cuda.get_device_properties(index), "gcnArchName", "")
-    return name.startswith("gfx10") and os.environ.get("PRAIN_GEMV", "1") != "0"
+    return gfx10(index) and os.environ.get("PRAIN_GEMV", "1") != "0"
 
 
 def _small(x, weight, bias):
@@ -60,8 +61,8 @@ def linear(x, weight, bias=None):
     x2 = x.reshape(-1, k).contiguous()
     m = x2.shape[0]
     y = torch.empty(m, n, dtype=x.dtype, device=x.device)
-    _gemv[lambda meta: (triton.cdiv(n, meta["BLOCK_N"]),)](
-        weight, x2, y, m, n, k, BLOCK_M=triton.next_power_of_2(m))
+    _gemv[(triton.cdiv(n, BLOCK_N),)](weight, x2, y, m, n, k, BLOCK_M=triton.next_power_of_2(m),
+                                      BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, num_warps=NUM_WARPS)
     return y.reshape(*x.shape[:-1], n)
 
 

prain/kernels.py

@@ -0,0 +1,133 @@
+"""Triton kernels for int4 experts on gfx10 (RDNA2), where vLLM's int4 fused-MoE kernel pads
+every token to a 16-row tile and dot-products in bf16, which RDNA2 lacks: 2.85 ms for one token.
+
+GPU int4 layout (prain.quant): q uint8 [rows, K/2], even k in the low nibble, bf16 scales
+[rows, K/64], w = (nibble - 8) * scale. Experts use gelu_tanh(gate) * up, gate = the first half
+of the w13 rows (Gemma 4).
+"""
+
+import functools
+
+import torch
+import triton
+import triton.language as tl
+
+GROUP = 64
+DECODE_ROWS = 4  # up to this many tokens per round use moe_int4_decode
+
+
+@functools.cache
+def gfx10(index):
+    """RDNA2 GPUs: no bf16 dot instructions, where vLLM's bf16 GEMV/int4 kernels are slow."""
+    return getattr(torch.cuda.get_device_properties(index), "gcnArchName", "").startswith("gfx10")
+
+
+@triton.jit
+def _gelu_tanh(x):  # 0.5 x (1 + tanh(z)) = x sigmoid(2z), z = sqrt(2/pi) (x + 0.044715 x^3)
+    return x * tl.sigmoid(1.5957691216057308 * (x + 0.044715 * x * x * x))
+
+
+@triton.jit
+def _gate_up(x_ptr, ids_ptr, table_ptr, q_ptr, s_ptr, h_ptr, TOPK: tl.constexpr,
+             H: tl.constexpr, N2: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    """h[pair] = W13[slot] @ x[token] for one (token, route) pair; routes off the table skip."""
+    pair = tl.program_id(0)
+    slot = tl.load(table_ptr + tl.load(ids_ptr + pair))
+    if slot >= 0:
+        m = pair // TOPK
+        rows = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
+        rmask = rows < N2
+        half = tl.arange(0, BLOCK_K // 2)
+        q_base = q_ptr + slot.to(tl.int64) * (N2 * (H // 2))
+        s_base = s_ptr + slot.to(tl.int64) * (N2 * (H // 64))
+        acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+        for k0 in range(0, H, BLOCK_K):
+            j = k0 // 2 + half
+            b = tl.load(q_base + rows[:, None] * (H // 2) + j[None, :], mask=rmask[:, None],
+                        other=0).to(tl.int32)
+            s = tl.load(s_base + rows[:, None] * (H // 64) + (j // 32)[None, :],
+                        mask=rmask[:, None], other=0.0).to(tl.float32)
+            xe = tl.load(x_ptr + m * H + 2 * j).to(tl.float32)
+            xo = tl.load(x_ptr + m * H + 2 * j + 1).to(tl.float32)
+            lo = (b & 15).to(tl.float32) - 8.0
+            hi = (b >> 4).to(tl.float32) - 8.0
+            acc += tl.sum((lo * xe[None, :] + hi * xo[None, :]) * s, axis=1)
+        tl.store(h_ptr + pair * N2 + rows, acc, mask=rmask)
+
+
+@triton.jit
+def _down(h_ptr, ids_ptr, table_ptr, w_ptr, q_ptr, s_ptr, y_ptr, H: tl.constexpr,
+          I: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    """y[pair] = weight * W2[slot] @ (gelu_tanh(gate) * up) for one (token, route) pair."""
+    pair = tl.program_id(0)
+    slot = tl.load(table_ptr + tl.load(ids_ptr + pair))
+    if slot >= 0:
+        rows = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
+        rmask = rows < H
+        half = tl.arange(0, BLOCK_K // 2)
+        q_base = q_ptr + slot.to(tl.int64) * (H * (I // 2))
+        s_base = s_ptr + slot.to(tl.int64) * (H * (I // 64))
+        h_base = h_ptr + pair * (2 * I)
+        acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+        for k0 in range(0, I, BLOCK_K):
+            j = k0 // 2 + half
+            kmask = 2 * j < I
+            ae = _gelu_tanh(tl.load(h_base + 2 * j, mask=kmask, other=0.0)) * \
+                tl.load(h_base + I + 2 * j, mask=kmask, other=0.0)
+            ao = _gelu_tanh(tl.load(h_base + 2 * j + 1, mask=kmask, other=0.0)) * \
+                tl.load(h_base + I + 2 * j + 1, mask=kmask, other=0.0)
+            mask = rmask[:, None] & kmask[None, :]
+            b = tl.load(q_base + rows[:, None] * (I // 2) + j[None, :], mask=mask,
+                        other=0).to(tl.int32)
+            s = tl.load(s_base + rows[:, None] * (I // 64) + (j // 32)[None, :], mask=mask,
+                        other=0.0).to(tl.float32)
+            lo = (b & 15).to(tl.float32) - 8.0
+            hi = (b >> 4).to(tl.float32) - 8.0
+            acc += tl.sum((lo * ae[None, :] + hi * ao[None, :]) * s, axis=1)
+        tl.store(y_ptr + pair * H + rows, acc * tl.load(w_ptr + pair), mask=rmask)
+
+
+def moe_int4_decode(x, weights, ids, table, w13, w13s, w2, w2s, blocks=(16, 256, 32, 128)):
+    """sum over routes r of weights[m, r] * expert(x[m]) for the routes whose expert `table`
+    maps to a slot (others add nothing). x [M, H] bf16 with few rows; ids [M, TOPK], padding
+    pointing at the table's last row (-1); slot tensors in the GPU int4 layout above."""
+    m, topk = ids.shape
+    hidden, n2 = x.shape[1], w13.shape[1]
+    pairs = m * topk
+    bn1, bk1, bn2, bk2 = blocks
+    assert hidden % bk1 == 0, "gate/up blocks must tile the hidden size (no k mask there)"
+    x, ids = x.contiguous(), ids.contiguous()
+    h = torch.empty(pairs, n2, dtype=torch.float32, device=x.device)
+    y = torch.zeros(pairs, hidden, dtype=torch.float32, device=x.device)
+    _gate_up[(pairs, triton.cdiv(n2, bn1))](x, ids, table, w13, w13s, h, TOPK=topk, H=hidden,
+                                            N2=n2, BLOCK_N=bn1, BLOCK_K=bk1, num_warps=4)
+    _down[(pairs, triton.cdiv(hidden, bn2))](h, ids, table, weights.float().contiguous(), w2,
+                                             w2s, y, H=hidden, I=n2 // 2, BLOCK_N=bn2,
+                                             BLOCK_K=bk2, num_warps=4)
+    return y.view(m, topk, hidden).sum(1).to(x.dtype)
+
+
+@triton.jit
+def _unpack(src_ptr, dst_ptr, N: tl.constexpr, K: tl.constexpr, BLOCK_R: tl.constexpr,
+            BLOCK_J: tl.constexpr):
+    """CPU int4 layout (torch's AVX2 order: [N/32, K, 16] bytes, rows d and d+16 share a byte)
+    -> GPU layout [N, K/2]; prain.quant.to_gpu_int4 in one pass."""
+    rows = tl.program_id(0) * BLOCK_R + tl.arange(0, BLOCK_R)
+    j = tl.program_id(1) * BLOCK_J + tl.arange(0, BLOCK_J)
+    mask = (rows < N)[:, None] & (j < K // 2)[None, :]
+    r = rows % 32
+    base = (rows // 32) * (K * 16) + r % 16
+    shift = ((r // 16) * 4)[:, None]
+    even = tl.load(src_ptr + base[:, None] + (2 * j * 16)[None, :], mask=mask, other=0)
+    odd = tl.load(src_ptr + base[:, None] + ((2 * j + 1) * 16)[None, :], mask=mask, other=0)
+    lo = (even.to(tl.int32) >> shift) & 15
+    hi = (odd.to(tl.int32) >> shift) & 15
+    tl.store(dst_ptr + rows[:, None] * (K // 2) + j[None, :], (lo | (hi << 4)).to(tl.uint8),
+             mask=mask)
+
+
+def unpack_int4(src, dst, block=(32, 64)):
+    """dst (uint8 [N, K/2], GPU layout) <- src (the same bytes in the CPU layout, flat)."""
+    n, half = dst.shape
+    _unpack[(triton.cdiv(n, block[0]), triton.cdiv(half, block[1]))](
+        src, dst, N=n, K=2 * half, BLOCK_R=block[0], BLOCK_J=block[1])

scripts/quality.py

@@ -33,10 +33,13 @@ def post(url, path, body):
 
 
 def greedy(url, model, ids, tokens):
+    """(text, token ids, their logprobs) of a greedy answer."""
     choice = post(url, "/v1/completions", {
         "model": model, "prompt": ids, "max_tokens": tokens, "temperature": 0, "logprobs": 1,
         "return_tokens_as_token_ids": True})["choices"][0]
-    return choice["text"], [int(t.split(":")[1]) for t in choice["logprobs"]["tokens"]]
+    logprobs = choice["logprobs"]
+    return (choice["text"], [int(t.split(":")[1]) for t in logprobs["tokens"]],
+            logprobs["token_logprobs"])
 
 
 def nll(url, model, ids, answer):
@@ -56,6 +59,9 @@ def main():
     parser.add_argument("--tokens", type=int, default=48, help="answer tokens")
     parser.add_argument("--code", default="prain/policy.py", help="file for the code question")
     parser.add_argument("--code-chars", type=int, default=900)
+    parser.add_argument("--greedy-only", action="store_true",
+                        help="score: skip the nll (prompt logprobs need ~0.5 GiB of free VRAM "
+                             "for the logits); compare greedy answers only")
     args = parser.parse_args()
     if args.action == "make":
         with open(args.code) as f:
@@ -65,7 +71,7 @@ def main():
             ids = post(args.url, "/tokenize", {"model": args.model, "add_generation_prompt": True,
                                                "messages": [{"role": "user", "content": prompt}]})
             ids = ids["tokens"]
-            text, answer = greedy(args.url, args.model, ids, args.tokens)
+            text, answer, _ = greedy(args.url, args.model, ids, args.tokens)
             refs.append({"prompt": prompt, "ids": ids, "answer": answer, "text": text})
             print(json.dumps({"prompt": prompt[:60], "answer": text[:120]}), flush=True)
         with open(args.ref, "w") as f:
@@ -73,15 +79,20 @@ def main():
         return
     with open(args.ref) as f:
         refs = json.load(f)
-    scores, same = [], []
+    scores, same, own = [], [], []
     for ref in refs:
-        scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
-        text, answer = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
+        if not args.greedy_only:
+            scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
+        text, answer, logprobs = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
+        own.append(sum(logprobs) / len(logprobs))
         same.append(next((i for i, (a, b) in enumerate(zip(answer, ref["answer"])) if a != b),
                          min(len(answer), len(ref["answer"]))))
-        print(json.dumps({"prompt": ref["prompt"][:40], "nll": round(scores[-1], 3),
-                          "same_tokens": same[-1], "answer": text[:100]}), flush=True)
-    print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3), "same_tokens": same,
+        print(json.dumps({"prompt": ref["prompt"][:40],
+                          "nll": round(scores[-1], 3) if scores else None,
+                          "own_lp": round(own[-1], 3), "same_tokens": same[-1],
+                          "answer": text[:100]}), flush=True)
+    print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3) if scores else None,
+                      "mean_own_lp": round(sum(own) / len(own), 3), "same_tokens": same,
                       "identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))}))
 
 

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()