vllm.models.deepseek_v41.nvidia.specialized_decode.attention.candidate_topk

Candidate-block selection with cooperative_topk for decode rows.

select_candidate_blocks scores every block of the padded logits row and runs torch.topk over all of them. At long context most of a decode row is padding, so here the score kernel also writes each row's valid block count and cooperative_topk scans only that prefix.

The selected set of non-negative ids matches select_candidate_blocks except for which blocks whose fp32 scores tie exactly at the selection boundary are kept. cooperative_topk keeps at most k (rows of up to 16384 blocks) or 1024 per cluster CTA (longer rows) candidates of its fp16 boundary bin, so a more crowded boundary bin is selected approximately. The pinned newest block is kept unless at least k blocks of the row score NaN or overflow fp16. Rows are not sorted and -1 padding may sit anywhere; every consumer of the candidate buffer treats a row as a set.

Functions:

• select_candidate_blocks_cooperative –

  select_candidate_blocks for decode rows, on cooperative_topk.

• supports_cooperative_candidates –

  Whether select_candidate_blocks_cooperative takes num_rows rows of

select_candidate_blocks_cooperative(logits, row_ends, block_size, out)

select_candidate_blocks for decode rows, on cooperative_topk.

Parameters:

• logits

  (Tensor) –

  [rows, width] fp32 logits; row i is valid below row_ends[i].

• row_ends

  (Tensor) –

  Contiguous [rows] int32 row lengths in logits columns.

• block_size

  (int) –

  Columns per candidate block.

• out

  (Tensor) –

  Contiguous [rows, k] int32 candidate block ids, -1 padded; see supports_cooperative_candidates for the supported shapes.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/candidate_topk.py

def select_candidate_blocks_cooperative(
    logits: torch.Tensor,
    row_ends: torch.Tensor,
    block_size: int,
    out: torch.Tensor,
) -> None:
    """`select_candidate_blocks` for decode rows, on ``cooperative_topk``.

    Args:
        logits: [rows, width] fp32 logits; row ``i`` is valid below
            ``row_ends[i]``.
        row_ends: Contiguous [rows] int32 row lengths in logits columns.
        block_size: Columns per candidate block.
        out: Contiguous [rows, k] int32 candidate block ids, -1 padded; see
            `supports_cooperative_candidates` for the supported shapes.

    """
    rows, width = logits.shape
    k = out.shape[1]
    nblocks = triton.cdiv(width, block_size)
    assert supports_cooperative_candidates(rows, nblocks, k)
    assert out.dtype == torch.int32 and out.shape[0] == rows and out.is_contiguous()
    scores = logits.new_empty((rows, nblocks))
    lengths = torch.empty(rows, dtype=torch.int32, device=logits.device)
    _block_scores_kernel[(rows, triton.cdiv(nblocks, 128))](
        logits,
        row_ends,
        scores,
        lengths,
        out,
        *logits.stride(),
        width,
        nblocks,
        block_size,
        k,
        128,
    )
    (workspace,) = current_workspace_manager().get_simultaneous(
        ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
    )
    # Rows of at most k blocks take all of them and pad with -1; the last
    # argument (max_seq_len) is unused by the kernel.
    torch.ops._C.cooperative_topk(scores, lengths, out, workspace, k, nblocks)
    _finalize_candidates_kernel[(rows, triton.cdiv(k, 1024))](
        out, scores, lengths, nblocks, k, 1024
    )

supports_cooperative_candidates(num_rows, num_blocks, k)

Whether select_candidate_blocks_cooperative takes num_rows rows of num_blocks blocks and k candidates.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/candidate_topk.py

def supports_cooperative_candidates(num_rows: int, num_blocks: int, k: int) -> bool:
    """Whether `select_candidate_blocks_cooperative` takes ``num_rows`` rows of
    ``num_blocks`` blocks and ``k`` candidates."""
    return (
        0 < num_rows <= AUTO_COOPERATIVE_MAX_ROWS
        and k in (512, 1024, 2048)
        # cooperative_topk reads the scores with TMA: 16-byte aligned rows.
        and num_blocks % 4 == 0
    )