Skip to content

vllm.models.deepseek_v41.amd.ops.topk_pack

One-launch packing of the DeepSeek V4.1 decode top-k into global slots.

compute_global_topk_ragged_indices_and_indptr in vLLM's V4.1 ROCm model runs once a step for each index source layer, 8 times a step for DeepSeek-V4.1-Flash. For a batch that this file does not take, it calls _compute_global_topk_ragged_indices_and_indptr_old, which launches 4 kernels: the per-row count of valid top-k entries, the zero fill and the cumsum of indptr, and the mapping of the entries through the compressed cache's block table. On MI325X they took 17.4 us of each index layer at 128k context, inside the decode graph.

_topk_pack_kernel writes the same 3 tensors with one program: the row lengths (0 for a padded token), indptr, and the global slots of the first length entries of each row (-1 for a negative entry). It does not write the ragged entries past indptr[-1], which _compute_global_topk_ragged_indices_and_indptr_old also leaves uninitialized.

On gfx942 with --monokernel (ops.enabled()), batches of up to MAX_ROWS rows use it. tests/kernels/models/amd/test_dsv41_gfx942_decode_meta.py compares it with _compute_global_topk_ragged_indices_and_indptr_old.

Functions:

  • launch –

    (global_topk_ragged, topk_indptr, topk_lens) from one launch, as

  • pack –

    The fused packing, or None when the batch takes

launch(topk_indices, token_to_req_indices, block_table, block_size, is_valid_token, poison=False)

(global_topk_ragged, topk_indptr, topk_lens) from one launch, as compute_global_topk_ragged_indices_and_indptr returns them. poison fills the outputs with -7 first, so that a test sees every value the launch does not write, even when the allocator hands back the last call's memory.

Source code in vllm/models/deepseek_v41/amd/ops/topk_pack.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def launch(
    topk_indices,
    token_to_req_indices,
    block_table,
    block_size,
    is_valid_token,
    poison=False,
):
    """(global_topk_ragged, topk_indptr, topk_lens) from one launch, as
    compute_global_topk_ragged_indices_and_indptr returns them. poison fills
    the outputs with -7 first, so that a test sees every value the launch
    does not write, even when the allocator hands back the last call's
    memory."""
    num_tokens, topk = topk_indices.shape
    dev = topk_indices.device
    lens = torch.empty(num_tokens, dtype=torch.int32, device=dev)
    indptr = torch.empty(num_tokens + 1, dtype=torch.int32, device=dev)
    ragged = torch.empty(num_tokens * topk, dtype=torch.int32, device=dev)
    if poison:
        for t in (lens, indptr, ragged):
            t.fill_(-7)
    _topk_pack_kernel[(1,)](
        ragged,
        indptr,
        lens,
        topk_indices,
        topk_indices.stride(0),
        is_valid_token,
        token_to_req_indices,
        block_table,
        block_table.stride(0),
        block_size,
        num_tokens,
        topk,
        ROWS=triton.next_power_of_2(num_tokens),
        BLOCK=triton.next_power_of_2(topk),
        num_warps=8,
    )
    return ragged, indptr, lens

pack(topk_indices, token_to_req_indices, block_table, block_size, is_valid_token)

The fused packing, or None when the batch takes _compute_global_topk_ragged_indices_and_indptr_old.

Source code in vllm/models/deepseek_v41/amd/ops/topk_pack.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def pack(
    topk_indices,
    token_to_req_indices,
    block_table,
    block_size,
    is_valid_token,
):
    """The fused packing, or None when the batch takes
    _compute_global_topk_ragged_indices_and_indptr_old."""
    if not enabled():
        return None
    topk_indices = topk_indices.reshape(topk_indices.shape[0], -1)
    num_tokens = topk_indices.shape[0]
    if (
        not 0 < num_tokens <= MAX_ROWS
        or topk_indices.stride(1) != 1
        or topk_indices.dtype != torch.int32
        or is_valid_token.shape[0] != num_tokens
    ):
        return None
    return launch(
        topk_indices,
        token_to_req_indices,
        block_table,
        block_size,
        is_valid_token,
    )