[CI][ROCm] Reduce V1 attention test runtime (#49916)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 26, 2026 at 22:12 UTC
ac87549cbdff37f06c93f96f94005331772f2297
2 files changed
+66
-51
.buildkite/test-amd.yaml
+6
-20
@@ -364,22 +364,6 @@ steps:
364
commands:
365
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
366
367
-- label: V1 attention (H100-MI250) # TBD
368
- timeout_in_minutes: 180
369
- mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
370
- agent_pool: mi250_1
371
- working_dir: "/vllm-workspace/tests"
372
- source_file_dependencies:
373
- - vllm/config/attention.py
374
- - vllm/model_executor/layers/attention
375
- - vllm/v1/attention
376
- - tests/v1/attention
377
- - vllm/_aiter_ops.py
378
- - vllm/envs.py
379
- - vllm/platforms/rocm.py
380
- commands:
381
- - pytest -v -s v1/attention
382
-
367
- label: V1 others (CPU) # TBD
368
timeout_in_minutes: 180
369
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -2694,11 +2678,12 @@ steps:
2678
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
2679
- pytest -v -s v1/kv_connector/extract_hidden_states_integration
2680
2697
-- label: V1 attention (H100-MI300) # TBD
2681
+- label: V1 attention (H100-MI300) %N # TBD
2682
timeout_in_minutes: 180
2683
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2684
dind: false
2685
agent_pool: mi300_1
2686
+ parallelism: 2
2687
optional: true
2688
working_dir: "/vllm-workspace/tests"
2689
source_file_dependencies:
@@ -2710,7 +2695,7 @@ steps:
2695
- vllm/envs.py
2696
- vllm/platforms/rocm.py
2697
commands:
2713
- - pytest -v -s v1/attention
2698
+ - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
2699
2700
- label: V1 Core + KV + Metrics # TBD
2701
timeout_in_minutes: 180
@@ -3676,11 +3661,12 @@ steps:
3661
3662
#------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------#
3663
3679
-- label: V1 attention (B200-MI355) # TBD
3664
+- label: V1 attention (B200-MI355) %N # TBD
3665
timeout_in_minutes: 180
3666
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3667
dind: false
3668
agent_pool: mi355_1
3669
+ parallelism: 2
3670
working_dir: "/vllm-workspace/tests"
3671
source_file_dependencies:
3672
- vllm/config/attention.py
@@ -3691,7 +3677,7 @@ steps:
3677
- vllm/envs.py
3678
- vllm/platforms/rocm.py
3679
commands:
3694
- - pytest -v -s v1/attention
3680
+ - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
3681
3682
- label: V1 Core + KV + Metrics # TBD
3683
timeout_in_minutes: 180
tests/v1/attention/test_mla_backends.py
+60
-31
@@ -40,6 +40,10 @@ from vllm.v1.attention.backends.mla.prefill import (
40
MLAPrefillBackendEnum,
41
get_mla_prefill_backend,
42
)
43
+from vllm.v1.attention.backends.mla.prefill.base import MLADimensions
44
+from vllm.v1.attention.backends.mla.prefill.selector import (
45
+ MLAPrefillSelectorConfig,
46
+)
47
from vllm.v1.attention.backends.registry import AttentionBackendEnum
48
from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
49
from vllm.v1.kv_cache_interface import (
@@ -145,14 +149,67 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
149
torch.testing.assert_close(layer.W_UK_T, old_w_uk_t + 100)
150
151
148
-# Filtered per-test via validate_configuration (capability/deps/dims).
152
+# Validate parameter combinations during collection, before GPU fixtures run.
153
PREFILL_BACKENDS_TO_TEST = [
154
+ MLAPrefillBackendEnum.ROCM_AITER_FA,
155
MLAPrefillBackendEnum.FLASH_ATTN,
156
MLAPrefillBackendEnum.FLASHINFER,
157
MLAPrefillBackendEnum.TRTLLM_RAGGED,
158
MLAPrefillBackendEnum.TOKENSPEED_MLA,
159
]
160
161
+MLA_DIMENSIONS_TO_TEST = [
162
+ ("deepseek", 128, 128),
163
+ ("glm", 192, 256),
164
+]
165
+
166
+
167
+def _prefill_backend_dimension_params():
168
+ device_capability = current_platform.get_device_capability()
169
+ params = []
170
+ for prefill_backend in PREFILL_BACKENDS_TO_TEST:
171
+ for dimensions_id, qk_nope_head_dim, v_head_dim in MLA_DIMENSIONS_TO_TEST:
172
+ if device_capability is None:
173
+ invalid_reasons = ["device capability unavailable"]
174
+ else:
175
+ try:
176
+ invalid_reasons = (
177
+ prefill_backend.get_class().validate_configuration(
178
+ device_capability,
179
+ MLAPrefillSelectorConfig(
180
+ dtype=torch.bfloat16,
181
+ mla_dimensions=MLADimensions(
182
+ qk_nope_head_dim=qk_nope_head_dim,
183
+ qk_rope_head_dim=64,
184
+ v_head_dim=v_head_dim,
185
+ ),
186
+ ),
187
+ )
188
+ )
189
+ except ImportError:
190
+ invalid_reasons = ["ImportError"]
191
+
192
+ marks = []
193
+ if invalid_reasons:
194
+ marks.append(
195
+ pytest.mark.skip(
196
+ reason=(
197
+ f"Prefill backend {prefill_backend.name} unavailable: "
198
+ f"{invalid_reasons}"
199
+ )
200
+ )
201
+ )
202
+ params.append(
203
+ pytest.param(
204
+ prefill_backend,
205
+ qk_nope_head_dim,
206
+ v_head_dim,
207
+ id=f"{dimensions_id}-{prefill_backend}",
208
+ marks=marks,
209
+ )
210
+ )
211
+ return params
212
+
213
214
SPEC_DECODE_BACKENDS = []
215
for backend in BACKENDS_TO_TEST:
@@ -1098,11 +1155,9 @@ def run_attention_backend(
1155
@pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16])
1156
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
1157
@pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])
1101
-@pytest.mark.parametrize("prefill_backend", PREFILL_BACKENDS_TO_TEST)
1158
@pytest.mark.parametrize(
1103
- ("qk_nope_head_dim", "v_head_dim"),
1104
- [(128, 128), (192, 256)],
1105
- ids=["deepseek", "glm"],
1159
+ ("prefill_backend", "qk_nope_head_dim", "v_head_dim"),
1160
+ _prefill_backend_dimension_params(),
1161
)
1162
def test_backend_correctness(
1163
default_vllm_config,
@@ -1153,32 +1208,6 @@ def test_backend_correctness(
1208
if not backends_to_test:
1209
pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}")
1210
1156
- # Skip prefill backends that can't satisfy capability/deps/dimension constraints.
1157
- from vllm.v1.attention.backends.mla.prefill.base import MLADimensions
1158
- from vllm.v1.attention.backends.mla.prefill.selector import (
1159
- MLAPrefillSelectorConfig,
1160
- )
1161
-
1162
- try:
1163
- prefill_invalid_reasons = prefill_backend.get_class().validate_configuration(
1164
- current_platform.get_device_capability(),
1165
- MLAPrefillSelectorConfig(
1166
- dtype=torch.bfloat16,
1167
- mla_dimensions=MLADimensions(
1168
- qk_nope_head_dim=qk_nope_head_dim,
1169
- qk_rope_head_dim=64,
1170
- v_head_dim=v_head_dim,
1171
- ),
1172
- ),
1173
- )
1174
- except ImportError:
1175
- prefill_invalid_reasons = ["ImportError"]
1176
- if prefill_invalid_reasons:
1177
- pytest.skip(
1178
- f"Prefill backend {prefill_backend.name} unavailable: "
1179
- f"{prefill_invalid_reasons}"
1180
- )
1181
-
1211
batch_spec = BATCH_SPECS[batch_spec_name]
1212
is_spec_decode_test = batch_spec_name.startswith("spec_decode")
1213
unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES[b] for b in backends_to_test))