[ROCm][CI] Run the MLA attention+quant fusion test on ROCm (#58717)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Sep 25, 2026 at 21:45 UTC
31cc2264011a7f9dbbea4d6e13f696cb0eee5282
1 file changed
+16
-1
tests/compile/passes/test_mla_attn_quant_fusion.py
+16
-1
@@ -361,13 +361,27 @@ if current_platform.is_cuda():
361
]
362
BACKENDS_MLA_FP8 = [AttentionBackendEnum.TRITON_MLA]
363
BACKENDS_MLA_FP4 = [AttentionBackendEnum.TRITON_MLA]
364
+elif current_platform.is_rocm():
365
+ # ROCm supports the static-FP8 output-quant fusion through AITER MLA.
366
+ # Per-group FP8 still relies on the CUDA-only CUTLASS block-scaled
367
+ # kernel, and NVFP4 is NVIDIA-specific.
368
+ MLA_DIMS = [(16, 128, 64, 128, 512)]
369
+ PATTERN_TEST_MODELS_MLA_FP8 = [
370
+ (
371
+ "deepseek-ai/DeepSeek-V2-Lite",
372
+ TestMLAAttentionFp8StaticQuantPatternModel,
373
+ )
374
+ ]
375
+ BACKENDS_MLA_FP8 = [AttentionBackendEnum.ROCM_AITER_MLA]
376
377
378
@pytest.mark.parametrize(
379
"num_heads, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, kv_lora_rank",
380
MLA_DIMS,
381
)
370
-@pytest.mark.parametrize("batch_size", [7, 256] if current_platform.is_cuda() else [8])
382
+@pytest.mark.parametrize(
383
+ "batch_size", [7, 256] if current_platform.is_cuda_alike() else [8]
384
+)
385
@pytest.mark.parametrize("dtype", [torch.bfloat16])
386
@pytest.mark.parametrize(
387
"backend, model_name, model_class, custom_ops",
@@ -405,6 +419,7 @@ def test_mla_attention_quant_pattern(
419
backend: AttentionBackendEnum,
420
dist_init,
421
disable_vllm_compile_cache,
422
+ workspace_init,
423
):
424
"""Test MLA AttentionQuantPattern fusion pass."""
425
if (