Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,8 @@ def __init__(
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)
Expand Down Expand Up @@ -101,6 +103,7 @@ def _supports_activation(activation: MoEActivation) -> bool:
MoEActivation.SWIGLUOAI,
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
MoEActivation.RELU2_NO_MUL,
MoEActivation.SITU,
)

@staticmethod
Expand Down Expand Up @@ -171,6 +174,19 @@ def apply(
"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(
Expand All @@ -190,6 +206,8 @@ def apply(
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,
)
14 changes: 11 additions & 3 deletions vllm/models/kimi_k3/nvidia/low_latency_gemm.py
Original file line number Diff line number Diff line change
@@ -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`
Expand All @@ -12,7 +12,10 @@
: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
Expand Down Expand Up @@ -626,9 +629,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
Expand Down
Loading