vllm.models.kimi_k3.nvidia.low_latency_gemm

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 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.

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. SM107 (Rubin) reuses the SM103 table: the plan was validated end-to-end on SM107 hardware, but apart from 3584x7168 the per-M crossovers have not been re-measured there and may deserve their own table once retuned.

Functions:

• autotune_kda_qkvg –

  Autotune the supported QKVG GEMM before CUDA graph capture.

• enable_kimi_k3_low_latency_gemm –

  Install shape-selected low-latency GEMMs and register CuTe warmups.

• request_low_latency_warmup –

  Warm up the measured CuTe configs for a weight that no linear module owns

• run_kda_projection_overlap –

  Run the TP8 KDA decode projection branches concurrently.

• select_kimi_k3_backend –

  Backend for a local (N, K) at num_tokens, or None to fall back.

• try_low_latency_gemm –

  Run the shape-selected low-latency kernel, or None to fall back.

_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:
        super().__init__()
        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]

_kda_qkvg_flashinfer_backend()

Return the tested FlashInfer backend for the low-M QKVG branch.

FlashInfer's cute-dsl BF16 GEMM supports SM100 and SM103, but rejects SM90 during backend validation. Other capabilities use torch.mm for QKVG while retaining the concurrent F_A/beta and F_B branch.

Source code in vllm/models/kimi_k3/nvidia/low_latency_gemm.py

def _kda_qkvg_flashinfer_backend() -> Literal["cute-dsl"] | None:
    """Return the tested FlashInfer backend for the low-M QKVG branch.

    FlashInfer's ``cute-dsl`` BF16 GEMM supports SM100 and SM103, but rejects
    SM90 during backend validation. Other capabilities use ``torch.mm`` for
    QKVG while retaining the concurrent F_A/beta and F_B branch.
    """
    if current_platform.is_device_capability(
        (10, 0)
    ) or current_platform.is_device_capability((10, 3)):
        return "cute-dsl"
    return None

_low_latency_table()

Measured dispatch table for the current device, or None if unsupported.

Source code in vllm/models/kimi_k3/nvidia/low_latency_gemm.py

def _low_latency_table() -> dict[tuple[int, int], ProjectionSpec] | None:
    """Measured dispatch table for the current device, or None if unsupported."""
    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
    if current_platform.is_device_capability((9, 0)):
        return KIMI_K3_PROJECTIONS_SM90
    return None

autotune_kda_qkvg(model)

Autotune the supported QKVG GEMM before CUDA graph capture.

Source code in vllm/models/kimi_k3/nvidia/low_latency_gemm.py

def autotune_kda_qkvg(model: nn.Module) -> None:
    """Autotune the supported QKVG GEMM before CUDA graph capture."""
    from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention

    backend = _kda_qkvg_flashinfer_backend()
    children: list[KimiK3DeltaAttention] = []
    weights_by_shape: dict[
        tuple[int, int, int, torch.dtype, torch.device], torch.Tensor
    ] = {}
    for child in model.modules():
        if not isinstance(child, KimiK3DeltaAttention):
            continue
        if child._projection_overlap_max_tokens <= 0:
            continue
        children.append(child)
        if backend is not None:
            qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE]
            shape = (
                KDA_PROJECTION_OVERLAP_MAX_TOKENS,
                qkvg_weight.shape[0],
                qkvg_weight.shape[1],
                qkvg_weight.dtype,
                qkvg_weight.device,
            )
            weights_by_shape.setdefault(shape, qkvg_weight)

    if backend is not None:
        from flashinfer.gemm import mm_bf16

        for shape, qkvg_weight in weights_by_shape.items():
            num_tokens = shape[0]
            hidden_states = torch.empty(
                (num_tokens, qkvg_weight.shape[1]),
                dtype=qkvg_weight.dtype,
                device=qkvg_weight.device,
            )
            mm_bf16(
                hidden_states,
                qkvg_weight.t(),
                pdl=True,
                backend=backend,
            )
    for child in children:
        child._projection_overlap_max_tokens = KDA_PROJECTION_OVERLAP_MAX_TOKENS

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 the current device's measured table (:data:KIMI_K3_PROJECTIONS on SM103 and SM107, :data:KIMI_K3_PROJECTIONS_SM100 on SM100, :data:KIMI_K3_PROJECTIONS_SM90 on SM90).

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 the current device's measured table
    (:data:`KIMI_K3_PROJECTIONS` on SM103 and SM107,
    :data:`KIMI_K3_PROJECTIONS_SM100` on SM100, :data:`KIMI_K3_PROJECTIONS_SM90`
    on SM90).
    """
    if dtype != torch.bfloat16:
        return
    table = _low_latency_table()
    if table is None:
        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 = table.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 _enable_kda_projection_overlap(module):
        from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import (
            kda_skinny_gemm,
        )

        warmup_configs.update((*KDA_QKVG_CONFIGS.values(), KDA_M1_FAB_CONFIG))
        kda_skinny_gemm.request_warmup(
            set(range(2, KDA_SKINNY_N_MAX_TOKENS + 1)),
            set(range(2, KDA_SKINNY_K_MAX_TOKENS + 1)),
        )

    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
            )

request_low_latency_warmup(weight)

Warm up the measured CuTe configs for a weight that no linear module owns (enable_kimi_k3_low_latency_gemm only covers modules).

Source code in vllm/models/kimi_k3/nvidia/low_latency_gemm.py

def request_low_latency_warmup(weight: torch.Tensor) -> None:
    """Warm up the measured CuTe configs for a weight that no linear module owns
    (``enable_kimi_k3_low_latency_gemm`` only covers modules)."""
    table = _low_latency_table()
    spec = table.get(tuple(weight.shape)) if table is not None else None
    if spec is not None and shape_dynamic_skinny_gemm.is_available():
        shape_dynamic_skinny_gemm.request_warmup_configs(
            weight.dtype, (config for _, config in spec.cute_configs)
        )

run_kda_projection_overlap(hidden_states, packed_weight, f_b_weight, aux_stream, events)

Run the TP8 KDA decode projection branches concurrently.

Parameters:

• hidden_states

  (Tensor) –

  Packed BF16 input with shape [M, 7168].

• packed_weight

  (Tensor) –

  Existing Q/K/V/G/F_A/beta/pad weight with shape [6288, 7168].

• f_b_weight

  (Tensor) –

  F_B weight with shape [1536, 128].

• aux_stream

  (Stream) –

  Stream for the F_A/beta then F_B branch.

• events

  (tuple[Event, Event]) –

  Start and completion events for the stream fork and join.

Returns:

• tuple[Tensor, Tensor, Tensor] –

  The Q/K/V/G projection, F_B output, and beta projection.

Source code in vllm/models/kimi_k3/nvidia/low_latency_gemm.py

def run_kda_projection_overlap(
    hidden_states: torch.Tensor,
    packed_weight: torch.Tensor,
    f_b_weight: torch.Tensor,
    aux_stream: torch.cuda.Stream,
    events: tuple[torch.cuda.Event, torch.cuda.Event],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Run the TP8 KDA decode projection branches concurrently.

    Args:
        hidden_states: Packed BF16 input with shape ``[M, 7168]``.
        packed_weight: Existing Q/K/V/G/F_A/beta/pad weight with shape
            ``[6288, 7168]``.
        f_b_weight: F_B weight with shape ``[1536, 128]``.
        aux_stream: Stream for the F_A/beta then F_B branch.
        events: Start and completion events for the stream fork and join.

    Returns:
        The Q/K/V/G projection, F_B output, and beta projection.

    """
    num_tokens = hidden_states.shape[0]
    qkvg_weight = packed_weight[:_KDA_QKVG_SIZE]

    def run_qkvg() -> torch.Tensor:
        if config := KDA_QKVG_CONFIGS.get(num_tokens):
            return shape_dynamic_skinny_gemm(
                hidden_states,
                qkvg_weight,
                config,
                None,
            )
        if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS and (
            backend := _kda_qkvg_flashinfer_backend()
        ):
            from flashinfer.gemm import mm_bf16

            return mm_bf16(
                hidden_states,
                qkvg_weight.t(),
                pdl=True,
                backend=backend,
            )
        return torch.mm(hidden_states, qkvg_weight.t())

    def run_fab_fb() -> tuple[torch.Tensor, torch.Tensor]:
        if num_tokens == 1:
            projected_fab = shape_dynamic_skinny_gemm(
                hidden_states,
                packed_weight[_KDA_QKVG_SIZE : _KDA_QKVG_SIZE + _KDA_FAB_SIZE],
                KDA_M1_FAB_CONFIG,
                None,
            )
            f_a, beta = projected_fab.split([128, 12], dim=-1)
            f_a = f_a.as_strided((1, 128), (128, 1))
            g1 = torch.empty(
                (1, 1536), dtype=hidden_states.dtype, device=hidden_states.device
            )
            ops.dsv3_fused_a_gemm(g1, f_a, f_b_weight.t(), enable_pdl=True)
        else:
            from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import (
                kda_skinny_gemm,
            )

            if num_tokens <= KDA_SKINNY_N_MAX_TOKENS:
                projected_fab = kda_skinny_gemm.run_n(
                    hidden_states,
                    packed_weight[_KDA_QKVG_SIZE:],
                )
            else:
                projected_fab = torch.mm(
                    hidden_states,
                    packed_weight[_KDA_QKVG_SIZE:].t(),
                )
            beta = projected_fab[:, 128:140]
            if num_tokens <= KDA_SKINNY_K_MAX_TOKENS:
                g1 = kda_skinny_gemm.run_k(projected_fab, f_b_weight)
            else:
                g1 = torch.mm(projected_fab[:, :128], f_b_weight.t())
        return g1, beta

    projected_qkvg, (g1, beta) = maybe_execute_in_parallel(
        run_qkvg,
        run_fab_fb,
        events[0],
        events[1],
        aux_stream,
    )
    return projected_qkvg, g1, beta

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 _runtime_ok(x, weight):
        return None
    table = _low_latency_table()
    if table is None:
        return None
    spec = table.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)