vllm.models.kimi_k3.nvidia.specialized_decode.recoverssm.kernels

Triton kernels of the specialized decode's deferred RecoverSSM: the stock RecoverSSM verify with its records written to each request's record slot, and the commit-plan steps that defer (_defer_commit_kernel) or flush (_prepare_flush_plan_kernel) a request's pending replay.

Functions:

• kda_recoverssm_verify –

  Verify a KDA speculative window without modifying its checkpoint. Request

_defer_commit_kernel(state_indices_ptr, commit_lens_ptr, record_rows_ptr, pending_ptr, eager_ptr, null_block_id, stride_state_indices, HAS_EAGER)

Deferred RecoverSSM after a plan: request b's commit length becomes the pending replay of its record slot (and no commit runs for it), unless the row is committed now (eager: pending 0, the plan kept). Rows without a state slot (padding, dummy batches) or a record slot leave every pending count.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/recoverssm/kernels.py

@triton.heuristics({"HAS_EAGER": lambda args: args["eager_ptr"] is not None})
@triton.jit
def _defer_commit_kernel(
    state_indices_ptr,
    commit_lens_ptr,
    record_rows_ptr,
    pending_ptr,
    eager_ptr,
    null_block_id,
    stride_state_indices,
    HAS_EAGER: tl.constexpr,
):
    """Deferred RecoverSSM after a plan: request b's commit length becomes the
    pending replay of its record slot (and no commit runs for it), unless the row is
    committed now (``eager``: pending 0, the plan kept). Rows without a state slot
    (padding, dummy batches) or a record slot leave every pending count."""
    spec_idx = tl.program_id(0)
    record_row = tl.load(record_rows_ptr + spec_idx)
    source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices)
    if (record_row < 0) | (source_state_idx <= null_block_id):
        return
    commit_len = tl.load(commit_lens_ptr + spec_idx)
    eager = False
    if HAS_EAGER:
        eager = tl.load(eager_ptr + spec_idx) != 0
    if eager:
        tl.store(pending_ptr + record_row, 0)
    else:
        tl.store(pending_ptr + record_row, commit_len)
        tl.store(commit_lens_ptr + spec_idx, 0)

_prepare_flush_plan_kernel(state_indices_ptr, record_rows_ptr, pending_ptr, commit_lens_ptr, final_state_indices_ptr, boundary_state_indices_ptr, boundary_recovery_lens_ptr, null_block_id, stride_state_indices, SPEC_QUERY_LEN)

The commit plan folding each request's pending replay into its checkpoint in place (deferred RecoverSSM steps that the stock verify runs).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/recoverssm/kernels.py

@triton.jit
def _prepare_flush_plan_kernel(
    state_indices_ptr,
    record_rows_ptr,
    pending_ptr,
    commit_lens_ptr,
    final_state_indices_ptr,
    boundary_state_indices_ptr,
    boundary_recovery_lens_ptr,
    null_block_id,
    stride_state_indices,
    SPEC_QUERY_LEN: tl.constexpr,
):
    """The commit plan folding each request's pending replay into its checkpoint
    in place (deferred RecoverSSM steps that the stock verify runs)."""
    spec_idx = tl.program_id(0)
    source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices)
    record_row = tl.load(record_rows_ptr + spec_idx)
    commit_len = 0
    if record_row >= 0:
        commit_len = tl.load(pending_ptr + record_row)
    commit_len = tl.minimum(tl.maximum(commit_len, 0), SPEC_QUERY_LEN)
    valid = (source_state_idx > null_block_id) & (commit_len > 0)
    tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0))
    tl.store(final_state_indices_ptr + spec_idx, source_state_idx)
    tl.store(boundary_state_indices_ptr + spec_idx, null_block_id)
    tl.store(boundary_recovery_lens_ptr + spec_idx, 0)

kda_recoverssm_verify(q, k, v, raw_g, raw_beta, A_log, dt_bias, lower_bound, checkpoint_state, records, query_start_loc, state_indices, spec_query_len, out=None, record_indices=None)

Verify a KDA speculative window without modifying its checkpoint. Request b writes its verify records to row b of records, or to row record_indices[b] (deferred RecoverSSM: records indexed by request slot).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/recoverssm/kernels.py

def kda_recoverssm_verify(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    raw_g: torch.Tensor,
    raw_beta: torch.Tensor,
    A_log: torch.Tensor,
    dt_bias: torch.Tensor,
    lower_bound: float | None,
    checkpoint_state: torch.Tensor,
    records: "KDARecoverSSMRecords",
    query_start_loc: torch.Tensor,
    state_indices: torch.Tensor,
    spec_query_len: int,
    out: torch.Tensor | None = None,
    record_indices: torch.Tensor | None = None,
) -> torch.Tensor:
    """Verify a KDA speculative window without modifying its checkpoint. Request
    ``b`` writes its verify records to row ``b`` of ``records``, or to row
    ``record_indices[b]`` (deferred RecoverSSM: records indexed by request slot)."""
    if q.ndim != 4 or q.shape[0] != 1:
        raise ValueError("KDA RecoverSSM q must have shape [1, tokens, heads, dim]")
    _, total_tokens, num_heads, key_dim = q.shape
    value_dim = v.shape[-1]
    if k.shape != q.shape or v.shape != (1, total_tokens, num_heads, value_dim):
        raise ValueError("KDA RecoverSSM q, k, and v shapes are incompatible")
    if raw_g.shape != q.shape or raw_beta.shape != (1, total_tokens, num_heads):
        raise ValueError("KDA RecoverSSM gate or beta shape is incompatible")
    if any(tensor.stride()[2:] != (key_dim, 1) for tensor in (q, k, raw_g)):
        raise ValueError("KDA RecoverSSM q, k, and gate heads must be contiguous")
    if v.stride()[2:] != (value_dim, 1) or raw_beta.stride(2) != 1:
        raise ValueError("KDA RecoverSSM v and beta heads must be contiguous")
    if checkpoint_state.shape[1:] != (
        num_heads,
        value_dim,
        key_dim,
    ):
        raise ValueError("KDA RecoverSSM checkpoint shape is incompatible")
    correction_cache, kg_cache = records.correction, records.key_gate
    if A_log.shape != (num_heads,) or dt_bias.numel() != num_heads * key_dim:
        raise ValueError("KDA RecoverSSM gate parameters are incompatible")
    if not A_log.is_contiguous() or not dt_bias.is_contiguous():
        raise ValueError("KDA RecoverSSM gate parameters must be contiguous")
    batch = state_indices.shape[0]
    if query_start_loc.shape[0] != batch + 1:
        raise ValueError("KDA RecoverSSM query metadata is incompatible")
    if total_tokens > batch * spec_query_len:
        raise ValueError(
            "KDA RecoverSSM speculative decode input exceeds its activation capacity"
        )
    if out is None:
        out = torch.empty_like(v)
    if out.shape != v.shape:
        raise ValueError("KDA RecoverSSM output shape is incompatible")
    if out.stride()[2:] != (value_dim, 1):
        raise ValueError("KDA RecoverSSM output heads must be contiguous")
    device = q.device
    if any(
        tensor.device != device
        for tensor in (
            k,
            v,
            raw_g,
            raw_beta,
            A_log,
            dt_bias,
            checkpoint_state,
            correction_cache,
            kg_cache,
            query_start_loc,
            state_indices,
            out,
        )
    ):
        raise ValueError("KDA RecoverSSM inputs must be on the same device")
    if total_tokens == 0:
        return out

    block_k = triton.next_power_of_2(key_dim)
    block_v = min(triton.next_power_of_2(value_dim), 32)
    grid = (triton.cdiv(value_dim, block_v), batch, num_heads)
    _kda_recoverssm_verify_kernel[grid](
        q,
        k,
        v,
        raw_g,
        raw_beta,
        A_log,
        dt_bias,
        checkpoint_state,
        correction_cache,
        kg_cache,
        out,
        query_start_loc,
        state_indices,
        record_indices,
        lower_bound or 0.0,
        NULL_BLOCK_ID,
        q.stride(1),
        k.stride(1),
        v.stride(1),
        raw_g.stride(1),
        raw_beta.stride(1),
        checkpoint_state.stride(0),
        checkpoint_state.stride(1),
        checkpoint_state.stride(2),
        checkpoint_state.stride(3),
        correction_cache.stride(0),
        correction_cache.stride(1),
        correction_cache.stride(2),
        correction_cache.stride(3),
        kg_cache.stride(0),
        kg_cache.stride(1),
        kg_cache.stride(2),
        kg_cache.stride(3),
        out.stride(1),
        query_start_loc.stride(0),
        state_indices.stride(0),
        K=key_dim,
        V=value_dim,
        BK=block_k,
        BV=block_v,
        SPEC_QUERY_LEN=spec_query_len,
        USE_LOWER_BOUND=lower_bound is not None,
        num_warps=4,
        num_stages=2,
    )
    return out