Kimi-K3 decode GEMM selection for unquantized BF16 on SM103.
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 holding the winning backend per token count. The static part of the decision is resolved once per module at install time into a small {M: call} plan, so the per-forward path is a single dict lookup.
Functions:
_KimiK3LowLatencyApply
Mixin: try the precomputed plan, else defer to the base method.
Source code in vllm/models/kimi_k3/nvidia/low_latency_gemm.py
| class _KimiK3LowLatencyApply:
"""Mixin: try the precomputed plan, else defer to the base method."""
def __init__(self, plan: dict[int, ResolvedCall]) -> None:
self._plan = plan
def apply(
self,
layer: nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
if (
bias is None
and not envs.VLLM_BATCH_INVARIANT
and _runtime_ok(x, layer.weight)
):
output = _run_plan(self._plan, x, layer.weight)
if output is not None:
return output
return super().apply(layer, x, bias) # type: ignore[misc]
|
enable_kimi_k3_low_latency_gemm(module, dtype)
Install shape-selected low-latency GEMMs and register CuTe warmups.
Modules are matched purely by type, an exactly-unquantized method, and a local (N, K) present in :data:KIMI_K3_PROJECTIONS.
Source code in vllm/models/kimi_k3/nvidia/low_latency_gemm.py
| def enable_kimi_k3_low_latency_gemm(
module: nn.Module,
dtype: torch.dtype,
) -> None:
"""Install shape-selected low-latency GEMMs and register CuTe warmups.
Modules are matched purely by type, an exactly-unquantized method, and a
local ``(N, K)`` present in :data:`KIMI_K3_PROJECTIONS`.
"""
if dtype != torch.bfloat16 or not _is_sm103():
return
warmup_configs: set[SkinnyGemmConfig] = set()
residual_warmup_configs: set[SkinnyGemmConfig] = set()
for child in module.modules():
is_linear = (
isinstance(child, LinearBase)
and type(child.quant_method) is UnquantizedLinearMethod
)
# ParallelLMHead is a VocabParallelEmbedding subclass; embed_tokens is
# the parent type, so isinstance already excludes it.
is_head = (
isinstance(child, ParallelLMHead)
and type(child.quant_method) is UnquantizedEmbeddingMethod
)
if not (is_linear or is_head):
continue
weight = getattr(child, "weight", None)
if weight is None or weight.dim() != 2:
continue
spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1]))
if spec is None:
continue
if is_linear:
child.quant_method = KimiK3LowLatencyLinearMethod(
_build_plan(spec), _build_residual_plan(spec)
)
else:
child.quant_method = KimiK3LowLatencyEmbeddingMethod(_build_plan(spec))
# Warm up only the configs measured for this module's local (N, K) so a
# TP8 deployment does not compile TP4 configs and vice versa.
warmup_configs.update(config for _, config in spec.cute_configs)
residual_warmup_configs.update(config for _, config in spec.residual_configs)
if shape_dynamic_skinny_gemm.is_available():
if warmup_configs:
shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs)
if residual_warmup_configs:
shape_dynamic_skinny_gemm.request_warmup_configs(
dtype, residual_warmup_configs, has_residual=True
)
|
select_kimi_k3_backend(num_tokens, n, k, *, has_residual=False)
Backend for a local (N, K) at num_tokens, or None to fall back.
Source code in vllm/models/kimi_k3/nvidia/low_latency_gemm.py
| def select_kimi_k3_backend(
num_tokens: int,
n: int,
k: int,
*,
has_residual: bool = False,
) -> Backend | None:
"""Backend for a local ``(N, K)`` at ``num_tokens``, or None to fall back."""
spec = KIMI_K3_PROJECTIONS.get((n, k))
return _backend_for(spec, num_tokens, has_residual) if spec is not None else None
|
try_low_latency_gemm(x, weight, residual=None)
Run the shape-selected low-latency kernel, or None to fall back.
Resolves the plan from the shape table on each call; production installs a precomputed plan (see :func:enable_kimi_k3_low_latency_gemm) and does not use this path.
Source code in vllm/models/kimi_k3/nvidia/low_latency_gemm.py
| def try_low_latency_gemm(
x: torch.Tensor,
weight: torch.Tensor,
residual: torch.Tensor | None = None,
) -> torch.Tensor | None:
"""Run the shape-selected low-latency kernel, or None to fall back.
Resolves the plan from the shape table on each call; production installs a
precomputed plan (see :func:`enable_kimi_k3_low_latency_gemm`) and does not
use this path.
"""
if envs.VLLM_BATCH_INVARIANT or not _is_sm103() or not _runtime_ok(x, weight):
return None
spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1]))
if spec is None:
return None
if residual is None:
return _run_plan(_build_plan(spec), x, weight)
if not _residual_ok(x, weight, residual):
return None
return _run_residual_plan(_build_residual_plan(spec), x, weight, residual)
|