| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from types import SimpleNamespace |
| 4 | from typing import TYPE_CHECKING, Any, cast |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.config import VllmConfig, get_layers_from_vllm_config |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 11 | from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends |
| 12 | |
| 13 | if TYPE_CHECKING: |
| 14 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 15 | else: |
| 16 | AttentionLayerBase = object |
| 17 | |
| 18 | logger = init_logger(__name__) |
| 19 | |
| 20 | |
| 21 | def check_attention_cp_compatibility( |
| 22 | vllm_config: VllmConfig, |
| 23 | target_layer_names: set[str] | None = None, |
| 24 | ) -> None: |
| 25 | pcp_size = vllm_config.parallel_config.prefill_context_parallel_size |
| 26 | dcp_size = vllm_config.parallel_config.decode_context_parallel_size |
| 27 | interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size |
| 28 | if pcp_size * dcp_size > 1: |
| 29 | layer_type = cast(type[Any], AttentionLayerBase) |
| 30 | layers = get_layers_from_vllm_config(vllm_config, layer_type) |
| 31 | for layer_name, layer in layers.items(): |
| 32 | check_pcp = target_layer_names is None or layer_name in target_layer_names |
| 33 | get_attn_backend = getattr(layer, "get_attn_backend", None) |
| 34 | if pcp_size > 1 and check_pcp and get_attn_backend is not None: |
| 35 | backend = get_attn_backend() |
| 36 | assert backend.supports_pcp(), ( |
| 37 | "PCP requires attention backend support, " |
| 38 | f"but {backend.get_name()} does not support PCP." |
| 39 | ) |
| 40 | layer_impl = getattr(layer, "impl", None) |
| 41 | if layer_impl is None: |
| 42 | continue |
| 43 | if not check_pcp and layer_impl.dcp_world_size == 1: |
| 44 | continue |
| 45 | if vllm_config.speculative_config is not None and interleave_size > 1: |
| 46 | assert layer_impl.supports_mtp_with_cp_non_trivial_interleave_size, ( |
| 47 | "MTP with cp_kv_cache_interleave_size > 1 is not " |
| 48 | f"supported in {layer_impl.__class__.__name__}." |
| 49 | ) |
| 50 | if dcp_size > 1: |
| 51 | assert layer_impl.need_to_return_lse_for_decode, ( |
| 52 | "Decode Context Parallelism (DCP) requires attention " |
| 53 | "implementations to return the softmax LSE during decode, " |
| 54 | f"but {layer_impl.__class__.__name__} does not. " |
| 55 | "Try a different backend by setting " |
| 56 | "--attention-backend or disable DCP." |
| 57 | ) |
| 58 | |
| 59 | |
| 60 | def get_dcp_dummy_context_len( |
| 61 | dcp_world_size: int, |
| 62 | cp_kv_cache_interleave_size: int, |
| 63 | has_kv_cache_config: bool, |
| 64 | create_mixed_batch: bool, |
| 65 | is_graph_capturing: bool, |
| 66 | uniform_decode: bool, |
| 67 | ) -> int: |
| 68 | if ( |
| 69 | dcp_world_size <= 1 |
| 70 | or not has_kv_cache_config |
| 71 | or not (create_mixed_batch or (is_graph_capturing and uniform_decode)) |
| 72 | ): |
| 73 | return 0 |
| 74 | return dcp_world_size * cp_kv_cache_interleave_size |
| 75 | |
| 76 | |
| 77 | def prepare_dcp_dummy_context_metadata( |
| 78 | *, |
| 79 | input_batch: Any, |
| 80 | kv_cache_config: Any, |
| 81 | query_pos: Any, |
| 82 | positions: torch.Tensor, |
| 83 | query_start_loc: Any, |
| 84 | num_reqs: int, |
| 85 | num_tokens_unpadded: int, |
| 86 | dcp_dummy_context_len: int, |
| 87 | ) -> None: |
| 88 | """Populate valid fake KV metadata for DCP CUDA graph warmup/capture.""" |
| 89 | if dcp_dummy_context_len == 0: |
| 90 | return |
| 91 | |
| 92 | # DCP graph warmup may exercise context attention, so block-table entries |
| 93 | # must point at allocated KV blocks. |
| 94 | assert kv_cache_config is not None |
| 95 | max_valid_block_id = kv_cache_config.num_blocks - 1 |
| 96 | assert max_valid_block_id > 0 |
| 97 | for blk_table in input_batch.block_table.block_tables: |
| 98 | max_row_blocks = ( |
| 99 | blk_table.max_num_blocks_per_req // blk_table.blocks_per_kv_block |
| 100 | ) |
| 101 | block_ids = [ |
| 102 | (block_idx % max_valid_block_id) + 1 for block_idx in range(max_row_blocks) |
| 103 | ] |
| 104 | for req_idx in range(num_reqs): |
| 105 | blk_table.add_row(block_ids, req_idx) |
| 106 | blk_table.commit_block_table(num_reqs) |
| 107 | |
| 108 | query_pos.copy_to_gpu(num_tokens_unpadded) |
| 109 | positions[:num_tokens_unpadded] = ( |
| 110 | query_pos.gpu[:num_tokens_unpadded] + dcp_dummy_context_len |
| 111 | ) |
| 112 | input_batch.block_table.compute_slot_mapping( |
| 113 | num_reqs, |
| 114 | query_start_loc.gpu[: num_reqs + 1], |
| 115 | positions[:num_tokens_unpadded], |
| 116 | ) |
| 117 | |
| 118 | |
| 119 | def should_skip_dcp_context_attention(context_kv_lens_cpu: torch.Tensor) -> bool: |
| 120 | """Whether DCP context attention can be skipped for this batch. |
| 121 | |
| 122 | Must be computed from rank-invariant inputs only (the global context |
| 123 | lengths, NOT this rank's local share from get_dcp_local_seq_lens): the |
| 124 | non-skip path in _forward_with_dcp issues DCP collectives (query |
| 125 | all-gather + LSE combine), so every DCP rank must take the same branch. |
| 126 | A rank can hold zero local context tokens while other ranks still hold |
| 127 | context for the same batch. |
| 128 | """ |
| 129 | return int(context_kv_lens_cpu.max().item()) == 0 |
| 130 | |
| 131 | |
| 132 | def split_dcp_context_queries( |
| 133 | query_start_loc: torch.Tensor, |
| 134 | seq_lens_cpu_upper_bound: torch.Tensor | None, |
| 135 | max_query_len: int, |
| 136 | num_actual_tokens: int, |
| 137 | ) -> tuple[int, int, int, int]: |
| 138 | """Split reordered DCP context queries into decode and extend regions.""" |
| 139 | num_reqs = query_start_loc.shape[0] - 1 |
| 140 | if max_query_len <= 1: |
| 141 | return num_reqs, 0, num_actual_tokens, 0 |
| 142 | if seq_lens_cpu_upper_bound is None: |
| 143 | return 0, num_reqs, 0, num_actual_tokens |
| 144 | |
| 145 | common_attn_metadata = cast( |
| 146 | CommonAttentionMetadata, |
| 147 | SimpleNamespace( |
| 148 | max_query_len=max_query_len, |
| 149 | num_reqs=num_reqs, |
| 150 | num_actual_tokens=num_actual_tokens, |
| 151 | query_start_loc_cpu=query_start_loc, |
| 152 | seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, |
| 153 | is_prefilling=None, |
| 154 | ), |
| 155 | ) |
| 156 | ( |
| 157 | num_decodes, |
| 158 | num_extends, |
| 159 | _num_prefills, |
| 160 | num_decode_tokens, |
| 161 | num_extend_tokens, |
| 162 | _num_prefill_tokens, |
| 163 | ) = split_decodes_prefills_and_extends(common_attn_metadata) |
| 164 | return num_decodes, num_extends, num_decode_tokens, num_extend_tokens |
| 165 | |
| 166 | |
| 167 | def should_split_fa2_dcp_context_attention( |
| 168 | fa_version: int | None, |
| 169 | max_query_len: int, |
| 170 | num_reqs: int, |
| 171 | num_decode_reqs: int, |
| 172 | num_context_prefill_reqs: int, |
| 173 | ) -> bool: |
| 174 | num_prefills = num_reqs - num_decode_reqs |
| 175 | # TODO: Remove this FA2-only DCP compatibility path once FA4 supports |
| 176 | # the Qwen3.5 head_size=256 shape on Blackwell and can be used here. |
| 177 | # FA2 paged-varlen context attention can fail for DCP mixed batches when |
| 178 | # decode rows, context-bearing extend rows, and zero-context pure prefill |
| 179 | # rows are submitted together. |
| 180 | return ( |
| 181 | fa_version == 2 |
| 182 | and max_query_len > 1 |
| 183 | and num_prefills > 0 |
| 184 | and (num_decode_reqs > 0 or num_context_prefill_reqs < num_prefills) |
| 185 | ) |
| 186 | |
| 187 | |
| 188 | def run_split_fa2_dcp_context_attention( |
| 189 | flash_attn_varlen_func: Any, |
| 190 | query_across_dcp: torch.Tensor, |
| 191 | key_cache: torch.Tensor, |
| 192 | value_cache: torch.Tensor, |
| 193 | dcp_context_out: torch.Tensor, |
| 194 | cu_seqlens_q: torch.Tensor, |
| 195 | max_seqlen_q: int, |
| 196 | dcp_context_kv_lens: torch.Tensor, |
| 197 | max_dcp_context_kv_len: int, |
| 198 | softmax_scale: float, |
| 199 | alibi_slopes: torch.Tensor | None, |
| 200 | sliding_window_size: list[int] | None, |
| 201 | block_table: torch.Tensor, |
| 202 | softcap: float, |
| 203 | fa_version: int, |
| 204 | q_descale: torch.Tensor | None, |
| 205 | k_descale: torch.Tensor | None, |
| 206 | v_descale: torch.Tensor | None, |
| 207 | max_num_splits: int, |
| 208 | num_heads: int, |
| 209 | dcp_world_size: int, |
| 210 | num_decode_reqs: int, |
| 211 | num_context_prefill_reqs: int, |
| 212 | num_decode_tokens: int, |
| 213 | num_context_prefill_tokens: int, |
| 214 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 215 | dcp_context_out.zero_() |
| 216 | context_lse = torch.full( |
| 217 | (num_heads * dcp_world_size, query_across_dcp.shape[0]), |
| 218 | -torch.inf, |
| 219 | dtype=torch.float32, |
| 220 | device=query_across_dcp.device, |
| 221 | ) |
| 222 | |
| 223 | if num_decode_tokens > 0: |
| 224 | _, decode_context_lse = flash_attn_varlen_func( |
| 225 | q=query_across_dcp[:num_decode_tokens], |
| 226 | k=key_cache, |
| 227 | v=value_cache, |
| 228 | out=dcp_context_out[:num_decode_tokens], |
| 229 | cu_seqlens_q=cu_seqlens_q[: num_decode_reqs + 1], |
| 230 | max_seqlen_q=1, |
| 231 | seqused_k=dcp_context_kv_lens[:num_decode_reqs], |
| 232 | max_seqlen_k=max_dcp_context_kv_len, |
| 233 | softmax_scale=softmax_scale, |
| 234 | causal=False, |
| 235 | alibi_slopes=alibi_slopes, |
| 236 | window_size=sliding_window_size, |
| 237 | block_table=block_table[:num_decode_reqs], |
| 238 | softcap=softcap, |
| 239 | return_softmax_lse=True, |
| 240 | scheduler_metadata=None, |
| 241 | fa_version=fa_version, |
| 242 | q_descale=q_descale[:num_decode_reqs] if q_descale is not None else None, |
| 243 | k_descale=k_descale[:num_decode_reqs] if k_descale is not None else None, |
| 244 | v_descale=v_descale[:num_decode_reqs] if v_descale is not None else None, |
| 245 | num_splits=max_num_splits, |
| 246 | ) |
| 247 | context_lse[:, :num_decode_tokens] = decode_context_lse |
| 248 | |
| 249 | if num_context_prefill_tokens > 0: |
| 250 | prefill_start = num_decode_tokens |
| 251 | prefill_end = prefill_start + num_context_prefill_tokens |
| 252 | prefill_query_start_loc = ( |
| 253 | cu_seqlens_q[ |
| 254 | num_decode_reqs : num_decode_reqs + num_context_prefill_reqs + 1 |
| 255 | ] |
| 256 | - num_decode_tokens |
| 257 | ) |
| 258 | prefill_req_slice = slice( |
| 259 | num_decode_reqs, num_decode_reqs + num_context_prefill_reqs |
| 260 | ) |
| 261 | _, prefill_context_lse = flash_attn_varlen_func( |
| 262 | q=query_across_dcp[prefill_start:prefill_end], |
| 263 | k=key_cache, |
| 264 | v=value_cache, |
| 265 | out=dcp_context_out[prefill_start:prefill_end], |
| 266 | cu_seqlens_q=prefill_query_start_loc, |
| 267 | max_seqlen_q=max_seqlen_q, |
| 268 | seqused_k=dcp_context_kv_lens[prefill_req_slice], |
| 269 | max_seqlen_k=max_dcp_context_kv_len, |
| 270 | softmax_scale=softmax_scale, |
| 271 | causal=False, |
| 272 | alibi_slopes=alibi_slopes, |
| 273 | window_size=sliding_window_size, |
| 274 | block_table=block_table[prefill_req_slice], |
| 275 | softcap=softcap, |
| 276 | return_softmax_lse=True, |
| 277 | scheduler_metadata=None, |
| 278 | fa_version=fa_version, |
| 279 | q_descale=q_descale[prefill_req_slice] if q_descale is not None else None, |
| 280 | k_descale=k_descale[prefill_req_slice] if k_descale is not None else None, |
| 281 | v_descale=v_descale[prefill_req_slice] if v_descale is not None else None, |
| 282 | num_splits=max_num_splits, |
| 283 | ) |
| 284 | context_lse[:, prefill_start:prefill_end] = prefill_context_lse |
| 285 | |
| 286 | return dcp_context_out, context_lse |