iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Perf] Use AITER tuned configs to select a4w4 gemm kernel instead of env var (#61278)

Browse Source 
Signed-off-by: Andy Friedrich <afriedri@amd.com>
Signed-off-by: afriedri <afriedri@amd.com>
Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
afriedri
2026-10-17 04:00:28 -05:00
committed by GitHub
co-authored by Simon Danielsson Kevin H. Luu Andreas Karatzas
parent 3bf25c5097
commit 2067297b2f
4 changed files with 66 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/quantization/test_rocm_mxfp4.py
+41 -30
View File
@@ -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()
vllm/_aiter_ops.py
+16 -6
View File
@@ -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(
vllm/envs.py
-6
View File
@@ -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
vllm/model_executor/kernels/linear/mxfp4/aiter.py
+9 -1
View File
@@ -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
):