From 31cc2264011a7f9dbbea4d6e13f696cb0eee5282 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Fri, 25 Sep 2026 21:45:32 +0200 Subject: [PATCH] [ROCm][CI] Run the MLA attention+quant fusion test on ROCm (#58717) Signed-off-by: Stefan Koncarevic Co-authored-by: Andreas Karatzas Co-authored-by: Andreas Karatzas --- .../passes/test_mla_attn_quant_fusion.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index 56d2bc30cc..a945f00079 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -361,13 +361,27 @@ if current_platform.is_cuda(): ] BACKENDS_MLA_FP8 = [AttentionBackendEnum.TRITON_MLA] BACKENDS_MLA_FP4 = [AttentionBackendEnum.TRITON_MLA] +elif current_platform.is_rocm(): + # ROCm supports the static-FP8 output-quant fusion through AITER MLA. + # Per-group FP8 still relies on the CUDA-only CUTLASS block-scaled + # kernel, and NVFP4 is NVIDIA-specific. + MLA_DIMS = [(16, 128, 64, 128, 512)] + PATTERN_TEST_MODELS_MLA_FP8 = [ + ( + "deepseek-ai/DeepSeek-V2-Lite", + TestMLAAttentionFp8StaticQuantPatternModel, + ) + ] + BACKENDS_MLA_FP8 = [AttentionBackendEnum.ROCM_AITER_MLA] @pytest.mark.parametrize( "num_heads, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, kv_lora_rank", MLA_DIMS, ) -@pytest.mark.parametrize("batch_size", [7, 256] if current_platform.is_cuda() else [8]) +@pytest.mark.parametrize( + "batch_size", [7, 256] if current_platform.is_cuda_alike() else [8] +) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize( "backend, model_name, model_class, custom_ops", @@ -405,6 +419,7 @@ def test_mla_attention_quant_pattern( backend: AttentionBackendEnum, dist_init, disable_vllm_compile_cache, + workspace_init, ): """Test MLA AttentionQuantPattern fusion pass.""" if (