From bbee55e717a0b77090997e97931446e7af604cf4 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Fri, 16 Oct 2026 15:24:21 -0400 Subject: [PATCH] kimi enablement flydsl Kimi-K3 enablement on gfx1250: - Opt-in AITER FlyDSL MXFP4 SiTU MoE backend (--moe-backend aiter_flydsl_mxfp4_bf16); SiTUv2 auto resolves to a8w4 on gfx1250. - Pin the KDA chunk kernels to num_stages=1 on gfx1250. - Route gfx1250 MLA prefill to AITER's Triton flash attention. - Keep gfx1250 on the Triton KDA decode; AITER's Gluon KDA decode is not yet validated there. On gfx1250 serve with AITER_FORCE_A8W4=1 and --kda-prefill-backend triton: AITER picks a8w4 there from that flag alone (a4w4 otherwise), and the auto KDA prefill picks AITER's Gluon chunk_kda, which fails to compile there. Signed-off-by: Jaden Mathias Co-Authored-By: Claude Opus 5.5 (1M context) --- vllm/_aiter_ops.py | 10 +- vllm/config/kernel.py | 3 + .../experts/aiter_flydsl_mxfp4_moe.py | 142 ++++++++++++++++++ .../layers/fused_moe/oracle/mxfp4.py | 69 +++++++-- .../layers/quantization/mxfp4.py | 22 +-- .../kimi_k3/amd/ops/third_party/kda/chunk.py | 7 + .../ops/chunk_delta_h.py | 9 +- .../backends/mla/prefill/flash_attn.py | 10 ++ 8 files changed, 239 insertions(+), 33 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/experts/aiter_flydsl_mxfp4_moe.py diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index e26d9bcc9ffe..2faf142786ab 100644 --- a/vllm/_aiter_ops.py +++ b/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() ) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index b50bc32e9a31..7be2af131607 100644 --- a/vllm/config/kernel.py +++ b/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) diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_flydsl_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_flydsl_mxfp4_moe.py new file mode 100644 index 000000000000..94a3425a77bc --- /dev/null +++ b/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) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index bbb7d6e86d36..4f458fa6f925 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/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, ): diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 0a63cf002f3f..473cc6e02753 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/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 diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py index 2e975e9f7718..95eede6ec943 100644 --- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py +++ b/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( diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py index cc27437fac58..74ffff98524b 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py +++ b/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( diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index b6cf96e5b738..bf9ba894c7f2 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -395,6 +395,11 @@ def __init__( # Track whether we're using vllm's FA or upstream (for ROCm) self._is_vllm_fa = current_platform.is_cuda() or current_platform.is_xpu() + self._is_aiter_triton_fa = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + + self._is_aiter_triton_fa = on_gfx1250() def supports_quant_output(self, quant_key: "QuantKey") -> bool: device_capability = current_platform.get_device_capability() @@ -429,6 +434,9 @@ def _flash_attn_varlen_diff_headdims( kwargs["output_scale"] = output_scale if envs.VLLM_BATCH_INVARIANT: kwargs["num_splits"] = 1 + elif self._is_aiter_triton_fa: + kwargs["return_lse"] = return_softmax_lse + assert out is None and output_scale is None else: # ROCm leverages the upstream flash_attn, which takes a parameter # called "return_attn_probs" instead of return_softmax_lse @@ -448,6 +456,8 @@ def _flash_attn_varlen_diff_headdims( lse = None if isinstance(attn_out, tuple): attn_out, lse = attn_out[0], attn_out[1] + if self._is_aiter_triton_fa and lse is not None: + lse = lse.transpose(0, 1).contiguous() # Remain consistent with old `flash_attn_varlen_func` where there # is only one output tensor if `return_softmax_lse` is False.