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

[ROCm][Kimi-K3] Enable a4w4 flydsl kernels for KimiK3 (#53940)

Browse Source 
Signed-off-by: ppalanga <ppalanga@amd.com>
Signed-off-by: Poovaiah Palangappa <poovaiah.palangappa@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Poovaiah Palangappa <poovaiah.palangappa@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com>
This commit is contained in:
ppalanga
2026-09-16 01:38:02 +00:00
committed by GitHub
co-authored by Cursor Poovaiah Palangappa mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Hongxia Yang
parent c8d1cf077a
commit a4d2d9d95e
5 changed files with 127 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_rocm_aiter_moe.py
+68
View File
@@ -565,6 +565,74 @@ def test_aiter_moe_shared_experts_enablement_follows_env(
assert rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() is expected
def test_aiter_moe_situv2_syncs_aiter_a4w4_env(
monkeypatch: pytest.MonkeyPatch,
):
"""SiTUv2 MoE should route to AITER a4w4 and clear legacy a8w4 overrides."""
from vllm._aiter_ops import rocm_aiter_ops
_assert_aiter_supported()
with monkeypatch.context() as mp:
mp.delenv("AITER_SITUV2_A8W4", raising=False)
mp.delenv("AITER_SITUV2_A4W4", raising=False)
mp.setenv("VLLM_ROCM_USE_AITER", "1")
mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "1")
_reload_envs()
rocm_aiter_ops.refresh_env_variables()
import os
assert os.environ.get("AITER_SITUV2_A4W4") == "1"
assert "AITER_SITUV2_A8W4" not in os.environ
def test_aiter_moe_situv2_legacy_a8w4_alias_enables_a4w4(
monkeypatch: pytest.MonkeyPatch,
):
"""Deprecated VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 still enables SiTUv2 a4w4."""
from vllm._aiter_ops import rocm_aiter_ops
_assert_aiter_supported()
with monkeypatch.context() as mp:
mp.delenv("AITER_SITUV2_A8W4", raising=False)
mp.delenv("AITER_SITUV2_A4W4", raising=False)
mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", raising=False)
mp.setenv("VLLM_ROCM_USE_AITER", "1")
mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "1")
_reload_envs()
rocm_aiter_ops.refresh_env_variables()
import os
assert os.environ.get("AITER_SITUV2_A4W4") == "1"
assert "AITER_SITUV2_A8W4" not in os.environ
def test_aiter_moe_situv2_clears_aiter_env_when_disabled(
monkeypatch: pytest.MonkeyPatch,
):
from vllm._aiter_ops import rocm_aiter_ops
_assert_aiter_supported()
with monkeypatch.context() as mp:
mp.setenv("AITER_SITUV2_A4W4", "1")
mp.setenv("VLLM_ROCM_USE_AITER", "1")
mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", raising=False)
mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0")
_reload_envs()
rocm_aiter_ops.refresh_env_variables()
import os
assert "AITER_SITUV2_A4W4" not in os.environ
@pytest.mark.parametrize("moe_padding", [True, False])
def test_aiter_moe_padding_env_var(
moe_padding: bool,
vllm/_aiter_ops.py
+32 -6
View File
@@ -1752,6 +1752,30 @@ def _mhc_delayed_pre_tail(
_OPS_REGISTERED = False
def _sync_aiter_situv2_moe_env() -> None:
"""Mirror the SiTUv2 MoE toggle into AITER's a4w4 dispatch env.
AITER selects afp8 vs afp4 activation kernels via AITER_SITUV2_A8W4 /
AITER_SITUV2_A4W4 (see ROCm/aiter fused_moe.py, A8W4 checked first).
When VLLM_ROCM_USE_AITER_MOE_SITUV2 is enabled we route to a4w4
(afp4_wfp4_fp4 kernels) and clear any legacy AITER_SITUV2_A8W4 override.
Requires AITER with ROCm/aiter#4463 (first tagged in v0.1.20): a4w4
dispatch plus kimik3_a4w4_{un,}tuned_fmoe.csv. Older AITER still runs
afp4 FlyDSL but falls back to heuristic configs, which is not the
tuned a4w4 path this recipe is meant to use.
"""
import os
import vllm.envs as envs
if envs.VLLM_ROCM_USE_AITER_MOE_SITUV2:
os.environ["AITER_SITUV2_A4W4"] = "1"
os.environ.pop("AITER_SITUV2_A8W4", None)
else:
os.environ.pop("AITER_SITUV2_A4W4", None)
class rocm_aiter_ops:
"""ROCm AITER operations wrapper for AMD GPU acceleration in vLLM.
@@ -1774,7 +1798,7 @@ class rocm_aiter_ops:
VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply.
VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.
VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.
VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: Controls a8w4 SiTU fused MoE variant.
VLLM_ROCM_USE_AITER_MOE_SITUV2: Controls SiTUv2 FlyDSL MoE (a4w4).
VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM.
Note:
@@ -1841,7 +1865,7 @@ class rocm_aiter_ops:
# TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE
_TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
_MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
_MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4
_MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2
# TODO: Consolidate under _LINEAR_ENABLED
_TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
# Lazily probed: whether aiter.topk_softmax supports the
@@ -1870,7 +1894,8 @@ class rocm_aiter_ops:
cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
cls._MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4
cls._MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2
_sync_aiter_situv2_moe_env()
cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
@staticmethod
@@ -1983,10 +2008,10 @@ class rocm_aiter_ops:
@classmethod
@if_aiter_supported
def is_fused_moe_situv2_a8w4_enabled(cls) -> bool:
# _MOE_SITUV2_A8W4 is a variant of aiter fused moe, so aiter
def is_fused_moe_situv2_enabled(cls) -> bool:
# _MOE_SITUV2 is a variant of aiter fused moe, so aiter
# fused moe must be enabled as well.
return cls.is_fused_moe_enabled() and cls._MOE_SITUV2_A8W4
return cls.is_fused_moe_enabled() and cls._MOE_SITUV2
@classmethod
@if_aiter_supported
@@ -4029,3 +4054,4 @@ class rocm_aiter_ops:
rocm_aiter_ops.register_ops_once()
_sync_aiter_situv2_moe_env()
vllm/envs.py
+13 -7
View File
@@ -137,7 +137,7 @@ if TYPE_CHECKING:
VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False
VLLM_ROCM_USE_AITER_MOE: bool = True
VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0
VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: bool = False
VLLM_ROCM_USE_AITER_MOE_SITUV2: bool = False
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
VLLM_ROCM_USE_AITER_MLA: bool = True
VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto"
@@ -1251,12 +1251,18 @@ environment_variables: dict[str, Callable[[], Any]] = {
"VLLM_ROCM_USE_AITER_MOE": lambda: (
os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1")
),
# Route K3 SiTU MXFP4 MoE through the a8w4 (fp8 activation) gate/up-
# interleaved flydsl kernels instead of the default a16w4 separated path.
# This is the only flag users need: vLLM picks the kernels by passing
# gate_mode to AITER and sets the AITER-side workaround env at init.
"VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4": lambda: (
os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower() in ("true", "1")
# Route K3 SiTU MXFP4 MoE through the FlyDSL SiTUv2 path (a4w4 fp4
# activations, separated gate/up layout) instead of default a16w4. vLLM
# sets AITER_SITUV2_A4W4 at init when this flag is on and clears any
# legacy AITER_SITUV2_A8W4 override (AITER checks A8W4 first).
# VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 is a deprecated alias for existing
# recipes; it does not select a8w4 kernels.
# Needs AITER >= v0.1.20 (ROCm/aiter#4463) for the a4w4 dispatch flag
# and tuned kimik3_a4w4_*_fmoe.csv rows; otherwise FlyDSL uses heuristics.
"VLLM_ROCM_USE_AITER_MOE_SITUV2": lambda: (
os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0").lower() in ("true", "1")
or os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower()
in ("true", "1")
),
# MoE sorting dispatch policy for AITER fused MoE kernels.
# 0 = auto (default): single-pass for small batches, multi-pass
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+4 -8
View File
@@ -385,14 +385,10 @@ def rocm_aiter_fused_experts(
gate_mode = ""
if activation == MoEActivation.SITU:
# a8w4 (VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1) uses the gate/up-
# interleaved (_gui_) fp8 flydsl kernels; default a16w4 SiTU stays
# separated.
gate_mode = (
GateMode.INTERLEAVE.value
if rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled()
else GateMode.SEPARATED.value
)
# SiTUv2 flydsl (VLLM_ROCM_USE_AITER_MOE_SITUV2=1) uses a4w4
# fp4 activations with separated gate/up weights (AITER #4463);
# default a16w4 SiTU also stays separated.
gate_mode = GateMode.SEPARATED.value
elif quant_config.use_mxfp4_w4a16:
gate_mode = GateMode.INTERLEAVE.value
elif activation_interleave is not None:
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+10 -9
View File
@@ -1572,12 +1572,6 @@ def convert_weight_to_mxfp4_moe_kernel_format(
if w2_bias is not None:
w2_bias = w2_bias.data.to(torch.float32)
import os
# TODO: Remove this once AITER is fixed
# Necessary for AITER side from crashing
os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"
if activation == MoEActivation.SITU:
from aiter.utility.fp4_utils import e8m0_shuffle
@@ -1585,9 +1579,8 @@ def convert_weight_to_mxfp4_moe_kernel_format(
fp4_dtype = torch.float4_e2m1fn_x2
e8m0_dtype = torch.float8_e8m0fnu
# a8w4 uses gate/up-interleaved flydsl kernels;
# default a16w4 keeps the separated layout.
guinterleave = rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled()
# SiTUv2 flydsl uses separated gate/up layout (a4w4).
guinterleave = False
w13 = rocm_aiter_ops.shuffle_weight_a16w4(
w13_weight.data.view(fp4_dtype), 16, guinterleave
)
@@ -1606,6 +1599,14 @@ def convert_weight_to_mxfp4_moe_kernel_format(
w2.is_shuffled = True
return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)
import os
# Interleaved a16w4 only (DeepSeekV4 etc.). AITER uses this bound to
# pick bf16 vs fp8 activations when gate_mode is INTERLEAVE. SiTUv2
# a4w4 is separated and selects q_dtype_a independently, so the bound
# is unused on that path.
os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"
from aiter.ops.shuffle import shuffle_scale as _shuf_s
from aiter.ops.shuffle import shuffle_weight as _shuf_w