[ROCm] Restore `VLLM_ROCM_USE_AITER_FP4_ASM_GEMM` and default w4a4 ASM GEMM back to off (#57055)
Signed-off-by: Andy Friedrich <afriedri@amd.com>
afriedri committed
Sep 17, 2026 at 13:39 UTC
a9a7e45f3174781f2b20c03c7ac1beee2b98f9d5
3 files changed
+23
-5
tests/kernels/quantization/test_rocm_mxfp4.py
+12
-4
@@ -156,22 +156,26 @@ def test_fp4_env_defaults():
156
"""ROCm FP4 env defaults should stay stable for the AITER gates."""
157
import vllm.envs as envs
158
159
+ assert envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM is False
160
assert envs.VLLM_ROCM_USE_AITER_FP4BMM is True
161
162
163
@pytest.mark.parametrize(
164
(
165
"use_aiter",
166
+ "use_fp4_asm_gemm",
167
"use_fp4bmm",
168
),
169
[
168
- (True, True),
169
- (True, False),
170
- (False, True),
170
+ (True, True, True),
171
+ (True, True, False),
172
+ (True, False, True),
173
+ (False, True, True),
174
],
175
)
176
def test_rocm_aiter_fp4_enablement_follows_env_and_arch(
177
use_aiter,
178
+ use_fp4_asm_gemm,
179
use_fp4bmm,
180
monkeypatch,
181
):
@@ -187,11 +191,15 @@ def test_rocm_aiter_fp4_enablement_follows_env_and_arch(
191
_assert_aiter_supported()
192
193
on_gfx950_value = on_gfx950()
190
- expected_asm_gemm = use_aiter and on_gfx950_value
194
+ expected_asm_gemm = use_aiter and use_fp4_asm_gemm and on_gfx950_value
195
expected_fp4bmm = use_aiter and use_fp4bmm and on_gfx950_value
196
197
with monkeypatch.context() as mp:
198
mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
199
+ mp.setenv(
200
+ "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM",
201
+ "1" if use_fp4_asm_gemm else "0",
202
+ )
203
mp.setenv("VLLM_ROCM_USE_AITER_FP4BMM", "1" if use_fp4bmm else "0")
204
_reload_envs()
205
rocm_aiter_ops.refresh_env_variables()
vllm/_aiter_ops.py
+5
-1
@@ -1893,6 +1893,7 @@ class rocm_aiter_ops:
1893
VLLM_ROCM_USE_AITER_MHA: Controls MHA ops including flash_attn_varlen.
1894
VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: Controls Triton unified attention.
1895
VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply.
1896
+ VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM.
1897
VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.
1898
VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.
1899
VLLM_ROCM_USE_AITER_MOE_SITUV2: Controls SiTUv2 FlyDSL MoE (a4w4).
@@ -1960,6 +1961,8 @@ class rocm_aiter_ops:
1961
_FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
1962
_FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
1963
_LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
1964
+ # TODO: Consolidate under _LINEAR_ENABLED
1965
+ _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
1966
# TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE
1967
_TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
1968
_MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
@@ -1989,6 +1992,7 @@ class rocm_aiter_ops:
1992
cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
1993
cls._FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
1994
cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
1995
+ cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
1996
cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
1997
cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
1998
cls._MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2
@@ -2230,7 +2234,7 @@ class rocm_aiter_ops:
2234
def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool:
2235
from vllm.platforms.rocm import on_gfx950
2236
2233
- return cls._AITER_ENABLED and on_gfx950()
2237
+ return cls._AITER_ENABLED and cls._FP4_GEMM_DYNAMIC_QUANT_ASM and on_gfx950()
2238
2239
@classmethod
2240
@if_aiter_supported
vllm/envs.py
+6
@@ -143,6 +143,7 @@ if TYPE_CHECKING:
143
VLLM_ROCM_USE_AITER_MLA: bool = True
144
VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto"
145
VLLM_ROCM_USE_AITER_MHA: bool = True
146
+ VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False
147
VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False
148
VLLM_ROCM_USE_AITER_FP8BMM: bool = True
149
VLLM_ROCM_USE_AITER_FP4BMM: bool = True
@@ -1305,6 +1306,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1306
"VLLM_ROCM_USE_AITER_MHA": lambda: (
1307
os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1")
1308
),
1309
+ # Whether to use aiter fp4 gemm asm.
1310
+ # By default is disabled.
1311
+ "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: (
1312
+ os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1")
1313
+ ),
1314
# Whether to use aiter rope.
1315
# By default is disabled.
1316
"VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: (