[ROCm][CI] Speed Up ROCm Skinny GEMM Tests (reduced parameterizations, (#51877)
Signed-off-by: Micah Williamson <micah.williamson@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Micah Williamson committed
Aug 11, 2026 at 20:40 UTC
045a986f3bad4f68e6409b17fffd26dcbca0f720
2 files changed
+72
-35
.buildkite/test-amd.yaml
+2
-4
@@ -1654,7 +1654,6 @@ steps:
1654
dind: false
1655
agent_pool: mi300_1
1656
optional: true
1657
- parallelism: 2
1657
working_dir: "/vllm-workspace/tests"
1658
source_file_dependencies:
1659
- csrc/quantization/
@@ -1670,7 +1669,7 @@ steps:
1669
- vllm/model_executor/kernels/
1670
- vllm/v1/attention/backends/rocm_aiter_fa.py
1671
commands:
1673
- - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
1672
+ - pytest -v -s kernels/quantization
1673
1674
- label: Kernels FP8 MoE Test (2xH100-2xMI300) # TBD
1675
timeout_in_minutes: 180
@@ -3498,7 +3497,6 @@ steps:
3497
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3498
dind: false
3499
agent_pool: mi355_1
3501
- parallelism: 2
3500
working_dir: "/vllm-workspace/tests"
3501
source_file_dependencies:
3502
- csrc/quantization/
@@ -3514,7 +3512,7 @@ steps:
3512
- vllm/model_executor/kernels/
3513
- vllm/v1/attention/backends/rocm_aiter_fa.py
3514
commands:
3517
- - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
3515
+ - pytest -v -s kernels/quantization
3516
3517
- label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD
3518
timeout_in_minutes: 180
tests/kernels/quantization/test_rocm_skinny_gemms.py
+70
-31
@@ -7,12 +7,52 @@ import torch
7
8
import vllm._custom_ops as ops
9
from tests.kernels.quant_utils import ref_dynamic_per_tensor_fp8_quant
10
+from vllm.distributed import cleanup_dist_env_and_memory
11
from vllm.platforms import current_platform
12
from vllm.platforms.rocm import on_gfx950
13
from vllm.utils.platform_utils import num_compute_units
14
15
+# Global per-test cleanup costs more than the tests themselves.
16
+# These tests can just cleanup once at the end
17
+pytestmark = pytest.mark.skip_global_cleanup
18
+
19
+SEEDS = [0]
20
+
21
+# These options are independent in the kernels, so the sets below cover every
22
+# pair of them instead of the full product. bias_mode: 0 = none, 1 = (m,),
23
+# 2 = (n, m).
24
+OPTIONS_WVSPLITKRC = [
25
+ # dtype, padded_a, bias_mode, xnorm
26
+ (torch.float16, False, 0, False),
27
+ (torch.float16, True, 1, True),
28
+ (torch.float16, True, 2, False),
29
+ (torch.bfloat16, True, 0, True),
30
+ (torch.bfloat16, False, 1, False),
31
+ (torch.bfloat16, False, 2, True),
32
+]
33
+
34
+OPTIONS_WVSPLITK = [
35
+ # dtype, padded_a, padded_b, bias_mode, xnorm
36
+ (torch.float16, False, False, 0, False),
37
+ (torch.float16, False, True, 1, True),
38
+ (torch.float16, True, False, 2, True),
39
+ (torch.bfloat16, True, True, 0, True),
40
+ (torch.bfloat16, True, False, 1, False),
41
+ (torch.bfloat16, False, True, 2, False),
42
+]
43
+
44
+OPTIONS_WVSPLITK_FP8 = [
45
+ # dtype, padded_a, padded_b, biased, xnorm
46
+ (torch.float16, False, False, True, True),
47
+ (torch.bfloat16, True, False, False, True),
48
+ (torch.bfloat16, False, True, True, True),
49
+ (torch.float16, True, False, True, False),
50
+ (torch.float16, False, True, False, False),
51
+ (torch.bfloat16, True, True, False, False),
52
+]
53
+
54
DTYPES = [torch.bfloat16, torch.float16]
15
-BIAS_MODES = [0, 1, 2]
55
+
56
# Specific (N, K, M) combinations for targeted testing
57
NKM_FACTORS_LLMM1 = [
58
# Small, medium, large cases
@@ -52,6 +92,7 @@ NKM_FACTORS_WVSPLITK = [
92
(4, 256, 8),
93
]
94
95
+# N is bucketed up to its next ^2 (16/32/64/128) and the remainder is masked
96
N_FACTORS_WVSPLITKRC = [
97
13,
98
16,
@@ -70,7 +111,9 @@ N_FACTORS_WVSPLITKRC = [
111
117,
112
128,
113
]
114
+# K shards are 512 wide, evenly divided or not, +8 for a partial 8-element load
115
K_FACTORS_WVSPLITKRC = [2880, 2880 + 8, 3072, 3072 + 8]
116
+# M tiles are 64 rows, +16 for a partial tile
117
M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16]
118
119
NKM_FACTORS_WVSPLITK_FP8 = [
@@ -106,7 +149,11 @@ NKM_FACTORS_WVSPLITK_FP8 = [
149
(4, 32768 * 2 + 16, 28672 + 16),
150
]
151
109
-SEEDS = [0]
152
+
153
+@pytest.fixture(scope="module", autouse=True)
154
+def cleanup_after_all_tests():
155
+ yield
156
+ cleanup_dist_env_and_memory()
157
158
159
def pad_fp8(weight):
@@ -116,17 +163,21 @@ def pad_fp8(weight):
163
return F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
164
165
119
-@pytest.mark.parametrize("xnorm", [False, True])
166
+def make_bias(bias_mode, n, m, dtype):
167
+ if bias_mode == 0:
168
+ return None
169
+ shape = (m,) if bias_mode == 1 else (n, m)
170
+ return torch.rand(shape, dtype=dtype, device="cuda") * 2 - 1
171
+
172
+
173
@pytest.mark.parametrize("n", N_FACTORS_WVSPLITKRC)
174
@pytest.mark.parametrize("k", K_FACTORS_WVSPLITKRC)
175
@pytest.mark.parametrize("m", M_FACTORS_WVSPLITKRC)
123
-@pytest.mark.parametrize("dtype", DTYPES)
176
+@pytest.mark.parametrize("dtype,padded_a,bias_mode,xnorm", OPTIONS_WVSPLITKRC)
177
@pytest.mark.parametrize("seed", SEEDS)
125
-@pytest.mark.parametrize("padded_a", [False, True])
126
-@pytest.mark.parametrize("bias_mode", BIAS_MODES)
178
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
179
@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
129
-def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode):
180
+def test_rocm_wvsplitkrc_kernel(n, k, m, dtype, padded_a, bias_mode, xnorm, seed):
181
torch.manual_seed(seed)
182
cu_count = num_compute_units()
183
@@ -155,13 +206,7 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode
206
if padded_a:
207
A = pad_fp8(A)
208
158
- BIAS = None
159
- if bias_mode == 1:
160
- BIAS = torch.rand(m, dtype=dtype, device="cuda") * 2 - 1
161
- elif bias_mode == 2:
162
- BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
163
- elif bias_mode == 3:
164
- BIAS = torch.rand(1, m, dtype=dtype, device="cuda") * 2 - 1
209
+ BIAS = make_bias(bias_mode, n, m, dtype)
210
211
ref_out = torch.nn.functional.linear(A, B, BIAS)
212
out = ops.wvSplitKrc(A, B, cu_count, BIAS)
@@ -197,16 +242,15 @@ def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
242
torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
243
244
200
-@pytest.mark.parametrize("xnorm", [False, True])
245
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
202
-@pytest.mark.parametrize("dtype", DTYPES)
246
+@pytest.mark.parametrize(
247
+ "dtype,padded_a,padded_b,bias_mode,xnorm",
248
+ OPTIONS_WVSPLITK,
249
+)
250
@pytest.mark.parametrize("seed", SEEDS)
251
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
205
-@pytest.mark.parametrize("bias_mode", BIAS_MODES)
206
-@pytest.mark.parametrize("padded_a", [False, True])
207
-@pytest.mark.parametrize("padded_b", [False, True])
252
def test_rocm_wvsplitk_kernel(
209
- xnorm, n, k, m, dtype, seed, bias_mode, padded_a, padded_b
253
+ n, k, m, dtype, padded_a, padded_b, bias_mode, xnorm, seed
254
):
255
torch.manual_seed(seed)
256
cu_count = num_compute_units()
@@ -217,11 +261,7 @@ def test_rocm_wvsplitk_kernel(
261
A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
262
B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
263
220
- BIAS = None
221
- if bias_mode == 1:
222
- BIAS = torch.rand(m, dtype=dtype, device="cuda") * 2 - 1
223
- elif bias_mode == 2:
224
- BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
264
+ BIAS = make_bias(bias_mode, n, m, dtype)
265
266
if padded_a:
267
A = pad_fp8(A)
@@ -236,19 +276,18 @@ def test_rocm_wvsplitk_kernel(
276
torch.testing.assert_close(out, ref_out, atol=atol, rtol=1e-2)
277
278
239
-@pytest.mark.parametrize("xnorm", [False, True])
279
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
241
-@pytest.mark.parametrize("dtype", DTYPES)
280
+@pytest.mark.parametrize(
281
+ "dtype,padded_a,padded_b,biased,xnorm",
282
+ OPTIONS_WVSPLITK_FP8,
283
+)
284
@pytest.mark.parametrize("seed", SEEDS)
243
-@pytest.mark.parametrize("padded_a", [False, True])
244
-@pytest.mark.parametrize("padded_b", [False, True])
245
-@pytest.mark.parametrize("biased", [False, True])
285
@pytest.mark.skipif(
286
not (current_platform.is_rocm() and current_platform.supports_fp8()),
287
reason="only test for rocm fp8",
288
)
289
def test_rocm_wvsplitk_fp8_kernel(
251
- xnorm, n, k, m, dtype, seed, padded_a, padded_b, biased
290
+ n, k, m, dtype, padded_a, padded_b, biased, xnorm, seed
291
):
292
torch.manual_seed(seed)
293