[ROCm] remove VLLM_ROCM_USE_AITER_FP4_ASM_GEMM environment variable; make w4a4 use the preshuffle triton+asm by default (#53141)
Signed-off-by: Andy Friedrich <afriedri@amd.com> Signed-off-by: afriedri <afriedri@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
afriedri committed
Aug 28, 2026 at 13:59 UTC
2f1cba799e02209e9782a747f7620e55d8251e67
4 files changed
+9
-27
tests/kernels/quantization/test_rocm_mxfp4.py
+6
-14
@@ -147,26 +147,22 @@ def test_fp4_env_defaults():
147
"""ROCm FP4 env defaults should stay stable for the AITER gates."""
148
import vllm.envs as envs
149
150
- assert envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM is False
150
assert envs.VLLM_ROCM_USE_AITER_FP4BMM is True
151
152
153
@pytest.mark.parametrize(
154
(
155
"use_aiter",
157
- "use_fp4_asm_gemm",
156
"use_fp4bmm",
157
),
158
[
161
- (True, True, True),
162
- (True, True, False),
163
- (True, False, True),
164
- (False, True, True),
159
+ (True, True),
160
+ (True, False),
161
+ (False, True),
162
],
163
)
164
def test_rocm_aiter_fp4_enablement_follows_env_and_arch(
165
use_aiter,
169
- use_fp4_asm_gemm,
166
use_fp4bmm,
167
monkeypatch,
168
):
@@ -182,15 +178,11 @@ def test_rocm_aiter_fp4_enablement_follows_env_and_arch(
178
_assert_aiter_supported()
179
180
on_gfx950_value = on_gfx950()
185
- expected_asm_gemm = use_aiter and use_fp4_asm_gemm and on_gfx950_value
181
+ expected_asm_gemm = use_aiter and on_gfx950_value
182
expected_fp4bmm = use_aiter and use_fp4bmm and on_gfx950_value
183
184
with monkeypatch.context() as mp:
185
mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
190
- mp.setenv(
191
- "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM",
192
- "1" if use_fp4_asm_gemm else "0",
193
- )
186
mp.setenv("VLLM_ROCM_USE_AITER_FP4BMM", "1" if use_fp4bmm else "0")
187
_reload_envs()
188
rocm_aiter_ops.refresh_env_variables()
@@ -410,7 +402,7 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape):
402
from aiter import per_1x32_f4_quant_hip
403
from aiter.ops.shuffle import shuffle_weight
404
from aiter.ops.triton.gemm_afp4wfp4 import (
413
- gemm_afp4wfp4_preshuffled_weight_scales,
405
+ gemm_afp4wfp4_preshuffle,
406
)
407
from aiter.ops.triton.quant import dynamic_mxfp4_quant
408
@@ -446,7 +438,7 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape):
438
439
def run_preshuffled() -> torch.Tensor:
440
y = torch.empty(M, N, device="cuda", dtype=torch.bfloat16)
449
- return gemm_afp4wfp4_preshuffled_weight_scales(
441
+ return gemm_afp4wfp4_preshuffle(
442
A_q.contiguous().view(torch.uint8),
443
B_fp4.contiguous().view(torch.uint8).reshape(B_fp4.shape[0] // 16, -1),
444
A_s,
vllm/_aiter_ops.py
+1
-5
@@ -1688,7 +1688,6 @@ class rocm_aiter_ops:
1688
VLLM_ROCM_USE_AITER_MHA: Controls MHA ops including flash_attn_varlen.
1689
VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: Controls Triton unified attention.
1690
VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply.
1691
- VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM.
1691
VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.
1692
VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.
1693
VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: Controls a8w4 SiTU fused MoE variant.
@@ -1755,8 +1754,6 @@ class rocm_aiter_ops:
1754
_FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
1755
_FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
1756
_LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
1758
- # TODO: Consolidate under _LINEAR_ENABLED
1759
- _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
1757
# TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE
1758
_TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
1759
_MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
@@ -1787,7 +1784,6 @@ class rocm_aiter_ops:
1784
cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
1785
cls._FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
1786
cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
1790
- cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
1787
cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
1788
cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
1789
cls._MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4
@@ -2008,7 +2004,7 @@ class rocm_aiter_ops:
2004
def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool:
2005
from vllm.platforms.rocm import on_gfx950
2006
2011
- return cls._AITER_ENABLED and cls._FP4_GEMM_DYNAMIC_QUANT_ASM and on_gfx950()
2007
+ return cls._AITER_ENABLED and on_gfx950()
2008
2009
@classmethod
2010
@if_aiter_supported
vllm/envs.py
-6
@@ -142,7 +142,6 @@ if TYPE_CHECKING:
142
VLLM_ROCM_USE_AITER_MLA: bool = True
143
VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto"
144
VLLM_ROCM_USE_AITER_MHA: bool = True
145
- VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False
145
VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False
146
VLLM_ROCM_USE_AITER_FP8BMM: bool = True
147
VLLM_ROCM_USE_AITER_FP4BMM: bool = True
@@ -1300,11 +1299,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
1299
"VLLM_ROCM_USE_AITER_MHA": lambda: (
1300
os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1")
1301
),
1303
- # Whether to use aiter fp4 gemm asm.
1304
- # By default is disabled.
1305
- "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: (
1306
- os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1")
1307
- ),
1302
# Whether to use aiter rope.
1303
# By default is disabled.
1304
"VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: (
vllm/model_executor/kernels/linear/mxfp4/aiter.py
+2
-2
@@ -35,7 +35,7 @@ if is_aiter_found_and_supported():
35
) -> torch.Tensor:
36
from aiter.ops.triton.gemm_afp4wfp4 import (
37
gemm_afp4wfp4,
38
- gemm_afp4wfp4_preshuffled_weight_scales,
38
+ gemm_afp4wfp4_preshuffle,
39
)
40
from aiter.ops.triton.quant import dynamic_mxfp4_quant
41
@@ -63,7 +63,7 @@ if is_aiter_found_and_supported():
63
x_s = x_s[:M, ...].view(torch.uint8)
64
65
y = torch.empty(M, N, device=x_q.device, dtype=out_dtype)
66
- gemm_afp4wfp4_preshuffled_weight_scales(
66
+ gemm_afp4wfp4_preshuffle(
67
x_q.view(torch.uint8),
68
weight.view(torch.uint8).view(weight.shape[0] // 16, -1),
69
x_s,