Skip to content

vllm.v1.worker.cp_utils

Functions:

prepare_dcp_dummy_context_metadata(*, input_batch, kv_cache_config, query_pos, positions, query_start_loc, num_reqs, num_tokens_unpadded, dcp_dummy_context_len)

Populate valid fake KV metadata for DCP CUDA graph warmup/capture.

Source code in vllm/v1/worker/cp_utils.py 
 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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def prepare_dcp_dummy_context_metadata(
    *,
    input_batch: Any,
    kv_cache_config: Any,
    query_pos: Any,
    positions: torch.Tensor,
    query_start_loc: Any,
    num_reqs: int,
    num_tokens_unpadded: int,
    dcp_dummy_context_len: int,
) -> None:
    """Populate valid fake KV metadata for DCP CUDA graph warmup/capture."""
    if dcp_dummy_context_len == 0:
        return

    # DCP graph warmup may exercise context attention, so block-table entries
    # must point at allocated KV blocks.
    assert kv_cache_config is not None
    max_valid_block_id = kv_cache_config.num_blocks - 1
    assert max_valid_block_id > 0
    for blk_table in input_batch.block_table.block_tables:
        max_row_blocks = (
            blk_table.max_num_blocks_per_req // blk_table.blocks_per_kv_block
        )
        block_ids = [
            (block_idx % max_valid_block_id) + 1 for block_idx in range(max_row_blocks)
        ]
        for req_idx in range(num_reqs):
            blk_table.add_row(block_ids, req_idx)
        blk_table.commit_block_table(num_reqs)

    query_pos.copy_to_gpu(num_tokens_unpadded)
    positions[:num_tokens_unpadded] = (
        query_pos.gpu[:num_tokens_unpadded] + dcp_dummy_context_len
    )
    input_batch.block_table.compute_slot_mapping(
        num_reqs,
        query_start_loc.gpu[: num_reqs + 1],
        positions[:num_tokens_unpadded],
    )

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 
123
124
125
126
127
128
129
130
131
132
133
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 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
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