Skip to content

vllm.v1.worker.cp_utils

Functions:

should_skip_dcp_context_attention(context_kv_lens_cpu)

Whether DCP context attention can be skipped for this batch.

Must be computed from rank-invariant inputs only (the global context lengths, NOT this rank's local share from get_dcp_local_seq_lens): the non-skip path in _forward_with_dcp issues DCP collectives (query all-gather + LSE combine), so every DCP rank must take the same branch. A rank can hold zero local context tokens while other ranks still hold context for the same batch.

Source code in vllm/v1/worker/cp_utils.py 
60
61
62
63
64
65
66
67
68
69
70
def should_skip_dcp_context_attention(context_kv_lens_cpu: torch.Tensor) -> bool:
    """Whether DCP context attention can be skipped for this batch.

    Must be computed from rank-invariant inputs only (the global context
    lengths, NOT this rank's local share from get_dcp_local_seq_lens): the
    non-skip path in _forward_with_dcp issues DCP collectives (query
    all-gather + LSE combine), so every DCP rank must take the same branch.
    A rank can hold zero local context tokens while other ranks still hold
    context for the same batch.
    """
    return int(context_kv_lens_cpu.max().item()) == 0

split_dcp_context_queries(query_start_loc, seq_lens_cpu_upper_bound, max_query_len, num_actual_tokens)

Split reordered DCP context queries into decode and extend regions.

Source code in vllm/v1/worker/cp_utils.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
100
101
102
103
104
105
def split_dcp_context_queries(
    query_start_loc: torch.Tensor,
    seq_lens_cpu_upper_bound: torch.Tensor | None,
    max_query_len: int,
    num_actual_tokens: int,
) -> tuple[int, int, int, int]:
    """Split reordered DCP context queries into decode and extend regions."""
    num_reqs = query_start_loc.shape[0] - 1
    if max_query_len <= 1:
        return num_reqs, 0, num_actual_tokens, 0
    if seq_lens_cpu_upper_bound is None:
        return 0, num_reqs, 0, num_actual_tokens

    common_attn_metadata = cast(
        CommonAttentionMetadata,
        SimpleNamespace(
            max_query_len=max_query_len,
            num_reqs=num_reqs,
            num_actual_tokens=num_actual_tokens,
            query_start_loc_cpu=query_start_loc,
            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
            is_prefilling=None,
        ),
    )
    (
        num_decodes,
        num_extends,
        _num_prefills,
        num_decode_tokens,
        num_extend_tokens,
        _num_prefill_tokens,
    ) = split_decodes_prefills_and_extends(common_attn_metadata)
    return num_decodes, num_extends, num_decode_tokens, num_extend_tokens