[ROCm][Perf] Use wvSplitK for single-output GEMMs (#53283)
Signed-off-by: tangzzycc <3081129260@qq.com>
00CC committed
Sep 23, 2026 at 14:38 UTC
eb2e91d3bac83fb8fbb37601590a25a5e5ebc06e
2 files changed
+21
-5
tests/model_executor/layers/test_rocm_unquantized_gemm.py
+20
-4
@@ -17,9 +17,10 @@ if current_platform.is_cuda():
17
from vllm.model_executor.layers import utils
18
19
20
-def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
20
+@pytest.mark.parametrize("m", [1, 128])
21
+def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch, m):
22
x = torch.randn(1, 64, dtype=torch.float16)
22
- weight = torch.randn(128, 64, dtype=torch.float16)
23
+ weight = torch.randn(m, 64, dtype=torch.float16)
24
25
monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
26
monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
@@ -42,6 +43,20 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
43
assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
44
45
46
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only kernel test")
47
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
48
+@pytest.mark.parametrize("n", range(1, 6))
49
+def test_rocm_unquantized_gemm_single_output_real_kernel(dtype, n):
50
+ torch.manual_seed(0)
51
+ x = torch.randn(n, 2048, device="cuda", dtype=dtype)
52
+ weight = torch.randn(1, 2048, device="cuda", dtype=dtype)
53
+
54
+ out = utils.rocm_unquantized_gemm_impl(x, weight, None)
55
+ ref = torch.nn.functional.linear(x, weight, None)
56
+
57
+ torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
58
+
59
+
60
def test_rocm_unquantized_gemm_makes_skinny_activation_contiguous(monkeypatch):
61
x = torch.randn(64, 4, dtype=torch.float16).t()
62
weight = torch.randn(128, 64, dtype=torch.float16)
@@ -149,11 +164,12 @@ def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch,
164
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
165
166
152
-def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch):
167
+@pytest.mark.parametrize("m", [1, 128])
168
+def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch, m):
169
# wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must
170
# fall back to torch.nn.functional.linear.
171
x = torch.randn(6, 64, dtype=torch.float16)
156
- weight = torch.randn(128, 64, dtype=torch.float16)
172
+ weight = torch.randn(m, 64, dtype=torch.float16)
173
174
monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
175
monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
vllm/model_executor/layers/utils.py
+1
-1
@@ -347,7 +347,7 @@ def rocm_unquantized_gemm_impl(
347
# The skinny kernels assume contiguous K elements. A shape-preserving
348
# reshape can retain a transposed activation's non-contiguous strides.
349
x_view = x.reshape(-1, x.size(-1)).contiguous()
350
- if m > 8 and 0 < n <= 5:
350
+ if (m == 1 or m > 8) and 0 < n <= 5:
351
cu_count = num_compute_units()
352
out = ops.wvSplitK(weight, x_view, cu_count, bias)
353
return out.reshape(*x.shape[:-1], weight.shape[0])