Skip to content

`vllm.models.kimi_k3.nvidia.kda_metadata` ¶

Kimi-K3 specialization of GDN attention metadata.

The request classification and cudagraph staging intentionally mirror GDNAttentionMetadataBuilder. Kimi-K3 builds the metadata required by its prefill KDA kernel internally, so this builder omits the shared FLA chunk metadata construction.

Functions:

stage_spec_decode_metadata –

Stage speculative-decode metadata into CUDA-graph buffers.

`stage_spec_decode_metadata(state_indices, query_start_loc, num_accepted_tokens, staged_state_indices, staged_query_start_loc, staged_num_accepted_tokens, *, num_spec_decodes)` ¶

Stage speculative-decode metadata into CUDA-graph buffers.

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

def stage_spec_decode_metadata(
    state_indices: torch.Tensor,
    query_start_loc: torch.Tensor,
    num_accepted_tokens: torch.Tensor,
    staged_state_indices: torch.Tensor,
    staged_query_start_loc: torch.Tensor,
    staged_num_accepted_tokens: torch.Tensor,
    *,
    num_spec_decodes: int,
) -> None:
    """Stage speculative-decode metadata into CUDA-graph buffers."""
    assert state_indices.is_cuda
    assert state_indices.ndim == 2
    batch_size, num_state_slots = staged_state_indices.shape
    BLOCK_ROWS = 32
    grid = (triton.cdiv(batch_size + 1, BLOCK_ROWS),)
    _stage_spec_decode_metadata_kernel[grid](
        state_indices,
        query_start_loc,
        num_accepted_tokens,
        staged_state_indices,
        staged_query_start_loc,
        staged_num_accepted_tokens,
        state_indices.stride(0),
        state_indices.stride(1),
        staged_state_indices.stride(0),
        staged_state_indices.stride(1),
        num_spec_decodes,
        batch_size,
        NUM_STATE_SLOTS=num_state_slots,
        BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots),
        NULL_STATE_ID=NULL_BLOCK_ID,
        BLOCK_ROWS=BLOCK_ROWS,
        num_warps=1,
        launch_pdl=_metadata_launch_pdl(),
    )