@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3137ff0773

[ROCm][Kimi-K3] Make VLLM_ROCM_USE_AITER_MOE_SITUV2 select a4w4/a8w4/a16w4 (#58201)

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

Hongxia Yang committed Sep 26, 2026 at 22:26 UTC 3137ff07739f0cd19e06119506c6432a4a7e77ad

5 files changed +103 -88

tests/kernels/moe/test_rocm_aiter_moe.py

+38 -46

@@ -566,72 +566,64 @@ def test_aiter_moe_shared_experts_enablement_follows_env(

566 assert rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() is expected

567

568

569 -def test_aiter_moe_situv2_syncs_aiter_a4w4_env(

569 +@pytest.mark.parametrize(

570 + "value,expected_act,expected_env",

571 + [

572 + # Unset / auto / legacy "1" all mean the a4w4 default.

573 + (None, "a4w4", {"AITER_SITUV2_A4W4": "1"}),

574 + ("auto", "a4w4", {"AITER_SITUV2_A4W4": "1"}),

575 + ("1", "a4w4", {"AITER_SITUV2_A4W4": "1"}),

576 + ("a4w4", "a4w4", {"AITER_SITUV2_A4W4": "1"}),

577 + ("0", "a16w4", {}),

578 + ("a16w4", "a16w4", {}),

579 + ("A8W4", "a8w4", {"AITER_SITUV2_A8W4": "1"}),

580 + ],

581 +)

582 +def test_aiter_moe_situv2_activation_syncs_aiter_env(

583 + value: str | None,

584 + expected_act: str,

585 + expected_env: dict[str, str],

586 monkeypatch: pytest.MonkeyPatch,

587 ):

572 - """SiTUv2 MoE should route to AITER a4w4 and clear legacy a8w4 overrides."""

573 - from vllm._aiter_ops import rocm_aiter_ops

574 -

575 - _assert_aiter_supported()

576 -

577 - with monkeypatch.context() as mp:

578 - mp.delenv("AITER_SITUV2_A8W4", raising=False)

579 - mp.delenv("AITER_SITUV2_A4W4", raising=False)

580 - mp.setenv("VLLM_ROCM_USE_AITER", "1")

581 - mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")

582 - mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "1")

583 - _reload_envs()

584 - rocm_aiter_ops.refresh_env_variables()

585 -

586 - import os

587 -

588 - assert os.environ.get("AITER_SITUV2_A4W4") == "1"

589 - assert "AITER_SITUV2_A8W4" not in os.environ

590 -

588 + """Each SiTUv2 activation choice sets exactly its AITER dispatch env."""

589 + import os

590

592 -def test_aiter_moe_situv2_legacy_a8w4_alias_enables_a4w4(

593 - monkeypatch: pytest.MonkeyPatch,

594 -):

595 - """Deprecated VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 still enables SiTUv2 a4w4."""

591 from vllm._aiter_ops import rocm_aiter_ops

592

593 _assert_aiter_supported()

594

595 with monkeypatch.context() as mp:

601 - mp.delenv("AITER_SITUV2_A8W4", raising=False)

602 - mp.delenv("AITER_SITUV2_A4W4", raising=False)

603 - mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", raising=False)

596 + # Stale AITER overrides must be replaced by the vLLM selection.

597 + mp.setenv("AITER_SITUV2_A8W4", "1")

598 + mp.setenv("AITER_SITUV2_A4W4", "1")

599 mp.setenv("VLLM_ROCM_USE_AITER", "1")

600 mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")

606 - mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "1")

601 + if value is None:

602 + mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", raising=False)

603 + else:

604 + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", value)

605 _reload_envs()

606 rocm_aiter_ops.refresh_env_variables()

607

610 - import os

611 -

612 - assert os.environ.get("AITER_SITUV2_A4W4") == "1"

613 - assert "AITER_SITUV2_A8W4" not in os.environ

608 + assert rocm_aiter_ops.get_fused_moe_situv2_activation() == expected_act

609 + assert rocm_aiter_ops.is_fused_moe_situv2_gate_up_interleaved() is (

610 + expected_act == "a8w4"

611 + )

612 + assert rocm_aiter_ops.is_fused_moe_situv2_enabled() is (expected_act != "a16w4")

613 + for name in ("AITER_SITUV2_A8W4", "AITER_SITUV2_A4W4"):

614 + assert os.environ.get(name) == expected_env.get(name)

615

616

616 -def test_aiter_moe_situv2_clears_aiter_env_when_disabled(

617 +def test_aiter_moe_situv2_rejects_unknown_activation(

618 monkeypatch: pytest.MonkeyPatch,

619 ):

619 - from vllm._aiter_ops import rocm_aiter_ops

620 -

621 - _assert_aiter_supported()

622 -

620 with monkeypatch.context() as mp:

624 - mp.setenv("AITER_SITUV2_A4W4", "1")

625 - mp.setenv("VLLM_ROCM_USE_AITER", "1")

626 - mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")

627 - mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", raising=False)

628 - mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0")

621 + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "a8w8")

622 _reload_envs()

630 - rocm_aiter_ops.refresh_env_variables()

631 -

632 - import os

623 + import vllm.envs as envs

624

634 - assert "AITER_SITUV2_A4W4" not in os.environ

625 + with pytest.raises(ValueError, match="VLLM_ROCM_USE_AITER_MOE_SITUV2"):

626 + _ = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2

627

628

629 @pytest.mark.parametrize("moe_padding", [True, False])

vllm/_aiter_ops.py

+51 -23

@@ -1870,28 +1870,39 @@ def _mhc_delayed_pre_tail(

1870 _OPS_REGISTERED = False

1871

1872

1873 -def _sync_aiter_situv2_moe_env() -> None:

1874 - """Mirror the SiTUv2 MoE toggle into AITER's a4w4 dispatch env.

1873 +_AITER_SITUV2_ACT_ENV = {

1874 + "a8w4": "AITER_SITUV2_A8W4",

1875 + "a4w4": "AITER_SITUV2_A4W4",

1876 +}

1877 +

1878

1876 - AITER selects afp8 vs afp4 activation kernels via AITER_SITUV2_A8W4 /

1877 - AITER_SITUV2_A4W4 (see ROCm/aiter fused_moe.py, A8W4 checked first).

1878 - When VLLM_ROCM_USE_AITER_MOE_SITUV2 is enabled we route to a4w4

1879 - (afp4_wfp4_fp4 kernels) and clear any legacy AITER_SITUV2_A8W4 override.

1879 +def _resolve_situv2_activation() -> str:

1880 + """VLLM_ROCM_USE_AITER_MOE_SITUV2 -> a4w4 | a8w4 | a16w4.

1881

1881 - Requires AITER with ROCm/aiter#4463 (first tagged in v0.1.20): a4w4

1882 - dispatch plus kimik3_a4w4_{un,}tuned_fmoe.csv. Older AITER still runs

1883 - afp4 FlyDSL but falls back to heuristic configs, which is not the

1884 - tuned a4w4 path this recipe is meant to use.

1882 + auto (default) and the legacy 1 mean a4w4; legacy 0 means a16w4.

1883 """

1886 - import os

1884 + value = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2.lower()

1885 + if value in ("auto", "1"):

1886 + return "a4w4"

1887 + if value == "0":

1888 + return "a16w4"

1889 + return value

1890

1888 - import vllm.envs as envs

1891

1890 - if envs.VLLM_ROCM_USE_AITER_MOE_SITUV2:

1891 - os.environ["AITER_SITUV2_A4W4"] = "1"

1892 - os.environ.pop("AITER_SITUV2_A8W4", None)

1893 - else:

1894 - os.environ.pop("AITER_SITUV2_A4W4", None)

1892 +def _sync_aiter_situv2_moe_env() -> None:

1893 + """Set the AITER_SITUV2_* env matching VLLM_ROCM_USE_AITER_MOE_SITUV2.

1894 +

1895 + AITER reads AITER_SITUV2_A8W4 / AITER_SITUV2_A4W4 and defaults to a16w4,

1896 + so exactly one is set for a8w4/a4w4 and both are cleared for a16w4.

1897 + """

1898 + import os

1899 +

1900 + selected = _AITER_SITUV2_ACT_ENV.get(_resolve_situv2_activation())

1901 + for name in _AITER_SITUV2_ACT_ENV.values():

1902 + if name == selected:

1903 + os.environ[name] = "1"

1904 + else:

1905 + os.environ.pop(name, None)

1906

1907

1908 class rocm_aiter_ops:

@@ -1917,7 +1928,8 @@ class rocm_aiter_ops:

1928 VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM.

1929 VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.

1930 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.

1920 - VLLM_ROCM_USE_AITER_MOE_SITUV2: Controls SiTUv2 FlyDSL MoE (a4w4).

1931 + VLLM_ROCM_USE_AITER_MOE_SITUV2: SiTUv2 FlyDSL MoE activation

1932 + dtype (a16w4 | a8w4 | a4w4).

1933 VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM.

1934

1935 Note:

@@ -1987,7 +1999,7 @@ class rocm_aiter_ops:

1999 # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE

2000 _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE

2001 _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

1990 - _MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2

2002 + _MOE_SITUV2 = _resolve_situv2_activation()

2003 # TODO: Consolidate under _LINEAR_ENABLED

2004 _TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM

2005 # Lazily probed: whether aiter.topk_softmax supports the

@@ -2016,7 +2028,7 @@ class rocm_aiter_ops:

2028 cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM

2029 cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE

2030 cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

2019 - cls._MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2

2031 + cls._MOE_SITUV2 = _resolve_situv2_activation()

2032 _sync_aiter_situv2_moe_env()

2033 cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM

2034 cls._MOE_DISPATCH_POLICY = envs.VLLM_ROCM_AITER_MOE_DISPATCH_POLICY

@@ -2132,9 +2144,25 @@ class rocm_aiter_ops:

2144 @classmethod

2145 @if_aiter_supported

2146 def is_fused_moe_situv2_enabled(cls) -> bool:

2135 - # _MOE_SITUV2 is a variant of aiter fused moe, so aiter

2136 - # fused moe must be enabled as well.

2137 - return cls.is_fused_moe_enabled() and cls._MOE_SITUV2

2147 + """True when a low-precision (a8w4/a4w4) SiTUv2 activation is selected."""

2148 + return cls.is_fused_moe_enabled() and cls._MOE_SITUV2 != "a16w4"

2149 +

2150 + @classmethod

2151 + @if_aiter_supported

2152 + def get_fused_moe_situv2_activation(cls) -> str:

2153 + """SiTUv2 activation dtype name: a16w4, a8w4 or a4w4."""

2154 + return cls._MOE_SITUV2

2155 +

2156 + @classmethod

2157 + @if_aiter_supported

2158 + def is_fused_moe_situv2_gate_up_interleaved(cls) -> bool:

2159 + """Whether the SiTUv2 kernels expect gate/up-interleaved w13.

2160 +

2161 + The a8w4 FlyDSL stage-1 kernels (_gui_) are interleaved; a16w4 and

2162 + a4w4 are separated. Weight shuffle (oracle/mxfp4.py) and gate_mode

2163 + (experts/rocm_aiter_moe.py) must both use this.

2164 + """

2165 + return cls._MOE_SITUV2 == "a8w4"

2166

2167 @classmethod

2168 @if_aiter_supported

vllm/envs.py

+8 -13

@@ -141,7 +141,7 @@ if TYPE_CHECKING:

141 VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False

142 VLLM_ROCM_USE_AITER_MOE: bool = True

143 VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0

144 - VLLM_ROCM_USE_AITER_MOE_SITUV2: bool = False

144 + VLLM_ROCM_USE_AITER_MOE_SITUV2: Literal["auto", "a4w4", "a8w4", "a16w4"] = "auto"

145 VLLM_ROCM_USE_AITER_RMSNORM: bool = True

146 VLLM_ROCM_USE_AITER_MLA: bool = True

147 VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto"

@@ -1280,18 +1280,13 @@ environment_variables: dict[str, Callable[[], Any]] = {

1280 "VLLM_ROCM_USE_AITER_MOE": lambda: (

1281 os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1")

1282 ),

1283 - # Route K3 SiTU MXFP4 MoE through the FlyDSL SiTUv2 path (a4w4 fp4

1284 - # activations, separated gate/up layout) instead of default a16w4. vLLM

1285 - # sets AITER_SITUV2_A4W4 at init when this flag is on and clears any

1286 - # legacy AITER_SITUV2_A8W4 override (AITER checks A8W4 first).

1287 - # VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 is a deprecated alias for existing

1288 - # recipes; it does not select a8w4 kernels.

1289 - # Needs AITER >= v0.1.20 (ROCm/aiter#4463) for the a4w4 dispatch flag

1290 - # and tuned kimik3_a4w4_*_fmoe.csv rows; otherwise FlyDSL uses heuristics.

1291 - "VLLM_ROCM_USE_AITER_MOE_SITUV2": lambda: (

1292 - os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0").lower() in ("true", "1")

1293 - or os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower()

1294 - in ("true", "1")

1283 + # Activation dtype for the Kimi-K3 SiTU MXFP4 MoE (AITER FlyDSL SiTUv2):

1284 + # auto (= a4w4), a4w4, a8w4 or a16w4. Legacy 1/0 mean a4w4/a16w4.

1285 + "VLLM_ROCM_USE_AITER_MOE_SITUV2": env_with_choices(

1286 + "VLLM_ROCM_USE_AITER_MOE_SITUV2",

1287 + "auto",

1288 + ["auto", "a4w4", "a8w4", "a16w4", "0", "1"],

1289 + case_sensitive=False,

1290 ),

1291 # MoE sorting dispatch policy for AITER fused MoE kernels.

1292 # 0 = auto (default): single-pass for small batches, multi-pass

vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

+5 -4

@@ -385,10 +385,11 @@ def rocm_aiter_fused_experts(

385

386 gate_mode = ""

387 if activation == MoEActivation.SITU:

388 - # SiTUv2 flydsl (VLLM_ROCM_USE_AITER_MOE_SITUV2=1) uses a4w4

389 - # fp4 activations with separated gate/up weights (AITER #4463);

390 - # default a16w4 SiTU also stays separated.

391 - gate_mode = GateMode.SEPARATED.value

388 + gate_mode = (

389 + GateMode.INTERLEAVE.value

390 + if rocm_aiter_ops.is_fused_moe_situv2_gate_up_interleaved()

391 + else GateMode.SEPARATED.value

392 + )

393 elif quant_config.use_mxfp4_w4a16:

394 gate_mode = GateMode.INTERLEAVE.value

395 elif activation_interleave is not None:

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

+1 -2

@@ -1596,8 +1596,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1596

1597 fp4_dtype = torch.float4_e2m1fn_x2

1598 e8m0_dtype = torch.float8_e8m0fnu

1599 - # SiTUv2 flydsl uses separated gate/up layout (a4w4).

1600 - guinterleave = False

1599 + guinterleave = rocm_aiter_ops.is_fused_moe_situv2_gate_up_interleaved()

1600 w13 = rocm_aiter_ops.shuffle_weight_a16w4(

1601 w13_weight.data.view(fp4_dtype), 16, guinterleave

1602 )