prain: Triton GEMV for decode GEMMs on gfx10 (rocBLAS ran bf16 GEMVs at 6% of the bandwidth; agent setup 1.58 -> 2.77 tok/s, same answers)

prain/gemv.py

@@ -0,0 +1,89 @@
+"""Decode GEMMs (at most MAX_ROWS rows) for gfx10 GPUs as a Triton GEMV.
+
+vLLM sends unquantized linear layers on ROCm through rocBLAS, except on gfx9/gfx11/gfx12 where it
+has skinny-GEMM kernels. On an RX 6600 XT (gfx1032) rocBLAS runs a bf16 GEMV at 13-16 GB/s
+(RDNA2 has no bf16 dot instructions), so the dense part of a decode step reads its weights at
+6% of the card's bandwidth. This kernel reads them at ~210 GB/s (measured, lm_head 262144 x 2816:
+93 ms -> 6.7 ms). install() wraps vLLM's dispatch_unquantized_gemm; PRAIN_GEMV=0 turns it off.
+"""
+
+import functools
+import os
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+MAX_ROWS = 4
+
+
+@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):
+    rows = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N)
+    ms = tl.arange(0, BLOCK_M)
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k0 in range(0, K, BLOCK_K):
+        cols = k0 + tl.arange(0, BLOCK_K)
+        w = tl.load(w_ptr + rows[:, None] * K + cols[None, :],
+                    mask=(rows[:, None] < N) & (cols[None, :] < K), other=0.0)
+        x = tl.load(x_ptr + ms[:, None] * K + cols[None, :],
+                    mask=(ms[:, None] < M) & (cols[None, :] < K), other=0.0)
+        acc += tl.sum(x.to(tl.float32)[:, None, :] * w.to(tl.float32)[None, :, :], axis=2)
+    tl.store(y_ptr + ms[:, None] * N + rows[None, :], acc.to(y_ptr.dtype.element_ty),
+             mask=(ms[:, None] < M) & (rows[None, :] < N))
+
+
+@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"
+
+
+def _small(x, weight, bias):
+    return (bias is None and x.is_cuda and x.dtype in (torch.bfloat16, torch.float16)
+            and weight.dtype == x.dtype and weight.is_contiguous()
+            and x.numel() <= MAX_ROWS * x.shape[-1] and fast_device(x.device.index or 0))
+
+
+def linear(x, weight, bias=None):
+    """F.linear(x, weight, bias); the Triton GEMV for up to MAX_ROWS rows on gfx10."""
+    if not _small(x, weight, bias):
+        return F.linear(x, weight, bias)
+    n, k = weight.shape
+    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))
+    return y.reshape(*x.shape[:-1], n)
+
+
+def wrap(dispatch):
+    """vLLM's dispatch_unquantized_gemm(), with decode shapes sent to linear()."""
+    def wrapped(*args, **kwargs):
+        fallback = dispatch(*args, **kwargs)
+
+        def gemm(layer, x, weight, bias=None):
+            if _small(x, weight, bias):
+                return linear(x, weight)
+            return fallback(layer, x, weight, bias)
+        return gemm
+    return wrapped
+
+
+def install():
+    """Before the model is built: every module that looks up dispatch_unquantized_gemm."""
+    from vllm.model_executor.layers import linear as linear_layers, utils
+    from vllm.model_executor.layers import vocab_parallel_embedding
+    from vllm.model_executor.layers.fused_moe.runner import moe_runner
+
+    wrapped = wrap(utils.dispatch_unquantized_gemm)
+    for module in (linear_layers, vocab_parallel_embedding, moe_runner):
+        module.dispatch_unquantized_gemm = wrapped

prain/plugin.py

@@ -10,7 +10,9 @@ def register():
 
     if "RoutedExperts" in op_registry_oot:  # vLLM may call register() more than once
         return
+    from prain import gemv
     from prain.experts import TieredRoutedExperts
 
     # vLLM's documented out-of-tree hook: every RoutedExperts(...) now builds ours.
     PluggableLayer.register_oot(TieredRoutedExperts, name="RoutedExperts")
+    gemv.install()  # decode GEMVs on gfx10

tests/test_gemv.py

@@ -0,0 +1,55 @@
+import unittest
+
+import torch
+import torch.nn.functional as F
+
+from prain import gemv
+
+GPU = torch.cuda.is_available() and gemv.fast_device(0)
+
+
+def rel(a, b):
+    return ((a.float() - b.float()).norm() / b.float().norm()).item()
+
+
+@unittest.skipUnless(GPU, "needs a gfx10 GPU")
+class TestGemv(unittest.TestCase):
+    def test_up_to_four_rows_match_a_float32_linear(self):
+        torch.manual_seed(0)
+        for n, k in ((1000, 2816), (2816, 2112), (130, 70)):
+            w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") * 0.02
+            for rows in (1, 2, 3, 4):
+                x = torch.randn(rows, k, dtype=torch.bfloat16, device="cuda")
+                y = gemv.linear(x, w)
+                self.assertEqual((y.shape, y.dtype), ((rows, n), torch.bfloat16))
+                self.assertLess(rel(y, F.linear(x.float(), w.float())), 1e-2, (n, k, rows))
+        x = torch.randn(2, 1, 70, dtype=torch.float16, device="cuda")  # leading dims, fp16
+        w = torch.randn(130, 70, dtype=torch.float16, device="cuda")
+        self.assertEqual(gemv.linear(x, w).shape, (2, 1, 130))
+        self.assertLess(rel(gemv.linear(x, w), F.linear(x.float(), w.float())), 1e-2)
+
+    def test_more_rows_bias_or_strided_weights_use_f_linear(self):
+        w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
+        bias = torch.randn(64, dtype=torch.bfloat16, device="cuda")
+        x = torch.randn(5, 96, dtype=torch.bfloat16, device="cuda")
+        self.assertTrue(torch.equal(gemv.linear(x, w), F.linear(x, w)))
+        self.assertTrue(torch.equal(gemv.linear(x[:1], w, bias), F.linear(x[:1], w, bias)))
+        strided = torch.randn(96, 64, dtype=torch.bfloat16, device="cuda").t()
+        self.assertTrue(torch.equal(gemv.linear(x[:1], strided), F.linear(x[:1], strided)))
+
+    def test_wrapped_dispatch_sends_only_decode_shapes_to_the_kernel(self):
+        calls = []
+
+        def fallback(layer, x, weight, bias=None):
+            calls.append(x.shape[0])
+            return F.linear(x, weight, bias)
+
+        gemm = gemv.wrap(lambda *args, **kwargs: fallback)("auto")
+        w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
+        gemm(None, torch.randn(1, 96, dtype=torch.bfloat16, device="cuda"), w)
+        gemm(None, torch.randn(9, 96, dtype=torch.bfloat16, device="cuda"), w)
+        self.assertEqual(calls, [9])
+
+
+if __name__ == "__main__":
+    unittest.main()