[ROCm][CI] Restore attention coverage after KV-cache layout refactor (#53351)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Aug 22, 2026 at 20:00 UTC
30b34171b113887e0e08d7f6d06e2e5a5c33b9d2
3 files changed
+50
-25
tests/kernels/attention/test_minimax_m3.py
+15
@@ -971,6 +971,21 @@ def test_aiter_sparse_pa_layout_contract(monkeypatch):
971
assert value_cache.is_contiguous()
972
973
974
+def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch):
975
+ """Do not pair AITER's separated cache layout with the Triton fallback."""
976
+ import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
977
+
978
+ monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
979
+ monkeypatch.setattr(
980
+ sparse_attn_mod.rocm_aiter_ops,
981
+ "is_shuffle_kv_cache_enabled",
982
+ lambda: True,
983
+ )
984
+
985
+ with pytest.raises(ValueError, match="num_kv_heads == 1"):
986
+ sparse_attn_mod.minimax_m3_use_aiter_sparse_pa(2)
987
+
988
+
989
def test_indexer_cache_squeezes_to_contiguous_3d():
990
"""The indexer side cache is standardized 4D with H=1: under both layouts
991
the allocator's logical view stays contiguous and squeezes (as
tests/v1/attention/test_attention_backends.py
+20
-9
@@ -42,9 +42,11 @@ BACKENDS_TO_TEST = [
42
]
43
44
45
-def _actual_backend(backend) -> AttentionBackendEnum:
45
+def _actual_backend(backend: AttentionBackendEnum | str) -> AttentionBackendEnum:
46
"""Resolve pseudo-backends (FLEX_ATTENTION_SLOW) to their real enum."""
47
- if backend == "FLEX_ATTENTION_SLOW":
47
+ if isinstance(backend, str):
48
+ if backend != "FLEX_ATTENTION_SLOW":
49
+ raise ValueError(f"Unknown pseudo-backend: {backend}")
50
return AttentionBackendEnum.FLEX_ATTENTION
51
return backend
52
@@ -266,7 +268,7 @@ def _clone_kv_cache_in_layout(
268
269
270
def run_attention_backend(
269
- backend: AttentionBackendEnum,
271
+ backend: AttentionBackendEnum | str,
272
kv_cache_spec: FullAttentionSpec,
273
layer_names: list[str],
274
vllm_config,
@@ -283,9 +285,7 @@ def run_attention_backend(
285
) -> torch.Tensor:
286
"""Run attention computation using the specified backend's AttentionImpl."""
287
286
- use_direct_block_mask = not current_platform.is_rocm() and is_torch_equal_or_newer(
287
- "2.9.0.dev0"
288
- )
288
+ use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
289
if backend == "FLEX_ATTENTION_SLOW":
290
use_direct_block_mask = False
291
backend = _actual_backend(backend)
@@ -377,7 +377,7 @@ def run_attention_backend(
377
def _test_backend_correctness(
378
batch_spec: BatchSpec,
379
model: str,
380
- backend_to_test: list[AttentionBackendEnum],
380
+ backend_to_test: list[AttentionBackendEnum | str],
381
mask_mod,
382
*,
383
causal: bool = True,
@@ -655,7 +655,7 @@ def _test_backend_correctness(
655
)
656
657
# Check numerical similarity
658
- def error_msg(msg: str, backend_name: str):
658
+ def error_msg(msg: str, backend_name: AttentionBackendEnum | str):
659
return f"[{backend_name}] output differs from SDPA baseline. {msg}"
660
661
torch.testing.assert_close(
@@ -1046,14 +1046,24 @@ if current_platform.is_rocm():
1046
SLIDING_WINDOW_BACKENDS_TO_TEST = [
1047
AttentionBackendEnum.FLEX_ATTENTION,
1048
AttentionBackendEnum.TRITON_ATTN,
1049
+ "FLEX_ATTENTION_SLOW",
1050
]
1051
else:
1052
SLIDING_WINDOW_BACKENDS_TO_TEST = [
1053
AttentionBackendEnum.FLASH_ATTN,
1054
AttentionBackendEnum.FLEX_ATTENTION,
1055
AttentionBackendEnum.TRITON_ATTN,
1056
+ "FLEX_ATTENTION_SLOW",
1057
]
1058
1059
+# Encoder-only FlexAttention always uses the slow builder, so the pseudo-backend
1060
+# would run an identical implementation twice.
1061
+SLIDING_WINDOW_ENCODER_BACKENDS_TO_TEST = [
1062
+ backend
1063
+ for backend in SLIDING_WINDOW_BACKENDS_TO_TEST
1064
+ if backend != "FLEX_ATTENTION_SLOW"
1065
+]
1066
+
1067
1068
@pytest.mark.parametrize(
1069
"batch_spec_name",
@@ -1162,7 +1172,7 @@ def test_sliding_window_encoder_backend_correctness(
1172
_test_backend_correctness(
1173
batch_spec,
1174
model,
1165
- SLIDING_WINDOW_BACKENDS_TO_TEST,
1175
+ SLIDING_WINDOW_ENCODER_BACKENDS_TO_TEST,
1176
sliding_window_mask_mod_fn,
1177
causal=False,
1178
attn_type=AttentionType.ENCODER_ONLY,
@@ -1173,6 +1183,7 @@ def test_sliding_window_encoder_backend_correctness(
1183
NON_CAUSAL_BACKENDS_TO_TEST = [
1184
AttentionBackendEnum.FLASH_ATTN,
1185
AttentionBackendEnum.FLEX_ATTENTION,
1186
+ "FLEX_ATTENTION_SLOW",
1187
]
1188
1189
if current_platform.is_rocm():
tests/v1/attention/test_mla_backends.py
+15
-16
@@ -62,6 +62,9 @@ BACKENDS_TO_TEST = [
62
AttentionBackendEnum.TOKENSPEED_MLA,
63
]
64
65
+if current_platform.is_rocm():
66
+ BACKENDS_TO_TEST.append(AttentionBackendEnum.ROCM_AITER_MLA)
67
+
68
DEVICE_TYPE = current_platform.device_type
69
70
@@ -201,13 +204,18 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
204
205
206
# Validate parameter combinations during collection, before GPU fixtures run.
204
-PREFILL_BACKENDS_TO_TEST = [
205
- MLAPrefillBackendEnum.ROCM_AITER_FA,
206
- MLAPrefillBackendEnum.FLASH_ATTN,
207
- MLAPrefillBackendEnum.FLASHINFER,
208
- MLAPrefillBackendEnum.TRTLLM_RAGGED,
209
- MLAPrefillBackendEnum.TOKENSPEED_MLA,
210
-]
207
+PREFILL_BACKENDS_TO_TEST: list[MLAPrefillBackendEnum] = []
208
+if current_platform.is_cuda():
209
+ PREFILL_BACKENDS_TO_TEST.extend(
210
+ [
211
+ MLAPrefillBackendEnum.FLASH_ATTN,
212
+ MLAPrefillBackendEnum.FLASHINFER,
213
+ MLAPrefillBackendEnum.TRTLLM_RAGGED,
214
+ MLAPrefillBackendEnum.TOKENSPEED_MLA,
215
+ ]
216
+ )
217
+elif current_platform.is_rocm():
218
+ PREFILL_BACKENDS_TO_TEST.append(MLAPrefillBackendEnum.ROCM_AITER_FA)
219
220
MLA_DIMENSIONS_TO_TEST = [
221
("deepseek", 128, 128),
@@ -250,15 +258,6 @@ def _prefill_backend_dimension_params():
258
)
259
)
260
)
253
- elif (
254
- current_platform.is_rocm()
255
- and prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN
256
- ):
257
- # Pre-existing on main: flash-attn MLA prefill returns wrong
258
- # outputs and faults on ROCm, so it fails on the base commit too.
259
- marks.append(
260
- pytest.mark.skip(reason="FLASH_ATTN MLA prefill broken on ROCm")
261
- )
261
params.append(
262
pytest.param(
263
prefill_backend,