vllm.models.kimi_k3.nvidia.low_latency_gemm ¶

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

def autotune_kda_qkvg(model: nn.Module) -> None:
    """Autotune the FlashInfer QKVG GEMM before CUDA graph capture."""
    from flashinfer.gemm import mm_bf16

    from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention

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

    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="cute-dsl",
        )
    for child in children:
        child._projection_overlap_max_tokens = KDA_PROJECTION_OVERLAP_MAX_TOKENS

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, :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
            )

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:
            from flashinfer.gemm import mm_bf16

            return mm_bf16(
                hidden_states,
                qkvg_weight.t(),
                pdl=True,
                backend="cute-dsl",
            )
        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

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

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)