[CI][ROCm] Reduce kernel test runtime (#49915)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 27, 2026 at 01:44 UTC
e09900436ca2c531e980827cea2e6eaad625ae00
3 files changed
+101
-37
.buildkite/test-amd.yaml
+4
-3
@@ -1561,11 +1561,12 @@ steps:
1561
commands:
1562
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
1563
1564
-- label: Kernels Core Operation Test # TBD
1564
+- label: Kernels Core Operation Test %N # TBD
1565
timeout_in_minutes: 180
1566
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
1567
dind: false
1568
agent_pool: mi300_1
1569
+ parallelism: 3
1570
working_dir: "/vllm-workspace/tests"
1571
source_file_dependencies:
1572
- csrc/
@@ -1576,7 +1577,7 @@ steps:
1577
- vllm/_aiter_ops.py
1578
- vllm/platforms/rocm.py
1579
commands:
1579
- - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
1580
+ - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
1581
1582
- label: Kernels KDA Test # TBD
1583
timeout_in_minutes: 180
@@ -2145,7 +2146,7 @@ steps:
2146
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2147
dind: false
2148
agent_pool: mi300_1
2148
- parallelism: 4
2149
+ parallelism: 8
2150
optional: true
2151
working_dir: "/vllm-workspace/"
2152
source_file_dependencies:
tests/kernels/core/test_fused_quant_layernorm.py
+92
-34
@@ -31,7 +31,7 @@ NUM_TOKENS_HIDDEN_SIZES = [
31
32
ADD_RESIDUAL = [False, True]
33
SCALE_UBS = [True, False]
34
-GROUP_SIZES = [None, [1, 64], [1, 128]]
34
+GROUP_SIZES = [[1, 64], [1, 128]]
35
TMA_ALIGNMENTS = [0, 4]
36
SEEDS = [0]
37
CUDA_DEVICES = [
@@ -40,6 +40,86 @@ CUDA_DEVICES = [
40
41
EPS = 1e-6
42
43
+
44
+def _is_valid_config(
45
+ hidden_size: int,
46
+ has_scale_ub: bool,
47
+ quant_dtype: torch.dtype,
48
+ group_size: list[int] | None,
49
+ tma_alignment: int,
50
+) -> bool:
51
+ if group_size is not None and hidden_size % group_size[1] != 0:
52
+ return False
53
+ if group_size is not None and has_scale_ub:
54
+ return False
55
+ if (
56
+ group_size is None or quant_dtype != current_platform.fp8_dtype()
57
+ ) and tma_alignment != 0:
58
+ return False
59
+ if (
60
+ group_size is not None
61
+ and tma_alignment != 0
62
+ and hidden_size // group_size[1] % tma_alignment == 0
63
+ ):
64
+ return False
65
+ return not (has_scale_ub and quant_dtype != current_platform.fp8_dtype())
66
+
67
+
68
+def _config_id(
69
+ num_tokens: int,
70
+ hidden_size: int,
71
+ has_scale_ub: bool,
72
+ quant_dtype: torch.dtype,
73
+ group_size: list[int] | None,
74
+ tma_alignment: int,
75
+) -> str:
76
+ quant = str(quant_dtype).removeprefix("torch.")
77
+ group = "per-token" if group_size is None else f"group-{group_size[1]}"
78
+ return (
79
+ f"{num_tokens}x{hidden_size}-{quant}-{group}-"
80
+ f"tma-{tma_alignment}-scale-ub-{has_scale_ub}"
81
+ )
82
+
83
+
84
+# Filter unsupported combinations during collection. Letting each case call
85
+# pytest.skip still runs the global per-test teardown and dominates this suite.
86
+RMS_NORM_CONFIGS = [
87
+ pytest.param(
88
+ num_tokens,
89
+ hidden_size,
90
+ has_scale_ub,
91
+ quant_dtype,
92
+ group_size,
93
+ tma_alignment,
94
+ id=_config_id(
95
+ num_tokens,
96
+ hidden_size,
97
+ has_scale_ub,
98
+ quant_dtype,
99
+ group_size,
100
+ tma_alignment,
101
+ ),
102
+ )
103
+ for (
104
+ (num_tokens, hidden_size),
105
+ has_scale_ub,
106
+ quant_dtype,
107
+ (group_size, tma_alignment),
108
+ ) in itertools.product(
109
+ NUM_TOKENS_HIDDEN_SIZES,
110
+ SCALE_UBS,
111
+ QUANT_DTYPES,
112
+ [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
113
+ )
114
+ if _is_valid_config(
115
+ hidden_size,
116
+ has_scale_ub,
117
+ quant_dtype,
118
+ group_size,
119
+ tma_alignment,
120
+ )
121
+]
122
+
123
## Helpers
124
125
@@ -154,15 +234,19 @@ def ops_impl(
234
)
235
236
157
-@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
158
-@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
159
-@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
160
-@pytest.mark.parametrize("dtype", DTYPES)
161
-@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
237
@pytest.mark.parametrize(
163
- "group_size, tma_alignment",
164
- [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
238
+ (
239
+ "num_tokens",
240
+ "hidden_size",
241
+ "has_scale_ub",
242
+ "quant_dtype",
243
+ "group_size",
244
+ "tma_alignment",
245
+ ),
246
+ RMS_NORM_CONFIGS,
247
)
248
+@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
249
+@pytest.mark.parametrize("dtype", DTYPES)
250
@pytest.mark.parametrize("seed", SEEDS)
251
@pytest.mark.parametrize("device", CUDA_DEVICES)
252
@pytest.mark.parametrize("strided_input", [False, True])
@@ -185,32 +269,6 @@ def test_rms_norm(
269
torch.set_default_device(device)
270
torch.accelerator.set_device_index(device)
271
188
- if group_size is not None and hidden_size % group_size[1] != 0:
189
- # skip
190
- pytest.skip("Skip non-divisible group sizes")
191
-
192
- if group_size is not None and has_scale_ub:
193
- # blockwise baseline doesn't support scale_ub
194
- pytest.skip("scale_ub not supported for blockwise/group quantization")
195
-
196
- if (
197
- group_size is None or quant_dtype != current_platform.fp8_dtype()
198
- ) and tma_alignment != 0:
199
- # TMA alignment is only supported for groupwise fp8 kernels
200
- pytest.skip("tma alignment not supported for per-token or int8 quantization")
201
-
202
- if (
203
- group_size is not None
204
- and tma_alignment != 0
205
- and hidden_size // group_size[1] % tma_alignment == 0
206
- ):
207
- # Skip tests where TMA alignment doesn't create extra padding to save time
208
- pytest.skip("Skip TMA alignment cases where no extra padding is added")
209
-
210
- if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
211
- # skip
212
- pytest.skip("scale_ub only supported for fp8 quantization")
213
-
272
layer = RMSNorm(hidden_size, EPS).to(dtype=dtype)
273
274
# Make weights
tests/kernels/core/test_vit_fp8_scaling.py
+5
@@ -18,6 +18,11 @@ from vllm.utils.flashinfer import (
18
is_flashinfer_cudnn_fp8_prefill_attn_supported,
19
)
20
21
+pytestmark = pytest.mark.skipif(
22
+ not is_flashinfer_cudnn_fp8_prefill_attn_supported(),
23
+ reason="FlashInfer cuDNN FP8 prefill attention not supported",
24
+)
25
+
26
LAYER_0 = "visual.blocks.0.attn.attn"
27
LAYER_1 = "visual.blocks.1.attn.attn"
28
NUM_HEADS = 16