@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8f50685c48

[ROCm][Kimi-K3] aiter moe environment variable cleanup (#50582)

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

Hongxia Yang committed Aug 3, 2026 at 14:54 UTC 8f50685c48d6bc88d199d9560ab67615e1b2c2a4

5 files changed +71 -56

vllm/_aiter_ops.py

+10

@@ -1530,6 +1530,7 @@ class rocm_aiter_ops:

1530 VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM.

1531 VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.

1532 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.

1533 + VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: Controls a8w4 SiTU fused MoE variant.

1534 VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM.

1535

1536 Note:

@@ -1598,6 +1599,7 @@ class rocm_aiter_ops:

1599 # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE

1600 _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE

1601 _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

1602 + _MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4

1603 # TODO: Consolidate under _LINEAR_ENABLED

1604 _TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM

1605 # Lazily probed: whether aiter.topk_softmax supports the

@@ -1627,6 +1629,7 @@ class rocm_aiter_ops:

1629 cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM

1630 cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE

1631 cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

1632 + cls._MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4

1633 cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM

1634

1635 @staticmethod

@@ -1719,6 +1722,13 @@ class rocm_aiter_ops:

1722 def is_fusion_moe_shared_experts_enabled(cls) -> bool:

1723 return cls.is_fused_moe_enabled() and cls._MOE_SHARED_EXPERTS_ENABLED

1724

1725 + @classmethod

1726 + @if_aiter_supported

1727 + def is_fused_moe_situv2_a8w4_enabled(cls) -> bool:

1728 + # _MOE_SITUV2_A8W4 is a variant of aiter fused moe, so aiter

1729 + # fused moe must be enabled as well.

1730 + return cls.is_fused_moe_enabled() and cls._MOE_SITUV2_A8W4

1731 +

1732 @classmethod

1733 @if_aiter_supported

1734 def topk_softmax_supports_fused_sigmoid(cls) -> bool:

vllm/envs.py

+5 -4

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

130 VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False

131 VLLM_ROCM_USE_AITER_MOE: bool = True

132 VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0

133 - AITER_SITUV2_A8W4: bool = False

133 + VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: bool = False

134 VLLM_ROCM_USE_AITER_RMSNORM: bool = True

135 VLLM_ROCM_USE_AITER_MLA: bool = True

136 VLLM_ROCM_USE_AITER_MHA: bool = True

@@ -1211,9 +1211,10 @@ environment_variables: dict[str, Callable[[], Any]] = {

1211 ),

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

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

1214 - # Shared with the AITER runtime, which reads the same env var directly.

1215 - "AITER_SITUV2_A8W4": lambda: (

1216 - os.getenv("AITER_SITUV2_A8W4", "0").lower() in ("true", "1")

1214 + # This is the only flag users need: vLLM picks the kernels by passing

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

1216 + "VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4": lambda: (

1217 + os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower() in ("true", "1")

1218 ),

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

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

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

+4 -4

@@ -5,7 +5,6 @@ from functools import lru_cache

6 import torch

8 -import vllm.envs as envs

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm._aiter_ops import rocm_aiter_ops

10 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

@@ -376,11 +375,12 @@ def rocm_aiter_fused_experts(

375

376 gate_mode = ""

377 if activation == MoEActivation.SITU:

379 - # a8w4 (AITER_SITUV2_A8W4=1) uses the gate/up-interleaved (_gui_)

380 - # fp8 flydsl kernels; default a16w4 SiTU stays separated.

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

379 + # interleaved (_gui_) fp8 flydsl kernels; default a16w4 SiTU stays

380 + # separated.

381 gate_mode = (

382 GateMode.INTERLEAVE.value

383 - if envs.AITER_SITUV2_A8W4

383 + if rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled()

384 else GateMode.SEPARATED.value

385 )

386 elif quant_config.use_mxfp4_w4a16:

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

+13 -2

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe import (

14 FusedMoEConfig,

15 RoutedExperts,

16 )

17 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

18 from vllm.model_executor.layers.fused_moe.all2all_utils import (

19 maybe_make_prepare_finalize,

20 )

@@ -634,10 +635,20 @@ def select_deepseek_v4_mxfp4_moe_backend(

635

636

637 def mxfp4_round_up_hidden_size_and_intermediate_size(

637 - backend: Mxfp4MoeBackend, hidden_size: int, intermediate_size: int

638 + backend: Mxfp4MoeBackend,

639 + hidden_size: int,

640 + intermediate_size: int,

641 + activation: MoEActivation | None = None,

642 ) -> tuple[int, int]:

643 """Round up hidden_size and intermediate_size based on backend requirements."""

640 - if backend == Mxfp4MoeBackend.EMULATION:

644 + if backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and activation == MoEActivation.SITU:

645 + # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size

646 + # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 (a

647 + # no-op for K3's shapes) rather than the generic ROCm 256 round-up,

648 + # which would inflate weights and OOM.

649 + intermediate_size = round_up(intermediate_size, 128)

650 + hidden_size = round_up(hidden_size, 128)

651 + elif backend == Mxfp4MoeBackend.EMULATION:

652 # Emulation has no kernel tile; it only needs OCP MX block alignment so the

653 # per-block scale buffers (`dim // OCP_MX_BLOCK_SIZE`) aren't floor-truncated

654 # by a non-block-aligned TP/DP shard (e.g. 2880 // 4 = 720).

vllm/model_executor/layers/quantization/mxfp4.py

+39 -46

@@ -1,9 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import os

5 +

6 import torch

6 -import vllm.envs as envs

8 from vllm.logger import init_logger

9 from vllm.model_executor.layers.attention import Attention

10 from vllm.model_executor.layers.fused_moe import (

@@ -515,6 +516,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

516 self.mxfp4_backend = Mxfp4MoeBackend.AITER_MXFP4_BF16

517 self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0]

518 logger.info_once("Using AITER_MXFP4_BF16 for Kimi-K3 SiTU MXFP4 MoE.")

519 + from vllm._aiter_ops import rocm_aiter_ops

520 +

521 + if rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled():

522 + # AITER keeps bf16 activations below this token count, which

523 + # would not match the fp8 a8w4 kernels the interleaved SiTU

524 + # path is tuned for. The a16w4 path never reads it.

525 + # TODO: Remove once AITER takes this as a kernel argument.

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

527 else:

528 self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(

529 moe

@@ -560,13 +569,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

569 act_dtype=act_dtype,

570 moe_parallel_config=moe_parallel_config,

571 )

563 - if self.is_k3_situ_aiter:

564 - # K3's AITER A16W4 kernel handles K3's native intermediate size

565 - # (moe_intermediate 3072; e.g. 384/partition at TP8); the generic

566 - # 256 round-up would inflate weights and OOM.

567 - return hidden_size, intermediate_size_per_partition

572 return mxfp4_round_up_hidden_size_and_intermediate_size(

569 - self.mxfp4_backend, hidden_size, intermediate_size_per_partition

573 + self.mxfp4_backend,

574 + hidden_size,

575 + intermediate_size_per_partition,

576 + activation=self.moe.activation,

577 )

578

579 def create_weights(

@@ -716,19 +723,24 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

723 )

724

725 # Convert weights to kernel format

719 - w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (

720 - convert_weight_to_mxfp4_moe_kernel_format(

721 - mxfp4_backend=self.mxfp4_backend,

722 - layer=layer,

723 - w13_weight=w13,

724 - w2_weight=w2,

725 - w13_weight_scale=w13_scale,

726 - w2_weight_scale=w2_scale,

727 - w13_bias=w13_bias,

728 - w2_bias=w2_bias,

729 - _cache_permute_indices=self._cache_permute_indices,

726 + if self.is_k3_situ_aiter:

727 + w13, w2, w13_scale, w2_scale = (

728 + self._convert_k3_situ_weight_to_kernel_format(layer)

729 + )

730 + else:

731 + w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (

732 + convert_weight_to_mxfp4_moe_kernel_format(

733 + mxfp4_backend=self.mxfp4_backend,

734 + layer=layer,

735 + w13_weight=w13,

736 + w2_weight=w2,

737 + w13_weight_scale=w13_scale,

738 + w2_weight_scale=w2_scale,

739 + w13_bias=w13_bias,

740 + w2_bias=w2_bias,

741 + _cache_permute_indices=self._cache_permute_indices,

742 + )

743 )

731 - )

744

745 # For TRITON backends, weights are wrapped tensors from triton_kernels

746 # that don't support .detach(). Manually assign parameters.

@@ -774,7 +786,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

786 layer=layer,

787 )

788

777 - def _setup_kernel_k3_situ(self, layer: RoutedExperts) -> None:

789 + def _convert_k3_situ_weight_to_kernel_format(

790 + self, layer: RoutedExperts

791 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

792 # K3's AITER A16W4 kernel wants the separated ([gate_all, up_all])

793 # stage-1 layout, unlike the interleaved gpt-oss/DeepSeek path in

794 # convert_weight_to_mxfp4_moe_kernel_format. Preshuffle once here.

@@ -786,10 +800,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

800 e8m0_dtype = torch.float8_e8m0fnu

801 num_experts = layer.w13_weight.shape[0]

802

789 - # a8w4 (AITER_SITUV2_A8W4=1) uses the gate/up-interleaved (_gui_) fp8

790 - # flydsl kernels, which need w13 weight+scale in interleave layout.

791 - # Default a16w4 keeps the separated layout.

792 - guinterleave = envs.AITER_SITUV2_A8W4

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

804 + # interleaved (_gui_) fp8 flydsl kernels, which need w13 weight+scale

805 + # in interleave layout. Default a16w4 keeps the separated layout.

806 + guinterleave = rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled()

807 w13 = rocm_aiter_ops.shuffle_weight_a16w4(

808 layer.w13_weight.data.view(fp4_dtype), 16, guinterleave

809 )

@@ -802,30 +816,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

816 w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), num_experts, guinterleave

817 )

818 w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1]))

805 -

806 - replace_parameter(layer, "w13_weight", w13)

807 - replace_parameter(layer, "w2_weight", w2)

808 - replace_parameter(layer, "w13_weight_scale", w13_scale)

809 - replace_parameter(layer, "w2_weight_scale", w2_scale)

810 - layer.w13_weight.is_shuffled = True

811 - layer.w2_weight.is_shuffled = True

812 -

813 - self.moe_quant_config = self.get_fused_moe_quant_config(layer)

814 - if self.moe_quant_config is not None and self.experts_cls is not None:

815 - self.moe_kernel = make_mxfp4_moe_kernel(

816 - moe_quant_config=self.moe_quant_config,

817 - moe_config=self.moe,

818 - mxfp4_backend=self.mxfp4_backend,

819 - experts_cls=self.experts_cls,

820 - routing_tables=layer._expert_routing_tables(),

821 - layer=layer,

822 - )

819 + return w13, w2, w13_scale, w2_scale

820

821 def process_weights_after_loading(self, layer):

825 - if self.is_k3_situ_aiter:

826 - self._setup_kernel_k3_situ(layer)

827 - return

828 -

822 w13 = layer.w13_weight

823 w2 = layer.w2_weight

824 w13_scale = layer.w13_weight_scale