From 81503550f379b5a83bc091073bc2a2f2ebe67ff5 Mon Sep 17 00:00:00 2001 From: xuebwang Date: Wed, 14 Oct 2026 18:46:02 +0800 Subject: [PATCH] [ROCm] Disable AITER FP8/FP4 BMM for MLA by default (#53123) Signed-off-by: xuebwang-amd --- tests/kernels/quantization/test_rocm_fp8.py | 2 +- tests/kernels/quantization/test_rocm_mxfp4.py | 2 +- vllm/envs.py | 17 +++++++++-------- .../layers/attention/mla_attention.py | 5 +++++ 4 files changed, 16 insertions(+), 10 deletions(-) diff --git a/tests/kernels/quantization/test_rocm_fp8.py b/tests/kernels/quantization/test_rocm_fp8.py index 47e37df119..9688b15d4d 100644 --- a/tests/kernels/quantization/test_rocm_fp8.py +++ b/tests/kernels/quantization/test_rocm_fp8.py @@ -30,7 +30,7 @@ fp8_only = pytest.mark.skipif( ENV_DEFAULTS = { "VLLM_ROCM_FP8_PADDING": True, - "VLLM_ROCM_USE_AITER_FP8BMM": True, + "VLLM_ROCM_USE_AITER_FP8BMM": False, "VLLM_ROCM_FP8_MFMA_PAGE_ATTN": False, } diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py index b2547c080a..fa39c641ba 100644 --- a/tests/kernels/quantization/test_rocm_mxfp4.py +++ b/tests/kernels/quantization/test_rocm_mxfp4.py @@ -148,7 +148,7 @@ def test_fp4_env_defaults(): import vllm.envs as envs assert envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM is False - assert envs.VLLM_ROCM_USE_AITER_FP4BMM is True + assert envs.VLLM_ROCM_USE_AITER_FP4BMM is False @pytest.mark.parametrize( diff --git a/vllm/envs.py b/vllm/envs.py index dbbc49ff59..4325b48823 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -153,8 +153,8 @@ if TYPE_CHECKING: VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: bool = False VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False - VLLM_ROCM_USE_AITER_FP8BMM: bool = True - VLLM_ROCM_USE_AITER_FP4BMM: bool = True + VLLM_ROCM_USE_AITER_FP8BMM: bool = False + VLLM_ROCM_USE_AITER_FP4BMM: bool = False VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS: bool = False @@ -1381,15 +1381,16 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_TRITON_ROPE", "False").lower() in ("true", "1") ), - # Whether to use aiter triton fp8 bmm kernel - # By default is enabled. + # Whether to use aiter triton fp8 bmm kernel for MLA kv_b_proj. + # By default is disabled: it re-quantizes the projection weights online, + # regardless of the precision the checkpoint asked for. "VLLM_ROCM_USE_AITER_FP8BMM": lambda: ( - os.getenv("VLLM_ROCM_USE_AITER_FP8BMM", "True").lower() in ("true", "1") + os.getenv("VLLM_ROCM_USE_AITER_FP8BMM", "False").lower() in ("true", "1") ), - # Whether to use aiter triton fp4 bmm kernel - # By default is enabled. + # Whether to use aiter triton fp4 bmm kernel for MLA kv_b_proj. + # By default is disabled, for the same reason as the fp8 variant above. "VLLM_ROCM_USE_AITER_FP4BMM": lambda: ( - os.getenv("VLLM_ROCM_USE_AITER_FP4BMM", "True").lower() in ("true", "1") + os.getenv("VLLM_ROCM_USE_AITER_FP4BMM", "False").lower() in ("true", "1") ), # Use AITER triton unified attention for V1 attention "VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION": lambda: ( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 806faba507..bc6bca99ff 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -709,6 +709,11 @@ class MLAAttention(nn.Module, AttentionLayerBase): use_pcp=self.use_pcp, ) + # TODO(ROCm): both BMM kernels re-quantize kv_b_proj online, ignoring + # the precision the checkpoint asked for, so they now default to off. + # Turning them back on needs long-context accuracy validation, and the + # per-layer decision belongs in --quantization-config.targets rather + # than in this layer. self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled() # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported