Enable Kimi-K3 for Gfx1250

vllm/_aiter_ops.py

@@ -2092,6 +2092,11 @@ def _resolve_situv2_activation() -> str:
     auto (default) and the legacy 1 mean a4w4; legacy 0 means a16w4.
     """
     value = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2.lower()
+    if value == "auto" and current_platform.is_rocm():
+        from vllm.platforms.rocm import on_gfx1250
+
+        if on_gfx1250():
+            return "a8w4"
     if value in ("auto", "1"):
         return "a4w4"
     if value == "0":
@@ -2738,11 +2743,14 @@ def is_kda_gluon_decode_available(cls) -> bool:
         RMSNorm in one launch, plain and speculative decode) ships CDNA4+
         (gfx950, gfx1250) builds only, and is absent in older aiter releases.
         """
-        from vllm.platforms.rocm import get_cdna_version
+        from vllm.platforms.rocm import get_cdna_version, on_gfx1250
 
+        # Not yet validated for Kimi-K3 on gfx1250, which keeps the Triton
+        # decode there.
         return (
             cls._AITER_ENABLED
             and get_cdna_version() >= 4
+            and not on_gfx1250()
             and cls._kda_gluon_decode_importable()
         )
 

vllm/config/kernel.py

@@ -140,6 +140,7 @@ def with_default(
     "aiter_mega_moe",
     "aiter_triton_mxfp4_bf16",
     "aiter_triton_mxfp4_mxfp4",
+    "aiter_flydsl_mxfp4_bf16",
     "flydsl",
     "hpc",
     "emulation",
@@ -336,6 +337,8 @@ class KernelConfig:
       (moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250)
     - "aiter_triton_mxfp4_mxfp4": Use the AITER Triton/Gluon MXFP4 W4A4
       (moe_gemm_a4w4) MoE kernel (ROCm gfx1250)
+    - "aiter_flydsl_mxfp4_bf16": Use the AITER FlyDSL MXFP4 SiTU MoE kernel
+      (a8w4/a4w4, ROCm gfx1250; Kimi-K3). Opt-in only, never auto-selected.
     - "flydsl": Use AMD FlyDSL kernels (ROCm only)
     - "rdna3": Use the fused RDNA3 W4A16 HIP kernel (ROCm gfx1100 only)
     - "hpc": Use HPC kernels (FP8 and Hopper only)

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

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MXFP4 SiTU MoE via AITER's FlyDSL grouped GEMM (gfx1250). Routes through
+``aiter.fused_moe`` with gate/up-interleaved w13; weights are preshuffled in
+``convert_weight_to_mxfp4_moe_kernel_format``.
+"""
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
+    AiterExperts,
+    QuantMethod,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kMxfp4Static,
+)
+
+
+class AiterFlydslMxfp4Experts(AiterExperts):
+    """MXFP4 weights with BF16 activations through AITER's FlyDSL MoE on
+    gfx1250. AITER quantizes activations internally to FP8 (a8w4) when
+    AITER_FORCE_A8W4=1 is set, otherwise FP4 (a4w4)."""
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        if (weight_key, activation_key) != (kMxfp4Static, None):
+            return False
+        from vllm.platforms.rocm import on_gfx1250
+
+        return on_gfx1250()
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation == MoEActivation.SITU
+
+    @staticmethod
+    def is_supported_config(
+        cls, moe_config, weight_key, activation_key, activation_format
+    ):
+        is_supported, reason = super().is_supported_config(
+            cls, moe_config, weight_key, activation_key, activation_format
+        )
+        situv2 = rocm_aiter_ops.get_fused_moe_situv2_activation()
+        if is_supported and situv2 not in ("a8w4", "a4w4"):
+            return False, (
+                "gfx1250 FlyDSL SiTU MoE only has a8w4 and a4w4 kernels; got "
+                f"VLLM_ROCM_USE_AITER_MOE_SITUV2={situv2}"
+            )
+        return is_supported, reason
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ):
+        from aiter.ops.flydsl.moe_common import GateMode
+
+        moe_config = self.moe_config
+        quant_config = self.quant_config
+        assert moe_config.hidden_dim_unpadded is not None
+        assert moe_config.intermediate_size_per_partition_unpadded is not None
+        hidden_size = hidden_states.shape[1]
+        if hidden_states.dtype == torch.float4_e2m1fn_x2:
+            # Pre-quantized MXFP4 from MoRI dispatch, packed two per byte.
+            hidden_size *= 2
+        # SiTU's FlyDSL kernel pads per gate/up half, so pads stay unrounded.
+        intermediate_pad = (
+            moe_config.intermediate_size_per_partition
+            - moe_config.intermediate_size_per_partition_unpadded
+            if moe_config.intermediate_pad is None
+            else moe_config.intermediate_pad
+        )
+
+        result = rocm_aiter_ops.fused_moe(
+            hidden_states,
+            w1,
+            w2,
+            topk_weights.to(torch.float32),
+            topk_ids.to(torch.int32),
+            # RoutedExperts hands AITER experts the 0/1 expert_mask here.
+            expert_mask=expert_map,
+            quant_method=QuantMethod.BLOCK_1X32.value,
+            activation_method=rocm_aiter_ops.get_aiter_activation_type("situ"),
+            w1_scale=quant_config.w1_scale,
+            w2_scale=quant_config.w2_scale,
+            a1_scale=quant_config.a1_scale if a1q_scale is None else a1q_scale,
+            a2_scale=quant_config.a2_scale,
+            doweight_stage1=apply_router_weight_on_input,
+            num_local_tokens=(
+                expert_tokens_meta.expert_num_tokens
+                if expert_tokens_meta is not None
+                else None
+            ),
+            output_dtype=output.dtype,
+            hidden_pad=hidden_size - moe_config.hidden_dim_unpadded,
+            intermediate_pad=intermediate_pad,
+            # gfx1250 FlyDSL only accepts gate/up-interleaved w13, matching the
+            # shuffle in oracle/mxfp4.py.
+            gate_mode=GateMode.INTERLEAVE.value,
+            bias1=quant_config.w1_bias,
+            bias2=quant_config.w2_bias,
+            moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(),
+            swiglu_limit=(
+                0.0
+                if moe_config.swiglu_limit is None
+                else float(moe_config.swiglu_limit)
+            ),
+            beta=moe_config.activation_situ_beta,
+            linear_beta=moe_config.activation_situ_linear_beta,
+        )
+        # avoid redundant copy when output is a view of the result
+        if (
+            output.shape == result.shape
+            and output.dtype == result.dtype
+            and output.device == result.device
+            and output.is_contiguous()
+            and result.is_contiguous()
+            and output._base is None
+        ):
+            output.set_(result)
+        else:
+            output.copy_(result)

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

@@ -153,6 +153,8 @@ class Mxfp4MoeBackend(Enum):
     AITER = "AITER_MXFP4_BF16"
     # W4A16: aiter Triton moe_gemm_a16w4 kernel (gfx942/gfx950/gfx1250)
     AITER_TRITON_MXFP4_BF16 = "AITER_TRITON_MXFP4_BF16"
+    # W4A16 interface: aiter FlyDSL SiTU MoE (gfx1250), fp8/fp4 acts internally
+    AITER_FLYDSL_MXFP4_BF16 = "AITER_FLYDSL_MXFP4_BF16"
     AITER_MXFP4_FP8 = "AITER_MXFP4_FP8"  # W4A8: triton kernel
     AITER_MXFP4_MXFP4 = "AITER_MXFP4_MXFP4"  # W4A4: CK kernel (gfx950)
     # W4A4: aiter Triton/Gluon moe_gemm_a4w4 kernel (gfx1250)
@@ -306,6 +308,13 @@ def backend_to_kernel_cls(
 
         return [AiterMxfp4Experts]
 
+    elif backend == Mxfp4MoeBackend.AITER_FLYDSL_MXFP4_BF16:
+        from vllm.model_executor.layers.fused_moe.experts.aiter_flydsl_mxfp4_moe import (  # noqa: E501
+            AiterFlydslMxfp4Experts,
+        )
+
+        return [AiterFlydslMxfp4Experts]
+
     elif backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16:
         from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import (
             AiterW4A16ExpertsMonolithic,
@@ -400,6 +409,7 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]:
         "aiter_mxfp4_fp8": [Mxfp4MoeBackend.AITER_MXFP4_FP8],
         "aiter_triton_mxfp4_mxfp4": [Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4],
         "aiter_mxfp4_mxfp4": [Mxfp4MoeBackend.AITER_MXFP4_MXFP4],
+        "aiter_flydsl_mxfp4_bf16": [Mxfp4MoeBackend.AITER_FLYDSL_MXFP4_BF16],
         "xpu": [Mxfp4MoeBackend.XPU],
         "cpu": [Mxfp4MoeBackend.CPU],
         "emulation": [Mxfp4MoeBackend.EMULATION],
@@ -902,7 +912,10 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
         # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128
         # rather than the generic ROCm 256 round-up, which would inflate
         # weights and OOM.
-        aiter_uses_128 = backend == Mxfp4MoeBackend.AITER_MXFP4_BF16
+        aiter_uses_128 = backend in (
+            Mxfp4MoeBackend.AITER_MXFP4_BF16,
+            Mxfp4MoeBackend.AITER_FLYDSL_MXFP4_BF16,
+        )
 
         # matmul_ogs uses block_k=128 for MXFP4 on pre-CDNA4 GPUs.
         # CDNA4's F16xMXFP4 configuration uses block_k=256.
@@ -1644,12 +1657,6 @@ def convert_weight_to_mxfp4_moe_kernel_format(
     Supports CUTLASS, DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin
     backends.
     """
-    is_gfx1250 = False
-    if current_platform.is_rocm():
-        from vllm.platforms.rocm import on_gfx1250
-
-        is_gfx1250 = on_gfx1250()
-
     if mxfp4_backend in B12X_BACKENDS:
         return (
             w13_weight.data,
@@ -1844,7 +1851,48 @@ def convert_weight_to_mxfp4_moe_kernel_format(
             w2_bias,
         )
 
-    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not is_gfx1250:
+    elif mxfp4_backend == Mxfp4MoeBackend.AITER_FLYDSL_MXFP4_BF16:
+        from aiter.ops.shuffle import (
+            interleave_gate_up_rows,
+            moe_shuffle_scale,
+            moe_shuffle_weight,
+        )
+
+        fp4_dtype = torch.float4_e2m1fn_x2
+        e8m0_dtype = torch.float8_e8m0fnu
+
+        w13_raw = w13_weight.data.view(fp4_dtype)
+        w2_raw = w2_weight.data.view(fp4_dtype)
+        w13_scale_raw = w13_weight_scale.data.view(e8m0_dtype)
+        w2_scale_raw = w2_weight_scale.data.view(e8m0_dtype)
+        w13_scale_raw = w13_scale_raw.view(-1, w13_scale_raw.shape[-1])
+        w2_scale_raw = w2_scale_raw.view(-1, w2_scale_raw.shape[-1])
+
+        # gfx1250 FlyDSL only accepts gate/up-interleaved w13; must match
+        # AiterFlydslMxfp4Experts' gate_mode.
+        w13 = moe_shuffle_weight(
+            w13_raw, num_experts, is_guinterleave=True, gate_up=True
+        )
+        w2 = moe_shuffle_weight(
+            w2_raw, num_experts, is_guinterleave=False, gate_up=False
+        )
+        w13_scale = moe_shuffle_scale(
+            w13_scale_raw, num_experts, is_guinterleave=True, gate_up=True
+        )
+        w2_scale = moe_shuffle_scale(
+            w2_scale_raw, num_experts, is_guinterleave=False, gate_up=False
+        )
+        w13.is_shuffled = True
+        w2.is_shuffled = True
+
+        if w13_bias is not None:
+            w13_bias = interleave_gate_up_rows(w13_bias.data.to(torch.float32))
+        if w2_bias is not None:
+            w2_bias = w2_bias.data.to(torch.float32)
+
+        return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)
+
+    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
         # Initially introduced for DeepSeekV4
 
         if w13_bias is not None:
@@ -1938,9 +1986,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
             w2_bias,
         )
 
-    elif mxfp4_backend in TRITON_BACKENDS or (
-        mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250
-    ):
+    elif mxfp4_backend in TRITON_BACKENDS:
         if mxfp4_backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16:
             # AITER moe_gemm_a16w4 needs gate/up interleaved
             def interleave_gate_up(w: torch.Tensor) -> torch.Tensor:
@@ -2256,6 +2302,7 @@ def make_mxfp4_moe_quant_config(
         Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
         Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
         Mxfp4MoeBackend.AITER_MXFP4_BF16,
+        Mxfp4MoeBackend.AITER_FLYDSL_MXFP4_BF16,
         Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16,
         Mxfp4MoeBackend.CPU,
     ):

vllm/model_executor/layers/quantization/mxfp4.py

@@ -47,7 +47,6 @@
     replace_parameter,
     set_weight_attrs,
 )
-from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
 
@@ -837,16 +836,7 @@ def _setup_kernel(
 
         # For TRITON backends, weights are wrapped tensors from triton_kernels
         # that don't support .detach(). Manually assign parameters.
-        is_gfx1250 = False
-        if current_platform.is_rocm():
-            from vllm.platforms.rocm import on_gfx1250
-
-            is_gfx1250 = on_gfx1250()
-
-        uses_triton_weight_format = self.mxfp4_backend in TRITON_BACKENDS or (
-            self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250
-        )
-        if not uses_triton_weight_format:
+        if self.mxfp4_backend not in TRITON_BACKENDS:
             replace_parameter(layer, "w13_weight", w13)
             replace_parameter(layer, "w2_weight", w2)
             replace_parameter(layer, "w13_weight_scale", w13_scale)
@@ -1046,15 +1036,7 @@ def get_fused_moe_quant_config(
         w2_bias = getattr(layer, "w2_bias", None)
         swiglu_limit = getattr(layer, "swiglu_limit", None)
 
-        is_gfx1250 = False
-        if current_platform.is_rocm():
-            from vllm.platforms.rocm import on_gfx1250
-
-            is_gfx1250 = on_gfx1250()
-
-        if self.mxfp4_backend in TRITON_BACKENDS or (
-            self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250
-        ):
+        if self.mxfp4_backend in TRITON_BACKENDS:
             # TRITON backends free w13/w2_weight_scale after swizzling; the
             # swizzled scales live inside the precision configs instead.
             assert self.w13_precision_config is not None

vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py

@@ -36,7 +36,14 @@
 # 4-stage pipeline over it emits a third, consumer-less async copy into LDS that
 # the `w` loop then reuses; at num_warps=4 that races on gfx950 and `u` comes
 # back with non-deterministic O(1e38) garbage once a batch reaches 4096 tokens.
+# gfx1250 separately returns wrong results at any depth above 1; root cause
+# not established.
 _RECOMPUTE_W_U_NUM_STAGES = [2, 3]
+if is_amd:
+    from vllm.platforms.rocm import on_gfx1250
+
+    if on_gfx1250():
+        _RECOMPUTE_W_U_NUM_STAGES = [1]
 
 
 @triton.heuristics(

vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py

@@ -18,7 +18,14 @@
 
 NUM_WARPS = [2, 4, 8, 16]
 # Triton's AMD backend fails to lower this kernel with num_stages=4.
-_CHUNK_DELTA_H_NUM_STAGES = [2, 3] if torch.version.hip else [2, 3, 4]
+# gfx1250 needs num_stages=1 on top of that: any deeper pipeline corrupts the
+# recurrent state and accuracy drops to zero. Root cause not established.
+if torch.version.hip:
+    from vllm.platforms.rocm import on_gfx1250
+
+    _CHUNK_DELTA_H_NUM_STAGES = [1] if on_gfx1250() else [2, 3]
+else:
+    _CHUNK_DELTA_H_NUM_STAGES = [2, 3, 4]
 
 
 @triton.heuristics(