Skip to content

vllm.models.kimi_k3.amd.kda_metadata

ROCm Kimi-K3 specialization of GDN attention metadata.

The request classification and cudagraph staging intentionally mirror GDNAttentionMetadataBuilder. Only the FLA chunk metadata is built differently on device rather than on the host.

Functions:

build_conv1d_chunk_schedule(query_start_loc_cpu, block_m, device)

Host-built AITER conv1d chunk schedule: no D2H sync, own buffers.

Source code in vllm/models/kimi_k3/amd/kda_metadata.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def build_conv1d_chunk_schedule(
    query_start_loc_cpu: torch.Tensor, block_m: int, device: torch.device
) -> dict[str, int | torch.Tensor]:
    """Host-built AITER conv1d chunk schedule: no D2H sync, own buffers."""
    assert query_start_loc_cpu.device.type == "cpu"
    nums = (-(-query_start_loc_cpu.diff() // block_m)).to(torch.int64)
    tot = int(nums.sum())
    batch_ptr = torch.repeat_interleave(
        torch.arange(nums.numel(), dtype=torch.int32), nums
    )
    starts = torch.repeat_interleave(nums.cumsum(0) - nums, nums)
    token_chunk_offset_ptr = (torch.arange(tot, dtype=torch.int64) - starts).to(
        torch.int32
    )
    return {
        "tot": tot,
        "batch_ptr": _to_device(batch_ptr, device),
        "token_chunk_offset_ptr": _to_device(token_chunk_offset_ptr, device),
    }

prepare_chunk_metadata_device(cu_seqlens, cu_seqlens_cpu, chunk_size)

Build FLA chunk metadata on device, with no host<->device transfer.

Source code in vllm/models/kimi_k3/amd/kda_metadata.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def prepare_chunk_metadata_device(
    cu_seqlens: torch.Tensor,
    cu_seqlens_cpu: torch.Tensor,
    chunk_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Build FLA chunk metadata on device, with no host<->device transfer."""
    num_seqs = cu_seqlens_cpu.numel() - 1
    seq_lens = cu_seqlens_cpu[1:] - cu_seqlens_cpu[:-1]
    num_chunks = int(((seq_lens + chunk_size - 1) // chunk_size).sum())

    chunk_indices = torch.empty(
        num_chunks, 2, dtype=cu_seqlens.dtype, device=cu_seqlens.device
    )
    chunk_offsets = torch.empty(
        num_seqs + 1, dtype=torch.int64, device=cu_seqlens.device
    )
    _chunk_metadata_kernel[(num_seqs,)](
        cu_seqlens,
        chunk_indices,
        chunk_offsets,
        num_seqs,
        BT=chunk_size,
        BLOCK_N=max(_MIN_BLOCK_N, next_power_of_2(num_seqs + 1)),
        BLOCK_T=_BLOCK_T,
        num_warps=4,
    )
    return chunk_indices, chunk_offsets