From 2067297b2f15f0776be6d2a7e1b435bd20af60c1 Mon Sep 17 00:00:00 2001 From: afriedri Date: Sat, 17 Oct 2026 04:00:28 -0500 Subject: [PATCH] [ROCm][Perf] Use AITER tuned configs to select a4w4 gemm kernel instead of env var (#61278) Signed-off-by: Andy Friedrich Signed-off-by: afriedri Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: Kevin H. Luu Co-authored-by: Andreas Karatzas --- tests/kernels/quantization/test_rocm_mxfp4.py | 71 +++++++++++-------- vllm/_aiter_ops.py | 22 ++++-- vllm/envs.py | 6 -- .../kernels/linear/mxfp4/aiter.py | 10 ++- 4 files changed, 66 insertions(+), 43 deletions(-) diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py index f83445b39c..2ff4415465 100644 --- a/tests/kernels/quantization/test_rocm_mxfp4.py +++ b/tests/kernels/quantization/test_rocm_mxfp4.py @@ -147,26 +147,15 @@ def test_fp4_env_defaults(): """ROCm FP4 env defaults should stay stable for the AITER gates.""" import vllm.envs as envs - assert envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM is False assert envs.VLLM_ROCM_USE_AITER_FP4BMM is False @pytest.mark.parametrize( - ( - "use_aiter", - "use_fp4_asm_gemm", - "use_fp4bmm", - ), - [ - (True, True, True), - (True, True, False), - (True, False, True), - (False, True, True), - ], + ("use_aiter", "use_fp4bmm"), + [(True, True), (True, False), (False, True)], ) def test_rocm_aiter_fp4_enablement_follows_env_and_arch( use_aiter, - use_fp4_asm_gemm, use_fp4bmm, monkeypatch, ): @@ -182,22 +171,16 @@ def test_rocm_aiter_fp4_enablement_follows_env_and_arch( _assert_aiter_supported() on_gfx950_value = on_gfx950() - expected_asm_gemm = use_aiter and use_fp4_asm_gemm and on_gfx950_value + expected_asm_gemm = use_aiter and on_gfx950_value expected_fp4bmm = use_aiter and use_fp4bmm and on_gfx950_value with monkeypatch.context() as mp: mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0") - mp.setenv( - "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", - "1" if use_fp4_asm_gemm else "0", - ) mp.setenv("VLLM_ROCM_USE_AITER_FP4BMM", "1" if use_fp4bmm else "0") _reload_envs() rocm_aiter_ops.refresh_env_variables() - assert ( - rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() is expected_asm_gemm - ) + assert rocm_aiter_ops.is_asm_fp4_gemm_supported() is expected_asm_gemm assert rocm_aiter_ops.is_fp4bmm_enabled() is expected_fp4bmm @@ -695,24 +678,52 @@ def test_asm_fp4_scale_swizzle_rejects_non_2d(): @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") -def test_aiter_mxfp4_process_weights_falls_back_to_triton_for_misaligned_scale(): - """A misaligned weight_scale must disable ASM and take the Triton path.""" +@pytest.mark.parametrize( + ("asm_tuned", "preshuffle_tuned", "weight_shape", "scale_shape", "expected"), + [ + (True, True, (1280, 4096), (1280, 256), True), + (True, False, (1280, 4096), (1280, 256), False), + (False, True, (1280, 4096), (1280, 256), False), + # Tuned, but the scale can't be swizzled for ASM (columns % 8 != 0). + (True, True, (2560, 160), (2560, 10), False), + ], + ids=["tuned", "preshuffle_untuned", "asm_untuned", "misaligned_scale"], +) +def test_aiter_mxfp4_asm_layout_follows_aiter_tuned_configs( + asm_tuned, preshuffle_tuned, weight_shape, scale_shape, expected, monkeypatch +): + """A layer takes the ASM layout only when aiter has tuned configs for both + the ASM GEMM and the preshuffled Triton GEMM, and the scale can be swizzled. + Otherwise it stays on the Triton layout.""" from torch.nn.parameter import Parameter + from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.kernels.linear.mxfp4.aiter import AiterMxfp4LinearKernel + monkeypatch.setattr(rocm_aiter_ops, "is_asm_fp4_gemm_tuned", lambda n, k: asm_tuned) + monkeypatch.setattr( + rocm_aiter_ops, + "is_triton_gemm_afp4wfp4_presh_ws_tuned", + lambda n, k: preshuffle_tuned, + ) kernel = object.__new__(AiterMxfp4LinearKernel) - kernel.use_asm_gemm = True + kernel.asm_gemm_supported = True + kernel.use_asm_gemm = False kernel.out_dtype = torch.bfloat16 layer = torch.nn.Module() + layer.weight = Parameter( + torch.zeros(weight_shape, dtype=torch.uint8, device="cuda"), + requires_grad=False, + ) layer.weight_scale = Parameter( - torch.zeros(2560, 10, dtype=torch.uint8, device="cpu"), requires_grad=False + torch.zeros(scale_shape, dtype=torch.uint8, device="cuda"), + requires_grad=False, ) - kernel.process_weights_after_loading(layer) - assert kernel.use_asm_gemm is False - # Triton path stores the transposed, contiguous scale. - assert tuple(layer.weight_scale.shape) == (10, 2560) - assert layer.weight_scale.is_contiguous() + assert kernel.use_asm_gemm is expected + if not expected: + # Triton path stores the transposed, contiguous scale. + assert tuple(layer.weight_scale.shape) == scale_shape[::-1] + assert layer.weight_scale.is_contiguous() diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 29364f204a..7b05fb4ac5 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2124,7 +2124,6 @@ class rocm_aiter_ops: VLLM_ROCM_USE_AITER_MHA: Controls MHA ops including flash_attn_varlen. VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: Controls Triton unified attention. VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply. - VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM. VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings. VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: Controls Triton sparse MLA (gfx950). VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion. @@ -2196,8 +2195,6 @@ class rocm_aiter_ops: _FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM _FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM _LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM - # TODO: Consolidate under _LINEAR_ENABLED - _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE _TRITON_SPARSE_MLA = envs.VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA @@ -2228,7 +2225,6 @@ class rocm_aiter_ops: cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM cls._FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM - cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE cls._TRITON_SPARSE_MLA = envs.VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS @@ -2565,10 +2561,10 @@ class rocm_aiter_ops: @classmethod @if_aiter_supported - def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool: + def is_asm_fp4_gemm_supported(cls) -> bool: from vllm.platforms.rocm import on_gfx950 - return cls._AITER_ENABLED and cls._FP4_GEMM_DYNAMIC_QUANT_ASM and on_gfx950() + return cls._AITER_ENABLED and on_gfx950() @classmethod @if_aiter_supported @@ -4360,6 +4356,20 @@ class rocm_aiter_ops: except (AssertionError, ImportError): return False + @staticmethod + @functools.cache + def is_asm_fp4_gemm_tuned(n: int, k_bytes: int) -> bool: + if not current_platform.is_rocm(): + return False + try: + from aiter.ops.gemm_op_a4w4 import get_GEMM_config + + return all( + get_GEMM_config(m, n, 2 * k_bytes) is not None for m in (1024, 8192) + ) + except ImportError: + return False + @staticmethod @functools.cache def is_shuffled_per_token_w8a8_gemm_tuned( diff --git a/vllm/envs.py b/vllm/envs.py index 5cd1d3eecb..0492c5a59d 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -146,7 +146,6 @@ if TYPE_CHECKING: VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto" VLLM_ROCM_AITER_MLA_DCP_VERIFY: Literal["auto", "asm", "segmented"] = "auto" VLLM_ROCM_USE_AITER_MHA: bool = True - 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 = False @@ -1336,11 +1335,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ROCM_USE_AITER_MHA": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1") ), - # Whether to use aiter fp4 gemm asm. - # By default is disabled. - "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: ( - os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1") - ), # Whether sparse MLA prefill and decode run on aiter's Triton kernel, which # reads the KV cache as stored (bf16 or fp8, paged or flat). Used by the # ROCM_AITER_MLA_SPARSE backend (DeepSeek V3.2, GLM-5.x) and DeepSeek diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py index 4cd09c1f90..938e04fff5 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/aiter.py +++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py @@ -159,7 +159,8 @@ class AiterMxfp4LinearKernel(MxFp4LinearKernel): def __init__(self, config: MxFp4LinearLayerConfig) -> None: super().__init__(config) - self.use_asm_gemm = rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() + self.asm_gemm_supported = rocm_aiter_ops.is_asm_fp4_gemm_supported() + self.use_asm_gemm = False self.out_dtype = torch.get_default_dtype() @classmethod @@ -199,6 +200,13 @@ class AiterMxfp4LinearKernel(MxFp4LinearKernel): return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Use preshuffled ASM (P) and Triton (D) GEMM when tuned in aiter + if self.asm_gemm_supported: + n, k_bytes = layer.weight.shape + self.use_asm_gemm = rocm_aiter_ops.is_asm_fp4_gemm_tuned( + n, k_bytes + ) and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(n, k_bytes) + if self.use_asm_gemm and not _asm_fp4_scale_swizzle_supported( layer.weight_scale.data ):