Add padding support to wvSplitK solution for skinny GEMMs (#33762)

Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com>

tests/kernels/quantization/test_rocm_skinny_gemms.py

@@ -30,15 +30,22 @@ NKM_FACTORS_LLMM1 = [
 
 NKM_FACTORS_WVSPLITK = [
     # Different batch sizes with key dimensions
-    (1, 16, 16),
+    (1, 32, 16),
     (1, 64, 64),
     (2, 256, 256),
     (3, 1024, 1024),
     (4, 4096, 4096),
+    (4, 4096, 4096 + 1),
+    (4, 4096 + 16, 4096),
+    (4, 4096 + 16, 4096 + 1),
     # Extended K values
     (1, 9216, 512),
     (2, 10240, 1024),
     (4, 16384, 8192),
+    (4, 16384 * 2, 8192),
+    (4, 16384 * 2, 8192 + 1),
+    (4, 16384 * 2 + 16, 8192),
+    (4, 16384 * 2 + 16, 8192 + 1),
     # Minimum M constraint validation (m >= 8)
     (1, 64, 8),
     (2, 128, 8),
@@ -180,59 +187,44 @@ def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
     torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
 
 
+@pytest.mark.parametrize("xnorm", [False, True])
 @pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("seed", SEEDS)
 @pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
-def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
-    torch.manual_seed(seed)
-    cu_count = num_compute_units()
-
-    A = torch.rand(n, k, dtype=dtype, device="cuda") - 0.5
-    B = torch.rand(m, k, dtype=dtype, device="cuda") - 0.5
-
-    ref_out = torch.nn.functional.linear(A, B)
-    out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count)
-
-    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
-
-
-@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
-def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):
+@pytest.mark.parametrize("bias_mode", BIAS_MODES)
+@pytest.mark.parametrize("padded_a", [False, True])
+@pytest.mark.parametrize("padded_b", [False, True])
+def test_rocm_wvsplitk_kernel(
+    xnorm, n, k, m, dtype, seed, bias_mode, padded_a, padded_b
+):
     torch.manual_seed(seed)
     cu_count = num_compute_units()
 
-    xavier = math.sqrt(2 / k)  # normalize to avoid large output-bias deltas
-    A = (torch.rand(n, k, dtype=dtype, device="cuda") - 0.5) * xavier
-    B = (torch.rand(m, k, dtype=dtype, device="cuda") - 0.5) * xavier
-    BIAS = torch.rand(m, dtype=dtype, device="cuda") - 0.5
-
-    ref_out = torch.nn.functional.linear(A, B, BIAS)
-    out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
-
-    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
-
+    xavier = (
+        math.sqrt(2 / k) if xnorm else 1
+    )  # normalize to avoid large output-bias deltas
+    A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
+    B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
 
-@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
-def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
-    torch.manual_seed(seed)
-    cu_count = num_compute_units()
+    BIAS = None
+    if bias_mode == 1:
+        BIAS = torch.rand(m, dtype=dtype, device="cuda") * 2 - 1
+    elif bias_mode == 2:
+        BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
 
-    xavier = math.sqrt(2 / k)  # normalize to avoid large output-bias deltas
-    A = (torch.rand(n, k, dtype=dtype, device="cuda") - 0.5) * xavier
-    B = (torch.rand(m, k, dtype=dtype, device="cuda") - 0.5) * xavier
-    BIAS = torch.rand(n, m, dtype=dtype, device="cuda") - 0.5
+    if padded_a:
+        A = pad_fp8(A)
+    if padded_b:
+        B = pad_fp8(B)
 
     ref_out = torch.nn.functional.linear(A, B, BIAS)
     out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
 
-    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
+    if xnorm:
+        assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
+    else:
+        assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-2)
 
 
 @pytest.mark.parametrize("xnorm", [False, True])

vllm/model_executor/layers/utils.py

@@ -191,7 +191,6 @@ def rocm_unquantized_gemm_impl(
         and on_gfx9()
         and x.dtype in [torch.float16, torch.bfloat16]
         and k % 8 == 0
-        and x.is_contiguous()
     )
 
     if use_skinny is not True: