diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index 89244bed1b..0967f9499e 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -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, diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 9bea4f8b3a..8f2b2a24ad 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -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() diff --git a/vllm/envs.py b/vllm/envs.py index 6158f86f89..78af9bcb72 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -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 diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index cb4225c92a..62df395f69 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -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: diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 3b415acd81..1c38e4378b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -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