vllm.models.kimi_k3.nvidia.ops.kda_conv

KDA prefill short convolution over the packed Q|K|V projection in one launch.

The stock prefill runs causal_conv1d_fn once per Q, K and V. Each call first casts the row-strided bf16 slice of the in_proj output to a dense FP32 copy (the conv state dtype), convolves it into an FP32 output and casts that back to bf16: per layer and rank at a 16K chunk that is 3 x (76 + 52 + 19) us, of which only the convolution reads and writes the data the layer needs.

This kernel reads the bf16 Q|K|V columns of the in_proj output in place, does the stock arithmetic in FP32 (bf16 -> FP32 is exact; same product order, FP32 weights, FP32 conv state, SiLU), and stores bf16 Q, K and V as three dense (tokens, dim) planes of one buffer, the layout the stock casts produce. FP32 -> bf16 rounds to nearest even, as the stock .to(bfloat16) does, so the outputs and the updated conv state are bit-identical to the stock path.

Programs map to (sequence, token chunk) on the device by a radix-32 search over query_start_loc, so no host-side chunk table is needed and the launch is CUDA graph safe. A program streams its tokens over one channel vector with the 3 history rows in registers and the next 8 input rows prefetched in a register queue. As in the stock kernel, only a sequence's first chunk reads its conv state, and that same program then writes the updated state.

Functions:

• packed_causal_conv1d_fn –

  Varlen causal depthwise conv of x (tokens, dim), dim = parts x part_dim.

• packed_causal_conv1d_supported –

  Whether packed_causal_conv1d_fn reproduces the stock conv bit for bit.

_config(part_dim)

(BLOCK_M unrolled tokens, NUM_BLOCKS blocks per program, BLOCK_N channels, warps). BLOCK_N divides part_dim so that a program writes one plane.

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

def _config(part_dim: int) -> tuple[int, int, int, int]:
    """(BLOCK_M unrolled tokens, NUM_BLOCKS blocks per program, BLOCK_N channels,
    warps). BLOCK_N divides part_dim so that a program writes one plane."""
    # Measured on GB300 at 2K-16K tokens x 3 x 1536 channels: one warp per
    # 128 channels and 64-token programs stream at ~4.8 TB/s of read + write.
    return 8, 8, 128 if part_dim % 128 == 0 else 64, 1

packed_causal_conv1d_fn(x, weight, bias, conv_states, query_start_loc, cache_indices, has_initial_state, part_dim, activation='silu', null_block_id=NULL_BLOCK_ID)

Varlen causal depthwise conv of x (tokens, dim), dim = parts x part_dim.

Same semantics as causal_conv1d_fn(x.T, weight, bias, conv_states, ...).T with conv_states (cache lines, dim, 3) updated in place; width-4 taps. Returns (dim // part_dim, tokens, part_dim) in x.dtype: plane i is the dense output of channels [i * part_dim, (i + 1) * part_dim).

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

def packed_causal_conv1d_fn(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None,
    conv_states: torch.Tensor,
    query_start_loc: torch.Tensor,
    cache_indices: torch.Tensor,
    has_initial_state: torch.Tensor,
    part_dim: int,
    activation: str | None = "silu",
    null_block_id: int | None = NULL_BLOCK_ID,
) -> torch.Tensor:
    """Varlen causal depthwise conv of ``x`` (tokens, dim), dim = parts x part_dim.

    Same semantics as ``causal_conv1d_fn(x.T, weight, bias, conv_states, ...).T``
    with ``conv_states`` (cache lines, dim, 3) updated in place; width-4 taps. Returns
    ``(dim // part_dim, tokens, part_dim)`` in ``x.dtype``: plane i is the dense
    output of channels [i * part_dim, (i + 1) * part_dim).
    """
    num_tokens, dim = x.shape
    num_seqs = query_start_loc.size(0) - 1
    out = torch.empty(
        (dim // part_dim, num_tokens, part_dim), dtype=x.dtype, device=x.device
    )
    if num_tokens == 0 or num_seqs == 0:
        return out
    block_m, num_blocks, block_n, num_warps = _config(part_dim)
    # Radix-32 levels of the kernel's program search: 32 ** steps >= num_seqs.
    search_steps = 0
    while 32**search_steps < num_seqs:
        search_steps += 1
    # Every sequence adds at most one partial chunk (see the kernel's program map).
    num_col_blocks = triton.cdiv(dim, block_n)
    grid = (
        (triton.cdiv(num_tokens, block_m * num_blocks) + num_seqs) * num_col_blocks,
    )
    _packed_causal_conv1d_fwd_kernel[grid](
        x,
        weight,
        bias,
        conv_states,
        cache_indices,
        has_initial_state,
        query_start_loc,
        out,
        num_seqs,
        search_steps,
        32 ** max(search_steps - 1, 0),
        num_col_blocks,
        part_dim,
        x.stride(0),
        weight.stride(0),
        weight.stride(1),
        conv_states.stride(0),
        conv_states.stride(1),
        conv_states.stride(2),
        cache_indices.stride(0),
        out.stride(0),
        out.stride(1),
        null_block_id if null_block_id is not None else -1,
        HAS_BIAS=bias is not None,
        HAS_NULL_BLOCK=null_block_id is not None,
        SILU_ACTIVATION=activation in ("silu", "swish", True),
        BLOCK_M=block_m,
        NUM_BLOCKS=num_blocks,
        BLOCK_N=block_n,
        num_warps=num_warps,
    )
    return out

packed_causal_conv1d_supported(x, weight, conv_states, part_dim)

Whether packed_causal_conv1d_fn reproduces the stock conv bit for bit.

The stock kernel computes in the conv state dtype; with an FP32 state that is FP32 for every input dtype, which this kernel does.

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

def packed_causal_conv1d_supported(
    x: torch.Tensor,
    weight: torch.Tensor,
    conv_states: torch.Tensor,
    part_dim: int,
) -> bool:
    """Whether ``packed_causal_conv1d_fn`` reproduces the stock conv bit for bit.

    The stock kernel computes in the conv state dtype; with an FP32 state that is
    FP32 for every input dtype, which this kernel does."""
    return (
        x.is_cuda
        and x.dim() == 2
        and x.stride(1) == 1
        and x.dtype in (torch.bfloat16, torch.float16, torch.float32)
        and conv_states.dtype == torch.float32
        and weight.dtype == torch.float32
        and weight.dim() == 2
        and weight.size(0) == x.size(1)
        and weight.size(1) == 4
        and conv_states.dim() == 3
        and conv_states.size(1) == x.size(1)
        and conv_states.size(2) >= 3
        and part_dim % 64 == 0
        and x.size(1) % part_dim == 0
    )