@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a4d2d9d95e

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

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>

ppalanga committed Sep 15, 2026 at 18:38 UTC a4d2d9d95edd2db39cc7dec0a3433c764ff2c578

5 files changed +127 -30

tests/kernels/moe/test_rocm_aiter_moe.py

+68

@@ -565,6 +565,74 @@ def test_aiter_moe_shared_experts_enablement_follows_env(

565 assert rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() is expected

566

567

568 +def test_aiter_moe_situv2_syncs_aiter_a4w4_env(

569 + monkeypatch: pytest.MonkeyPatch,

570 +):

571 + """SiTUv2 MoE should route to AITER a4w4 and clear legacy a8w4 overrides."""

572 + from vllm._aiter_ops import rocm_aiter_ops

573 +

574 + _assert_aiter_supported()

575 +

576 + with monkeypatch.context() as mp:

577 + mp.delenv("AITER_SITUV2_A8W4", raising=False)

578 + mp.delenv("AITER_SITUV2_A4W4", raising=False)

579 + mp.setenv("VLLM_ROCM_USE_AITER", "1")

580 + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")

581 + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "1")

582 + _reload_envs()

583 + rocm_aiter_ops.refresh_env_variables()

584 +

585 + import os

586 +

587 + assert os.environ.get("AITER_SITUV2_A4W4") == "1"

588 + assert "AITER_SITUV2_A8W4" not in os.environ

589 +

590 +

591 +def test_aiter_moe_situv2_legacy_a8w4_alias_enables_a4w4(

592 + monkeypatch: pytest.MonkeyPatch,

593 +):

594 + """Deprecated VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 still enables SiTUv2 a4w4."""

595 + from vllm._aiter_ops import rocm_aiter_ops

596 +

597 + _assert_aiter_supported()

598 +

599 + with monkeypatch.context() as mp:

600 + mp.delenv("AITER_SITUV2_A8W4", raising=False)

601 + mp.delenv("AITER_SITUV2_A4W4", raising=False)

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

603 + mp.setenv("VLLM_ROCM_USE_AITER", "1")

604 + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")

605 + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "1")

606 + _reload_envs()

607 + rocm_aiter_ops.refresh_env_variables()

608 +

609 + import os

610 +

611 + assert os.environ.get("AITER_SITUV2_A4W4") == "1"

612 + assert "AITER_SITUV2_A8W4" not in os.environ

613 +

614 +

615 +def test_aiter_moe_situv2_clears_aiter_env_when_disabled(

616 + monkeypatch: pytest.MonkeyPatch,

617 +):

618 + from vllm._aiter_ops import rocm_aiter_ops

619 +

620 + _assert_aiter_supported()

621 +

622 + with monkeypatch.context() as mp:

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

624 + mp.setenv("VLLM_ROCM_USE_AITER", "1")

625 + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")

626 + mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", raising=False)

627 + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0")

628 + _reload_envs()

629 + rocm_aiter_ops.refresh_env_variables()

630 +

631 + import os

632 +

633 + assert "AITER_SITUV2_A4W4" not in os.environ

634 +

635 +

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

637 def test_aiter_moe_padding_env_var(

638 moe_padding: bool,

vllm/_aiter_ops.py

+32 -6

@@ -1752,6 +1752,30 @@ def _mhc_delayed_pre_tail(

1752 _OPS_REGISTERED = False

1753

1754

1755 +def _sync_aiter_situv2_moe_env() -> None:

1756 + """Mirror the SiTUv2 MoE toggle into AITER's a4w4 dispatch env.

1757 +

1758 + AITER selects afp8 vs afp4 activation kernels via AITER_SITUV2_A8W4 /

1759 + AITER_SITUV2_A4W4 (see ROCm/aiter fused_moe.py, A8W4 checked first).

1760 + When VLLM_ROCM_USE_AITER_MOE_SITUV2 is enabled we route to a4w4

1761 + (afp4_wfp4_fp4 kernels) and clear any legacy AITER_SITUV2_A8W4 override.

1762 +

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

1764 + dispatch plus kimik3_a4w4_{un,}tuned_fmoe.csv. Older AITER still runs

1765 + afp4 FlyDSL but falls back to heuristic configs, which is not the

1766 + tuned a4w4 path this recipe is meant to use.

1767 + """

1768 + import os

1769 +

1770 + import vllm.envs as envs

1771 +

1772 + if envs.VLLM_ROCM_USE_AITER_MOE_SITUV2:

1773 + os.environ["AITER_SITUV2_A4W4"] = "1"

1774 + os.environ.pop("AITER_SITUV2_A8W4", None)

1775 + else:

1776 + os.environ.pop("AITER_SITUV2_A4W4", None)

1777 +

1778 +

1779 class rocm_aiter_ops:

1780 """ROCm AITER operations wrapper for AMD GPU acceleration in vLLM.

1781

@@ -1774,7 +1798,7 @@ class rocm_aiter_ops:

1798 VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply.

1799 VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.

1800 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.

1777 - VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: Controls a8w4 SiTU fused MoE variant.

1801 + VLLM_ROCM_USE_AITER_MOE_SITUV2: Controls SiTUv2 FlyDSL MoE (a4w4).

1802 VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM.

1803

1804 Note:

@@ -1841,7 +1865,7 @@ class rocm_aiter_ops:

1865 # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE

1866 _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE

1867 _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

1844 - _MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4

1868 + _MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2

1869 # TODO: Consolidate under _LINEAR_ENABLED

1870 _TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM

1871 # Lazily probed: whether aiter.topk_softmax supports the

@@ -1870,7 +1894,8 @@ class rocm_aiter_ops:

1894 cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM

1895 cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE

1896 cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

1873 - cls._MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4

1897 + cls._MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2

1898 + _sync_aiter_situv2_moe_env()

1899 cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM

1900

1901 @staticmethod

@@ -1983,10 +2008,10 @@ class rocm_aiter_ops:

2008

2009 @classmethod

2010 @if_aiter_supported

1986 - def is_fused_moe_situv2_a8w4_enabled(cls) -> bool:

1987 - # _MOE_SITUV2_A8W4 is a variant of aiter fused moe, so aiter

2011 + def is_fused_moe_situv2_enabled(cls) -> bool:

2012 + # _MOE_SITUV2 is a variant of aiter fused moe, so aiter

2013 # fused moe must be enabled as well.

1989 - return cls.is_fused_moe_enabled() and cls._MOE_SITUV2_A8W4

2014 + return cls.is_fused_moe_enabled() and cls._MOE_SITUV2

2015

2016 @classmethod

2017 @if_aiter_supported

@@ -4029,3 +4054,4 @@ class rocm_aiter_ops:

4054

4055

4056 rocm_aiter_ops.register_ops_once()

4057 +_sync_aiter_situv2_moe_env()

vllm/envs.py

+13 -7

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

137 VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False

138 VLLM_ROCM_USE_AITER_MOE: bool = True

139 VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0

140 - VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: bool = False

140 + VLLM_ROCM_USE_AITER_MOE_SITUV2: bool = False

141 VLLM_ROCM_USE_AITER_RMSNORM: bool = True

142 VLLM_ROCM_USE_AITER_MLA: bool = True

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

@@ -1251,12 +1251,18 @@ environment_variables: dict[str, Callable[[], Any]] = {

1251 "VLLM_ROCM_USE_AITER_MOE": lambda: (

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

1253 ),

1254 - # Route K3 SiTU MXFP4 MoE through the a8w4 (fp8 activation) gate/up-

1255 - # interleaved flydsl kernels instead of the default a16w4 separated path.

1256 - # This is the only flag users need: vLLM picks the kernels by passing

1257 - # gate_mode to AITER and sets the AITER-side workaround env at init.

1258 - "VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4": lambda: (

1259 - os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower() in ("true", "1")

1254 + # Route K3 SiTU MXFP4 MoE through the FlyDSL SiTUv2 path (a4w4 fp4

1255 + # activations, separated gate/up layout) instead of default a16w4. vLLM

1256 + # sets AITER_SITUV2_A4W4 at init when this flag is on and clears any

1257 + # legacy AITER_SITUV2_A8W4 override (AITER checks A8W4 first).

1258 + # VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 is a deprecated alias for existing

1259 + # recipes; it does not select a8w4 kernels.

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

1261 + # and tuned kimik3_a4w4_*_fmoe.csv rows; otherwise FlyDSL uses heuristics.

1262 + "VLLM_ROCM_USE_AITER_MOE_SITUV2": lambda: (

1263 + os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0").lower() in ("true", "1")

1264 + or os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower()

1265 + in ("true", "1")

1266 ),

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

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

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

+4 -8

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

385

386 gate_mode = ""

387 if activation == MoEActivation.SITU:

388 - # a8w4 (VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1) uses the gate/up-

389 - # interleaved (_gui_) fp8 flydsl kernels; default a16w4 SiTU stays

390 - # separated.

391 - gate_mode = (

392 - GateMode.INTERLEAVE.value

393 - if rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled()

394 - else GateMode.SEPARATED.value

395 - )

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

392 elif quant_config.use_mxfp4_w4a16:

393 gate_mode = GateMode.INTERLEAVE.value

394 elif activation_interleave is not None:

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

+10 -9

@@ -1572,12 +1572,6 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1572 if w2_bias is not None:

1573 w2_bias = w2_bias.data.to(torch.float32)

1574

1575 - import os

1576 -

1577 - # TODO: Remove this once AITER is fixed

1578 - # Necessary for AITER side from crashing

1579 - os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"

1580 -

1575 if activation == MoEActivation.SITU:

1576 from aiter.utility.fp4_utils import e8m0_shuffle

1577

@@ -1585,9 +1579,8 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1579

1580 fp4_dtype = torch.float4_e2m1fn_x2

1581 e8m0_dtype = torch.float8_e8m0fnu

1588 - # a8w4 uses gate/up-interleaved flydsl kernels;

1589 - # default a16w4 keeps the separated layout.

1590 - guinterleave = rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled()

1582 + # SiTUv2 flydsl uses separated gate/up layout (a4w4).

1583 + guinterleave = False

1584 w13 = rocm_aiter_ops.shuffle_weight_a16w4(

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

1586 )

@@ -1606,6 +1599,14 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1599 w2.is_shuffled = True

1600 return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)

1601

1602 + import os

1603 +

1604 + # Interleaved a16w4 only (DeepSeekV4 etc.). AITER uses this bound to

1605 + # pick bf16 vs fp8 activations when gate_mode is INTERLEAVE. SiTUv2

1606 + # a4w4 is separated and selects q_dtype_a independently, so the bound

1607 + # is unused on that path.

1608 + os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"

1609 +

1610 from aiter.ops.shuffle import shuffle_scale as _shuf_s

1611 from aiter.ops.shuffle import shuffle_weight as _shuf_w

1612