vllm.models.kimi_k3.nvidia.ops.dsa_indexer

Fused Triton prologues for the Kimi K3.1 DSA layer, replacing the small launch-bound kernels between the projection GEMMs and sparse_attn_indexer.

The quant helpers and dsa_fused_prologue follow fused_norm_rope in deepseek_v32/common/kernels.py (without RoPE); they are copied so this model does not depend on another model's private kernels.

Functions:

• dsa_fused_prologue –

  fused_norm_rope for the NoPE Kimi K3.1 DSA layer: the q_a / kv_a

• indexer_k_norm_quant_insert –

  LayerNorm + FP8 quant + cache insert of the indexer key, and the -1

• indexer_q_quant_fold_weights –

  Quantize index_q [T, H, D] to FP8 per head (ue8m0 scales) and

dsa_fused_prologue(q_c, kv_c, k_pe, index_k, q_norm_weight, kv_norm_weight, k_norm_weight, k_norm_bias, eps, topk_indices_buffer, mla_cache, mla_slot_mapping, index_cache, index_slot_mapping)

fused_norm_rope for the NoPE Kimi K3.1 DSA layer: the q_a / kv_a RMSNorms, the fp8_ds_mla insert of [kv_c_normed | k_pe], the indexer key LayerNorm + FP8 insert and the -1 fill of the top-k rows, in one launch.

mla_cache ([num_blocks, block_size, 656]) and index_cache ([num_blocks, block_size, head_dim + 4]) are uint8, possibly strided page views of a packed cache; None (profile runs) skips the insert. Rows past a slot mapping or with a negative slot are not inserted.

Returns (q_c_normed, kv_c_normed).

Source code in vllm/models/kimi_k3/nvidia/ops/dsa_indexer.py

def dsa_fused_prologue(
    q_c: torch.Tensor,
    kv_c: torch.Tensor,
    k_pe: torch.Tensor,
    index_k: torch.Tensor,
    q_norm_weight: torch.Tensor,
    kv_norm_weight: torch.Tensor,
    k_norm_weight: torch.Tensor,
    k_norm_bias: torch.Tensor,
    eps: float,
    topk_indices_buffer: torch.Tensor,
    mla_cache: torch.Tensor | None,
    mla_slot_mapping: torch.Tensor | None,
    index_cache: torch.Tensor | None,
    index_slot_mapping: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """``fused_norm_rope`` for the NoPE Kimi K3.1 DSA layer: the q_a / kv_a
    RMSNorms, the fp8_ds_mla insert of ``[kv_c_normed | k_pe]``, the indexer key
    LayerNorm + FP8 insert and the -1 fill of the top-k rows, in one launch.

    ``mla_cache`` (``[num_blocks, block_size, 656]``) and ``index_cache``
    (``[num_blocks, block_size, head_dim + 4]``) are uint8, possibly strided
    page views of a packed cache; ``None`` (profile runs) skips the insert. Rows
    past a slot mapping or with a negative slot are not inserted.

    Returns ``(q_c_normed, kv_c_normed)``.
    """
    num_tokens, q_dim = q_c.shape
    kv_dim, rope_dim, index_dim = kv_c.shape[1], k_pe.shape[1], index_k.shape[1]
    assert all(t.stride(-1) == 1 for t in (q_c, kv_c, k_pe, index_k))
    q_out = torch.empty((num_tokens, q_dim), dtype=q_c.dtype, device=q_c.device)
    kv_out = torch.empty((num_tokens, kv_dim), dtype=kv_c.dtype, device=kv_c.device)
    if num_tokens == 0:
        return q_out, kv_out

    # No cache (profile runs): an empty slot mapping disables the insert.
    def placeholder() -> tuple[torch.Tensor, torch.Tensor]:
        return (
            torch.empty(1, 1, 4, dtype=torch.uint8, device=q_c.device),
            torch.empty(0, dtype=torch.int64, device=q_c.device),
        )

    if mla_cache is None or mla_slot_mapping is None:
        mla_cache, mla_slot_mapping = placeholder()
    if index_cache is None or index_slot_mapping is None:
        index_cache, index_slot_mapping = placeholder()
    mla = mla_cache.view(torch.uint8)
    index = index_cache.view(torch.uint8)
    mx_scale = mqa_logits_use_mx_scales()
    topk = topk_indices_buffer.shape[-1]
    use_pdl = current_platform.is_arch_support_pdl()
    _dsa_fused_prologue_kernel[(num_tokens, 4)](
        q_c,
        q_c.stride(0),
        q_norm_weight,
        q_out,
        kv_c,
        kv_c.stride(0),
        kv_norm_weight,
        kv_out,
        k_pe,
        k_pe.stride(0),
        index_k,
        index_k.stride(0),
        k_norm_weight,
        k_norm_bias,
        eps,
        mla_slot_mapping,
        mla_slot_mapping.shape[0],
        mla.view(_FP8_DTYPE),
        mla.view(torch.float32),
        mla.view(torch.bfloat16),
        mla.shape[1],
        mla.stride(0),
        mla.stride(1),
        index_slot_mapping,
        index_slot_mapping.shape[0],
        index.view(_FP8_DTYPE),
        index.view(torch.int32 if mx_scale else torch.float32),
        index.shape[1],
        index.stride(0),
        topk_indices_buffer,
        topk_indices_buffer.stride(0),
        Q_DIM=q_dim,
        Q_BLOCK=triton.next_power_of_2(q_dim),
        KV_DIM=kv_dim,
        ROPE_DIM=rope_dim,
        INDEX_DIM=index_dim,
        TOPK=topk,
        TOPK_BLOCK=min(1024, triton.next_power_of_2(topk)),
        MX_SCALE=mx_scale,
        USE_PDL=use_pdl,
        FP8_MAX=_FP8_MAX,
        USE_FNUZ=_USE_FNUZ,
        launch_pdl=use_pdl,
    )
    return q_out, kv_out

indexer_k_norm_quant_insert(k, norm_weight, norm_bias, norm_eps, slot_mapping, kv_cache, topk_indices_buffer, num_tokens)

LayerNorm + FP8 quant + cache insert of the indexer key, and the -1 fill of the top-k rows [:num_tokens].

kv_cache is [num_blocks, block_size, head_dim + 4] uint8, possibly a strided page view of a packed cache. Rows past slot_mapping or with a negative slot are not inserted.

Source code in vllm/models/kimi_k3/nvidia/ops/dsa_indexer.py

def indexer_k_norm_quant_insert(
    k: torch.Tensor,
    norm_weight: torch.Tensor,
    norm_bias: torch.Tensor,
    norm_eps: float,
    slot_mapping: torch.Tensor,
    kv_cache: torch.Tensor,
    topk_indices_buffer: torch.Tensor,
    num_tokens: int,
) -> None:
    """LayerNorm + FP8 quant + cache insert of the indexer key, and the -1
    fill of the top-k rows ``[:num_tokens]``.

    ``kv_cache`` is ``[num_blocks, block_size, head_dim + 4]`` uint8, possibly
    a strided page view of a packed cache. Rows past ``slot_mapping`` or with
    a negative slot are not inserted.
    """
    if num_tokens == 0:
        return
    head_dim = k.shape[-1]
    assert k.stride(-1) == 1
    assert kv_cache.dim() == 3 and kv_cache.shape[-1] == head_dim + 4
    assert triton.next_power_of_2(head_dim) == head_dim
    mx_scale = mqa_logits_use_mx_scales()
    u8_cache = kv_cache.view(torch.uint8)
    topk = topk_indices_buffer.shape[-1]
    _indexer_k_norm_quant_insert_kernel[(num_tokens, 2)](
        k,
        k.stride(0),
        norm_weight,
        norm_bias,
        norm_eps,
        slot_mapping,
        slot_mapping.shape[0],
        u8_cache.view(_FP8_DTYPE),
        u8_cache.view(torch.int32 if mx_scale else torch.float32),
        u8_cache.shape[1],
        u8_cache.stride(0),
        topk_indices_buffer,
        topk_indices_buffer.stride(0),
        HEAD_DIM=head_dim,
        TOPK=topk,
        TOPK_BLOCK=min(1024, triton.next_power_of_2(topk)),
        MX_SCALE=mx_scale,
        FP8_MAX=_FP8_MAX,
        USE_FNUZ=_USE_FNUZ,
        num_warps=1,
    )

indexer_q_quant_fold_weights(index_q, index_weights, weight_scale, padded_num_heads)

Quantize index_q [T, H, D] to FP8 per head (ue8m0 scales) and fold each head's scale into index_weights [T, H] together with weight_scale.

Returns (q_fp8 [T, P, D], weights [T, P] fp32) where P = padded_num_heads; heads [H, P) are zero.

Source code in vllm/models/kimi_k3/nvidia/ops/dsa_indexer.py

def indexer_q_quant_fold_weights(
    index_q: torch.Tensor,
    index_weights: torch.Tensor,
    weight_scale: float,
    padded_num_heads: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Quantize ``index_q`` ``[T, H, D]`` to FP8 per head (ue8m0 scales) and
    fold each head's scale into ``index_weights`` ``[T, H]`` together with
    ``weight_scale``.

    Returns ``(q_fp8 [T, P, D], weights [T, P] fp32)`` where ``P =
    padded_num_heads``; heads ``[H, P)`` are zero.
    """
    num_tokens, num_heads, head_dim = index_q.shape
    assert index_q.stride(-1) == 1
    assert index_weights.shape == (num_tokens, num_heads)
    assert triton.next_power_of_2(head_dim) == head_dim
    alloc = torch.empty if padded_num_heads == num_heads else torch.zeros
    q_fp8 = alloc(
        (num_tokens, padded_num_heads, head_dim),
        dtype=_FP8_DTYPE,
        device=index_q.device,
    )
    weights = alloc(
        (num_tokens, padded_num_heads), dtype=torch.float32, device=index_q.device
    )
    if num_tokens == 0:
        return q_fp8, weights
    _indexer_q_quant_fold_weights_kernel[(num_tokens, num_heads)](
        index_q,
        index_q.stride(0),
        index_q.stride(1),
        index_weights,
        index_weights.stride(0),
        q_fp8,
        q_fp8.stride(0),
        q_fp8.stride(1),
        weights,
        weights.stride(0),
        weight_scale,
        HEAD_DIM=head_dim,
        FP8_MAX=_FP8_MAX,
        USE_FNUZ=_USE_FNUZ,
        num_warps=1,
    )
    return q_fp8, weights