Skip to content

vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl

Functions:

_dcp_merge_topk_and_interleave(vllm_config)

(topk, cp_interleave) the DCP merge runs with.

Kpool indexers (GLM-5.x) merge pool ids: topk and the interleave are both in units of index_kpool tokens.

Source code in vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def _dcp_merge_topk_and_interleave(vllm_config: Any) -> tuple[int, int]:
    """(topk, cp_interleave) the DCP merge runs with.

    Kpool indexers (GLM-5.x) merge pool ids: topk and the interleave are both
    in units of ``index_kpool`` tokens.
    """
    hf_text_config = vllm_config.model_config.hf_text_config
    topk = hf_text_config.index_topk
    cp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size
    kpool = getattr(hf_text_config, "index_kpool", None) or 1
    if kpool > 1:
        topk //= kpool
        cp_interleave //= kpool
    return topk, cp_interleave

dcp_topk_id_bits(num_positions)

Bits that hold every candidate id (a position below num_positions).

Source code in vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py 
51
52
53
def dcp_topk_id_bits(num_positions: int) -> int:
    """Bits that hold every candidate id (a position below num_positions)."""
    return max(1, min(32, (num_positions - 1).bit_length()))

dcp_topk_score_bits(logits_dtype)

High score bits that rank: BF16 logits, upcast to FP32 for the top-k, have a zero low half. Derived from the configured logits dtype, so every DCP rank (even one with no logits this step) and the warm-up agree.

Source code in vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py 
44
45
46
47
48
def dcp_topk_score_bits(logits_dtype: torch.dtype) -> int:
    """High score bits that rank: BF16 logits, upcast to FP32 for the top-k,
    have a zero low half. Derived from the configured logits dtype, so every
    DCP rank (even one with no logits this step) and the warm-up agree."""
    return 16 if logits_dtype == torch.bfloat16 else 32

stable_topk_from_gathered_candidates_cutedsl(gathered, topk, out=None, score_bits=32, id_bits=32)

Select each row's global top-topk from all ranks' candidates.

gathered is rank-major [world_size, rows, topk, 2], as produced by a dim-0 all-gather of each rank's [rows, topk, 2] candidates, and is read in place: a dim-1 gather into [rows, world_size * topk, 2] would cost a full reorder copy. The key ranks only the high score_bits of each score (see dcp_topk_score_bits) and the low id_bits of each id.

Source code in vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def stable_topk_from_gathered_candidates_cutedsl(
    gathered: torch.Tensor,
    topk: int,
    out: torch.Tensor | None = None,
    score_bits: int = 32,
    id_bits: int = 32,
) -> torch.Tensor:
    """Select each row's global top-``topk`` from all ranks' candidates.

    ``gathered`` is rank-major ``[world_size, rows, topk, 2]``, as produced by a
    dim-0 all-gather of each rank's ``[rows, topk, 2]`` candidates, and is read
    in place: a dim-1 gather into ``[rows, world_size * topk, 2]`` would cost a
    full reorder copy. The key ranks only the high ``score_bits`` of each score
    (see ``dcp_topk_score_bits``) and the low ``id_bits`` of each id.
    """
    assert gathered.dim() == 4 and gathered.shape[2:] == (topk, 2), (
        f"Expected [world_size, rows, {topk}, 2] candidates, "
        f"got {tuple(gathered.shape)}"
    )
    assert gathered.stride(-1) == 1 and gathered.stride(-2) == 2, (
        f"Candidate pairs must be contiguous, got strides {gathered.stride()}"
    )
    if out is None:
        out = torch.empty(
            (gathered.shape[1], topk),
            dtype=torch.int32,
            device=gathered.device,
        )
    _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL(
        gathered, out, topk=topk, score_bits=score_bits, id_bits=id_bits
    )
    return out