[Kernel] Enable Kimi-K3 SiTU on the CuteDSL MoE backend and the SM107 low-latency GEMM plan (#54606)
Signed-off-by: Bolin Sun <bolins@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Bolin Sun committed
Sep 3, 2026 at 12:49 UTC
d410fc12f30b4224fca71a2faca75dcd4a5c3557
2 files changed
+30
-4
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
+19
-1
@@ -62,6 +62,8 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
62
self.gemm1_alpha = quant_config.gemm1_alpha
63
self.gemm1_beta = quant_config.gemm1_beta
64
self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
65
+ self.situ_beta = moe_config.activation_situ_beta
66
+ self.situ_linear_beta = moe_config.activation_situ_linear_beta
67
68
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
69
layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
@@ -101,6 +103,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
103
MoEActivation.SWIGLUOAI,
104
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
105
MoEActivation.RELU2_NO_MUL,
106
+ MoEActivation.SITU,
107
)
108
109
@staticmethod
@@ -171,6 +174,19 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
174
"swiglu_beta": self.gemm1_beta,
175
"swiglu_limit": self.gemm1_clamp_limit,
176
}
177
+ elif activation == MoEActivation.SITU:
178
+ # The cute_dsl kernel keys SiTU on situ_beta and requires
179
+ # activation_type to stay a base type (ActivationType.Situ is
180
+ # rejected by normalize_cute_dsl_moe_activation_type), so the
181
+ # Swiglu base type is passed below and SiTU rides the betas.
182
+ if self.situ_beta is None:
183
+ raise ValueError(
184
+ "SITU activation requires moe_config.activation_situ_beta"
185
+ )
186
+ swiglu_params = {
187
+ "situ_beta": self.situ_beta,
188
+ "situ_linear_beta": self.situ_linear_beta,
189
+ }
190
swiglu_kwargs = {k: v for k, v in swiglu_params.items() if v is not None}
191
192
flashinfer_cute_dsl_fused_moe_nvfp4(
@@ -190,6 +206,8 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
206
num_local_experts=self.local_num_experts,
207
local_expert_offset=self.local_expert_offset,
208
moe_output=output,
193
- activation_type=activation_to_flashinfer_int(activation),
209
+ activation_type=activation_to_flashinfer_int(
210
+ MoEActivation.SILU if activation == MoEActivation.SITU else activation
211
+ ),
212
**swiglu_kwargs,
213
)
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+11
-3
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM90/SM100/SM103.
3
+"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM90/SM100/SM103/SM107.
4
5
Dispatch is purely by local ``(N, K)`` shape and token count ``M`` — the module
6
name plays no role. Each measured shape maps to a :class:`ProjectionSpec`
@@ -12,7 +12,10 @@ The supported capabilities carry separate measured tables:
12
:data:`KIMI_K3_PROJECTIONS` was tuned on B300 (SM103),
13
:data:`KIMI_K3_PROJECTIONS_SM100` on B200 (SM100), and
14
:data:`KIMI_K3_PROJECTIONS_SM90` on H200 (SM90). The per-(shape, M) winners
15
-genuinely differ between the parts, so the tables must not be merged.
15
+genuinely differ between the parts, so the tables must not be merged. SM107
16
+(Rubin) reuses the SM103 table: the plan was validated end-to-end on SM107
17
+hardware, but the per-M crossovers have not been re-measured there and may
18
+deserve their own table once retuned.
19
"""
20
21
from __future__ import annotations
@@ -646,9 +649,14 @@ def _is_sm103() -> bool:
649
return current_platform.is_device_capability((10, 3))
650
651
652
+def _is_sm107() -> bool:
653
+ return current_platform.is_device_capability((10, 7))
654
+
655
+
656
def _low_latency_table() -> dict[tuple[int, int], ProjectionSpec] | None:
657
"""Measured dispatch table for the current device, or None if unsupported."""
651
- if _is_sm103():
658
+ if _is_sm103() or _is_sm107():
659
+ # SM107 reuses the SM103 table; see the module docstring.
660
return KIMI_K3_PROJECTIONS
661
if current_platform.is_device_capability((10, 0)):
662
return KIMI_K3_PROJECTIONS_SM100