From d1ccc5c2944b612716ad1da42de3ba5d090ddb4b Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Thu, 29 Oct 2026 14:20:23 +0100 Subject: [PATCH] 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 | 123 +++++++++++++++++++++++++++++++++----- prain/gemv.py | 17 +++--- prain/kernels.py | 133 ++++++++++++++++++++++++++++++++++++++++++ scripts/quality.py | 27 ++++++--- tests/test_kernels.py | 111 +++++++++++++++++++++++++++++++++++ 5 files changed, 380 insertions(+), 31 deletions(-) create mode 100644 prain/kernels.py create mode 100644 tests/test_kernels.py diff --git a/prain/experts.py b/prain/experts.py index a91ff45..b75d597 100644 --- a/prain/experts.py +++ b/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: diff --git a/prain/gemv.py b/prain/gemv.py index 7b5d9df..7ee2efa 100644 --- a/prain/gemv.py +++ b/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) diff --git a/prain/kernels.py b/prain/kernels.py new file mode 100644 index 0000000..77cbb0e --- /dev/null +++ b/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]) diff --git a/scripts/quality.py b/scripts/quality.py index bec166a..8d3c3ee 100644 --- a/scripts/quality.py +++ b/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))})) diff --git a/tests/test_kernels.py b/tests/test_kernels.py new file mode 100644 index 0000000..1a76475 --- /dev/null +++ b/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()