biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Improvements to wvSplitKrc skinny GEMM solution (#34304)

Browse Source 
Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com>
This commit is contained in:
Hashem Hashemi
2026-03-10 09:14:27 -07:00
committed by GitHub
parent aefc59f088
commit 721ae79f50
3 changed files with 169 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/kernels/quantization/test_rocm_skinny_gemms.py
View File

@@ -70,7 +70,6 @@ N_FACTORS_WVSPLITKRC = [
117,
128,
]
K_FACTORS_WVSPLITKRC = [2880, 2880 + 8, 3072, 3072 + 8]
M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16]
@@ -123,10 +122,11 @@ def pad_fp8(weight):
@pytest.mark.parametrize("m", M_FACTORS_WVSPLITKRC)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("padded_a", [False, True])
@pytest.mark.parametrize("bias_mode", BIAS_MODES)
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode):
torch.manual_seed(seed)
cu_count = num_compute_units()
@@ -141,7 +141,8 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
# Given the above, how many CUs would we need?
CuNeeded = rndup_cus * GrpsShrB
# candidate for atomic reduce count splitk?
fits_wvsplitkrc = CuNeeded <= cu_count
fits_wvsplitkrc = (N_p2 * m * ((k + 512 - 1) // 512)) <= 128 * 1024 * 12
fits_wvsplitkrc &= CuNeeded <= cu_count
if not fits_wvsplitkrc:
pytest.skip("Too large for wvSplitKrc")
@@ -151,6 +152,8 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
) # 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
if padded_a:
A = pad_fp8(A)
BIAS = None
if bias_mode == 1:
@@ -159,7 +162,7 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
ref_out = torch.nn.functional.linear(A, B, BIAS)
out = ops.wvSplitKrc(B, A.view(-1, A.size(-1)), cu_count, BIAS)
out = ops.wvSplitKrc(A, B, cu_count, BIAS)
if xnorm:
torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)