diff --git a/prain/gemv.py b/prain/gemv.py new file mode 100644 index 0000000..7b5d9df --- /dev/null +++ b/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 diff --git a/prain/plugin.py b/prain/plugin.py index e014b75..6353fe7 100644 --- a/prain/plugin.py +++ b/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 diff --git a/tests/test_gemv.py b/tests/test_gemv.py new file mode 100644 index 0000000..c9eaba8 --- /dev/null +++ b/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()