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

stable_topk_from_gathered_candidates_cutedsl(gathered, topk, out=None)

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.

Source code in vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def stable_topk_from_gathered_candidates_cutedsl(
    gathered: torch.Tensor,
    topk: int,
    out: torch.Tensor | None = None,
) -> 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.
    """
    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)
    return out