[ROCm][Perf] Use wvSplitK for single-output GEMMs

tests/model_executor/layers/test_rocm_unquantized_gemm.py

@@ -17,9 +17,10 @@
 from vllm.model_executor.layers import utils
 
 
-def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
+@pytest.mark.parametrize("m", [1, 128])
+def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch, m):
     x = torch.randn(1, 64, dtype=torch.float16)
-    weight = torch.randn(128, 64, dtype=torch.float16)
+    weight = torch.randn(m, 64, dtype=torch.float16)
 
     monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
     monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
@@ -42,6 +43,20 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
     assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only kernel test")
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("n", range(1, 6))
+def test_rocm_unquantized_gemm_single_output_real_kernel(dtype, n):
+    torch.manual_seed(0)
+    x = torch.randn(n, 2048, device="cuda", dtype=dtype)
+    weight = torch.randn(1, 2048, device="cuda", dtype=dtype)
+
+    out = utils.rocm_unquantized_gemm_impl(x, weight, None)
+    ref = torch.nn.functional.linear(x, weight, None)
+
+    torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+
 def test_rocm_unquantized_gemm_makes_skinny_activation_contiguous(monkeypatch):
     x = torch.randn(64, 4, dtype=torch.float16).t()
     weight = torch.randn(128, 64, dtype=torch.float16)
@@ -149,11 +164,12 @@ def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch,
     torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
 
 
-def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch):
+@pytest.mark.parametrize("m", [1, 128])
+def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch, m):
     # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must
     # fall back to torch.nn.functional.linear.
     x = torch.randn(6, 64, dtype=torch.float16)
-    weight = torch.randn(128, 64, dtype=torch.float16)
+    weight = torch.randn(m, 64, dtype=torch.float16)
 
     monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
     monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)

vllm/model_executor/layers/utils.py

@@ -347,7 +347,7 @@ def rocm_unquantized_gemm_impl(
         # The skinny kernels assume contiguous K elements. A shape-preserving
         # reshape can retain a transposed activation's non-contiguous strides.
         x_view = x.reshape(-1, x.size(-1)).contiguous()
-        if m > 8 and 0 < n <= 5:
+        if (m == 1 or m > 8) and 0 < n <= 5:
             cu_count = num_compute_units()
             out = ops.wvSplitK(weight, x_view, cu_count, bias)
             return out.reshape(*x.shape[:-1], weight.shape[0])