From d410fc12f30b4224fca71a2faca75dcd4a5c3557 Mon Sep 17 00:00:00 2001 From: Bolin Sun Date: Thu, 3 Sep 2026 12:49:31 -0400 Subject: [PATCH] [Kernel] Enable Kimi-K3 SiTU on the CuteDSL MoE backend and the SM107 low-latency GEMM plan (#54606) Signed-off-by: Bolin Sun Co-authored-by: Claude Fable 5 --- .../experts/flashinfer_cutedsl_moe.py | 20 ++++++++++++++++++- .../models/kimi_k3/nvidia/low_latency_gemm.py | 14 ++++++++++--- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index e54a62124a..ec0e15bde5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -62,6 +62,8 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): self.gemm1_alpha = quant_config.gemm1_alpha self.gemm1_beta = quant_config.gemm1_beta self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit + self.situ_beta = moe_config.activation_situ_beta + self.situ_linear_beta = moe_config.activation_situ_linear_beta def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) @@ -101,6 +103,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): MoEActivation.SWIGLUOAI, MoEActivation.SWIGLUOAI_UNINTERLEAVE, MoEActivation.RELU2_NO_MUL, + MoEActivation.SITU, ) @staticmethod @@ -171,6 +174,19 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): "swiglu_beta": self.gemm1_beta, "swiglu_limit": self.gemm1_clamp_limit, } + elif activation == MoEActivation.SITU: + # The cute_dsl kernel keys SiTU on situ_beta and requires + # activation_type to stay a base type (ActivationType.Situ is + # rejected by normalize_cute_dsl_moe_activation_type), so the + # Swiglu base type is passed below and SiTU rides the betas. + if self.situ_beta is None: + raise ValueError( + "SITU activation requires moe_config.activation_situ_beta" + ) + swiglu_params = { + "situ_beta": self.situ_beta, + "situ_linear_beta": self.situ_linear_beta, + } swiglu_kwargs = {k: v for k, v in swiglu_params.items() if v is not None} flashinfer_cute_dsl_fused_moe_nvfp4( @@ -190,6 +206,8 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): num_local_experts=self.local_num_experts, local_expert_offset=self.local_expert_offset, moe_output=output, - activation_type=activation_to_flashinfer_int(activation), + activation_type=activation_to_flashinfer_int( + MoEActivation.SILU if activation == MoEActivation.SITU else activation + ), **swiglu_kwargs, ) diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py index b3a4853bd1..17c07b43a6 100644 --- a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py +++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM90/SM100/SM103. +"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM90/SM100/SM103/SM107. Dispatch is purely by local ``(N, K)`` shape and token count ``M`` — the module name plays no role. Each measured shape maps to a :class:`ProjectionSpec` @@ -12,7 +12,10 @@ The supported capabilities carry separate measured tables: :data:`KIMI_K3_PROJECTIONS` was tuned on B300 (SM103), :data:`KIMI_K3_PROJECTIONS_SM100` on B200 (SM100), and :data:`KIMI_K3_PROJECTIONS_SM90` on H200 (SM90). The per-(shape, M) winners -genuinely differ between the parts, so the tables must not be merged. +genuinely differ between the parts, so the tables must not be merged. SM107 +(Rubin) reuses the SM103 table: the plan was validated end-to-end on SM107 +hardware, but the per-M crossovers have not been re-measured there and may +deserve their own table once retuned. """ from __future__ import annotations @@ -646,9 +649,14 @@ def _is_sm103() -> bool: return current_platform.is_device_capability((10, 3)) +def _is_sm107() -> bool: + return current_platform.is_device_capability((10, 7)) + + def _low_latency_table() -> dict[tuple[int, int], ProjectionSpec] | None: """Measured dispatch table for the current device, or None if unsupported.""" - if _is_sm103(): + if _is_sm103() or _is_sm107(): + # SM107 reuses the SM103 table; see the module docstring. return KIMI_K3_PROJECTIONS if current_platform.is_device_capability((10, 0)): return KIMI_K3_PROJECTIONS_SM100