From 83694eacee630ec8154b278a00bc984505d282be Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Mon, 20 Jul 2026 17:25:12 +0800 Subject: [PATCH 01/18] rebased Signed-off-by: JaredforReal --- vllm/model_executor/layers/attention/pcp.py | 30 +++++ vllm/v1/attention/backends/flash_attn.py | 121 ++++++++++++++++++-- vllm/v1/worker/gpu/pcp_manager.py | 6 +- 3 files changed, 148 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/layers/attention/pcp.py b/vllm/model_executor/layers/attention/pcp.py index 75ab1c9e8e13..f12b5150106d 100644 --- a/vllm/model_executor/layers/attention/pcp.py +++ b/vllm/model_executor/layers/attention/pcp.py @@ -66,6 +66,36 @@ def maybe_gather_mla_latent_cache_inputs( return cache_kv_c, cache_k_pe, cache_slot_mapping +def maybe_gather_kv_cache_inputs( + key: torch.Tensor, + value: torch.Tensor, + slot_mapping: torch.Tensor | None, + num_decode_tokens: int | None, + use_pcp: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """GQA/MHA PCP KV-cache gather. + + All-gather the prefill portion of K/V across PCP ranks so every rank can + write the full prefill KV to its (replicated, ``dcp=1``) cache, while + keeping decode writes local. Returns contiguous K, V and the gathered + cache slot mapping ready for ``reshape_and_cache_flash``. No-op when PCP + is off. + + Gathering K and V separately (rather than ``cat``-then-``split``) keeps + each output contiguous, so the cache kernel's ``head stride == head_size`` + assumption holds. + """ + if not use_pcp or num_decode_tokens is None: + return key, value, slot_mapping + assert slot_mapping is not None + (cache_key, cache_value), cache_slot_mapping = _gather_prefill_cache_inputs( + (key, value), + slot_mapping, + num_decode_tokens, + ) + return cache_key, cache_value, cache_slot_mapping + + def maybe_gather_indexer_k( k: torch.Tensor, slot_mapping: torch.Tensor, diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 86198dda61c0..1a5cb418065d 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -28,7 +28,7 @@ is_flash_attn_varlen_func_available, ) from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens -from vllm.v1.attention.ops.common import cp_lse_ag_out_rs +from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.worker.workspace import current_workspace_manager @@ -48,6 +48,10 @@ ) from vllm.config.cache import CacheDType from vllm.distributed.parallel_state import get_dcp_group +from vllm.forward_context import ( + get_forward_context, + is_forward_context_available, +) from vllm.logger import init_logger from vllm.platforms.interface import DeviceCapability from vllm.utils.math_utils import cdiv, round_up @@ -351,6 +355,13 @@ def get_cudagraph_support( vllm_config: "VllmConfig", kv_cache_spec: "AttentionSpec", ) -> AttentionCGSupport: + # Under PCP, do_kv_cache_update runs a PCP all-gather to materialize the + # full replicated KV cache before attention. That collective is not + # captureable into a CUDA graph, so force NEVER and let the runner use + # piecewise graphs (rest of the model captured, attention eager) when + # PCP is on. + if vllm_config.parallel_config.prefill_context_parallel_size > 1: + return AttentionCGSupport.NEVER return cls._cudagraph_support def __init__( @@ -388,6 +399,16 @@ def __init__( self.dcp_world_size = 1 self.dcp_rank = 0 + try: + from vllm.distributed.parallel_state import get_pcp_group + + self.pcp_world_size = get_pcp_group().world_size + self.pcp_rank = get_pcp_group().rank_in_group + except AssertionError: + # PCP might not be initialized in testing. + self.pcp_world_size = 1 + self.pcp_rank = 0 + self.cp_kv_cache_interleave_size = ( self.parallel_config.cp_kv_cache_interleave_size ) @@ -731,6 +752,7 @@ def use_cascade_attention(self, *args, **kwargs) -> bool: class FlashAttentionImpl(AttentionImpl): can_return_lse_for_decode: bool = True + supports_pcp: bool = True def __init__( self, @@ -798,7 +820,20 @@ def __init__( and vllm_config.parallel_config.decode_context_parallel_size > 1 and vllm_config.parallel_config.dcp_comm_backend == "a2a" ) - self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs + # When DCP shares the PCP ranks (dcp == pcp), Q is replicated across the + # DCP ranks, so the partial attentions combine with an all-reduce (every + # rank ends with the full output) instead of the gathered-Q + + # reduce-scatter used when DCP reuses the TP ranks. + self.dcp_shares_pcp_ranks = ( + vllm_config is not None + and vllm_config.parallel_config.decode_context_parallel_size > 1 + and vllm_config.parallel_config.decode_context_parallel_size + == vllm_config.parallel_config.prefill_context_parallel_size + ) + if self.dcp_shares_pcp_ranks: + self.dcp_combine = cp_lse_ag_out_ar + else: + self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs self._dcp_dtype: torch.dtype | None = None self._dcp_max_num_tokens: int = 0 @@ -808,6 +843,10 @@ def __init__( vllm_config.scheduler_config.max_num_batched_tokens ) + # self.pcp_world_size / self.pcp_rank are auto-populated by + # AttentionImplBase.__new__ from get_pcp_group(). + self.use_pcp = self.pcp_world_size > 1 + def forward( self, layer: torch.nn.Module, @@ -1068,6 +1107,24 @@ def forward( ) return output + def _get_attn_metadata_for_layer( + self, layer: torch.nn.Module + ) -> FlashAttentionMetadata | None: + """Fetch this layer's FlashAttentionMetadata from the forward context. + + ``do_kv_cache_update`` does not receive ``attn_metadata``, so under PCP + we look it up by ``layer_name`` (the same key the runner uses to store + per-layer metadata in ``forward_context.attn_metadata``). + """ + if self.pcp_world_size <= 1 or not is_forward_context_available(): + return None + attn_metadata_map = get_forward_context().attn_metadata + layer_name = getattr(layer, "layer_name", None) + if not isinstance(attn_metadata_map, dict) or layer_name is None: + return None + meta = attn_metadata_map.get(layer_name) + return meta if isinstance(meta, FlashAttentionMetadata) else None + def do_kv_cache_update( self, layer: torch.nn.Module, @@ -1081,6 +1138,44 @@ def do_kv_cache_update( # we use direct Q, K, V tensors without caching return + if self.use_pcp: + # Under PCP each rank holds only its DualChunkSwap prefill chunks in + # the rank-local batch. All-gather the prefill K/V across PCP ranks + # (keeping decode writes local) so every rank's cache receives the + # full prefill KV (replicated when dcp=1; the DCP-local shard when + # pcp+dcp). maybe_gather_kv_cache_inputs gathers K/V separately -> + # contiguous (cache kernel head-stride contract), builds the gathered + # slot mapping, and handles the decode/prefill split + empty-prefill. + from vllm.model_executor.layers.attention.pcp import ( + maybe_gather_kv_cache_inputs, + ) + + attn_metadata = self._get_attn_metadata_for_layer(layer) + num_decode_tokens = ( + attn_metadata.num_decode_tokens if attn_metadata is not None else 0 + ) + key_cache, value_cache = kv_cache.transpose(1, 2).split( + self.head_size, dim=-1 + ) + cache_key, cache_value, cache_slot_mapping = maybe_gather_kv_cache_inputs( + key, + value, + slot_mapping, + num_decode_tokens, + self.use_pcp, + ) + reshape_and_cache_flash( + cache_key, + cache_value, + key_cache, + value_cache, + cache_slot_mapping, + self.kv_cache_dtype, + layer._k_scale, + layer._v_scale, + ) + return + # Scatter write into the KV cache using slot_mapping indices. # No TMA kernel is invoked here, so stride canonicalization is not needed. # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D)) @@ -1152,11 +1247,19 @@ def _forward_with_dcp( ) return output - query_across_dcp = get_dcp_group().all_gather(query, dim=1) + # When DCP shares the PCP ranks (dcp == pcp), Q is replicated across the + # DCP ranks, so attend directly with local heads and combine via + # all-reduce. Otherwise gather Q across heads and reduce-scatter. + if self.dcp_shares_pcp_ranks: + query_for_context = query + context_num_heads = self.num_heads + else: + query_for_context = get_dcp_group().all_gather(query, dim=1) + context_num_heads = self.num_heads * self.dcp_world_size sliding_window_size = ( list(self.sliding_window) if self.sliding_window is not None else None ) - n = query_across_dcp.shape[0] + n = query_for_context.shape[0] num_reqs = cu_seqlens_q.shape[0] - 1 num_decodes = attn_metadata.num_decode_reqs num_context_prefills = attn_metadata.num_prefill_reqs @@ -1173,7 +1276,7 @@ def _forward_with_dcp( dcp_context_out_spec = ( ( dcp_context_out_tokens, - self.num_heads * self.dcp_world_size, + context_num_heads, self.head_size, ), self._dcp_dtype, @@ -1186,12 +1289,16 @@ def _forward_with_dcp( if split_dcp_context: # TODO: Remove this DCP + FA2 mixed decode/prefill workaround once # FA4 supports this Qwen3.5 shape. + assert not self.dcp_shares_pcp_ranks, ( + "FA2 split-DCP context path does not support dcp_shares_pcp_ranks" + " (pcp+dcp); use FA3/FA4." + ) assert attn_metadata.dcp_context_kv_lens is not None assert attn_metadata.max_dcp_context_kv_len is not None assert self.vllm_flash_attn_version is not None context_attn_out, context_lse = run_split_fa2_dcp_context_attention( flash_attn_varlen_func, - query_across_dcp, + query_for_context, key_cache, value_cache, dcp_context_out, @@ -1218,7 +1325,7 @@ def _forward_with_dcp( ) else: context_attn_out, context_lse = flash_attn_varlen_func( - q=query_across_dcp, + q=query_for_context, k=key_cache, v=value_cache, out=dcp_context_out, diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index f50e48749421..2f157d9c234a 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -132,8 +132,10 @@ def validate_config( if pcp_size <= 1: return - if not model_config.use_mla: - raise NotImplementedError("MRV2 PCP currently supports MLA models only.") + # Non-MLA (GQA/MHA) models are supported on the FlashAttention backend, + # which opts in via AttentionImplBase.supports_pcp. MLA uses its own + # latent-attention PCP path. Per-backend capability is still enforced by + # check_attention_cp_compatibility in vllm/v1/worker/cp_utils.py. if parallel_config.pipeline_parallel_size > 1: raise NotImplementedError("MRV2 PCP does not support PP yet.") if model_config.is_encoder_decoder: From 6ce9ec74ccf2be086f537c2de8f299ccba7a4571 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Mon, 20 Jul 2026 18:16:05 +0800 Subject: [PATCH 02/18] fix assertion Signed-off-by: JaredforReal --- vllm/config/model.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/vllm/config/model.py b/vllm/config/model.py index 6b032ae76210..00a780979231 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1247,7 +1247,22 @@ def verify_with_parallel_config( ) decode_context_parallel_size = parallel_config.decode_context_parallel_size - if decode_context_parallel_size > 1 and not self.use_mla: + # When DCP spans the PCP axis (dcp == pcp, pcp > 1) it does NOT shard + # the TP-local KV heads: Q is replicated across PCP ranks and the + # partial attentions combine via all-reduce (see FlashAttention + # dcp_shares_pcp_ranks). So the TP-head DCP constraints below apply + # only when DCP reuses the TP ranks (pcp == 1) or spans the full + # TP x PCP axis (dcp == tp*pcp). + dcp_shares_pcp_ranks = ( + parallel_config.prefill_context_parallel_size > 1 + and decode_context_parallel_size + == parallel_config.prefill_context_parallel_size + ) + if ( + decode_context_parallel_size > 1 + and not self.use_mla + and not dcp_shares_pcp_ranks + ): total_num_kv_heads = self.get_total_num_kv_heads() assert tensor_parallel_size > total_num_kv_heads, ( f"tensor parallel size {tensor_parallel_size} must be greater " From 7c78f462b34fbfec6faedaa62f18ba494f2afbe2 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Mon, 20 Jul 2026 18:23:05 +0800 Subject: [PATCH 03/18] enable MRv2 PCP Signed-off-by: JaredforReal --- vllm/config/vllm.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 4b4a97f41b6b..0a26ad235770 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2130,10 +2130,11 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: model_config = self.model_config speculative_config = self.speculative_config - if self.parallel_config.prefill_context_parallel_size > 1 and not ( - model_config is not None and model_config.use_mla - ): - unsupported.append("prefill context parallelism") + # PCP on MRv2 is opt-in per attention backend via + # AttentionImplBase.supports_pcp, enforced by + # check_attention_cp_compatibility (vllm/v1/worker/cp_utils.py). MLA and + # GQA (FlashAttention) both opt in, so do not block here -- if no + # attention impl supports PCP the cp_utils check raises a clear error. if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE: unsupported.append("stock torch.compile") From dde82d75f390f8d19c94bd5df303e5797911ada8 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Tue, 21 Jul 2026 06:26:35 +0000 Subject: [PATCH 04/18] sotre Signed-off-by: JaredforReal --- vllm/v1/attention/backends/flash_attn.py | 261 ++++++++++++++++++- vllm/v1/core/kv_cache_coordinator.py | 9 +- vllm/v1/core/kv_cache_utils.py | 8 +- vllm/v1/core/single_type_kv_cache_manager.py | 8 +- vllm/v1/kv_cache_interface.py | 13 +- vllm/v1/worker/gpu/model_runner.py | 18 +- 6 files changed, 292 insertions(+), 25 deletions(-) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 1a5cb418065d..fd925dfe4c44 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -413,6 +413,17 @@ def __init__( self.parallel_config.cp_kv_cache_interleave_size ) + # DCP shares the PCP ranks (dcp == pcp): the MRv2 PCP+DCP path + # (_forward_dcp_mrv2) is used instead of _forward_with_dcp. Q is + # replicated across these ranks (no DCP head gather), so the + # scheduler_metadata for the prefill path must use num_heads_q without + # the DCP multiply. + self.dcp_shares_pcp_ranks = ( + self.dcp_world_size > 1 + and self.parallel_config.decode_context_parallel_size + == self.parallel_config.prefill_context_parallel_size + ) + self.use_full_cuda_graph = ( self.compilation_config.cudagraph_mode.has_full_cudagraphs() ) @@ -523,19 +534,24 @@ def build( max_num_splits = 1 def schedule( - batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal + batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal, + num_heads_q=None, ): cache_dtype = self.cache_config.cache_dtype if is_quantized_kv_cache(cache_dtype): qkv_dtype = current_platform.fp8_dtype() else: qkv_dtype = self.kv_cache_dtype + if num_heads_q is None: + # DCP that reuses the TP ranks gathers Q across DCP heads before + # the kernel, so the scheduler sees num_heads_q * dcp heads. + num_heads_q = self.num_heads_q * self.dcp_world_size if aot_schedule: return get_scheduler_metadata( batch_size=batch_size, max_seqlen_q=max_query_len, max_seqlen_k=max_seq_len, - num_heads_q=self.num_heads_q * self.dcp_world_size, + num_heads_q=num_heads_q, num_heads_kv=self.num_heads_kv, headdim=self.headdim, cache_seqlens=seqlens, @@ -563,7 +579,13 @@ def schedule( suffix_kv_lens = None prefix_scheduler_metadata = None - if self.dcp_world_size > 1: + if self.dcp_world_size > 1 and not self.dcp_shares_pcp_ranks: + # Pure-DCP path (DCP reuses the TP ranks; no PCP). The MRv2 PCP+DCP + # case (dcp_shares_pcp_ranks) keeps a REPLICATED cache + # (model_runner.pcp_dcp_replicated_cache) and runs the normal FA + # path, so it must skip this DCP block entirely -- otherwise + # split_dcp_context_queries sets a DCP-split num_decode_tokens that + # corrupts do_kv_cache_update's PCP gather and garbles the cache. query_lens = query_start_loc[1:] - query_start_loc[:-1] context_kv_lens = seq_lens - query_lens local_context_kv_lens = get_dcp_local_seq_lens( @@ -616,14 +638,29 @@ def schedule( (max_seq_len + num_partitions - 1) // num_partitions ) * self.cp_kv_cache_interleave_size - scheduler_metadata = schedule( - batch_size=num_reqs, - cu_query_lens=query_start_loc, - max_query_len=max_query_len, - seqlens=dcp_context_kv_lens, - max_seq_len=max_dcp_context_kv_len, - causal=False, - ) + if self.dcp_shares_pcp_ranks: + # MRv2 PCP+DCP: prefill attends the full replicated cache + # with standard FA (_forward_dcp_mrv2); decode does not use + # scheduler_metadata. Build the standard scheduler_metadata + # (full seq_lens; Q is replicated so no DCP head multiply). + scheduler_metadata = schedule( + batch_size=num_reqs, + cu_query_lens=query_start_loc, + max_query_len=max_query_len, + seqlens=seq_lens, + max_seq_len=max_seq_len, + causal=causal, + num_heads_q=self.num_heads_q, + ) + else: + scheduler_metadata = schedule( + batch_size=num_reqs, + cu_query_lens=query_start_loc, + max_query_len=max_query_len, + seqlens=dcp_context_kv_lens, + max_seq_len=max_dcp_context_kv_len, + causal=False, + ) elif use_cascade: cu_prefix_query_lens = torch.tensor( [0, num_actual_tokens], dtype=torch.int32, device=self.device @@ -956,7 +993,9 @@ def forward( k_descale = layer._k_scale.expand(descale_shape) v_descale = layer._v_scale.expand(descale_shape) - if self.dcp_world_size > 1: + if self.dcp_world_size > 1 and not self.dcp_shares_pcp_ranks: + # Pure DCP (DCP reuses the TP ranks; no PCP). MRv1-era sharded + # path -- the cache is DCP-sharded and writes are rank-local. self._forward_with_dcp( query[:num_actual_tokens], key[:num_actual_tokens], @@ -971,6 +1010,13 @@ def forward( ) return output else: + # dcp_world_size <= 1, or dcp_shares_pcp_ranks (MRv2 PCP+DCP). + # With the REPLICATED cache (model_runner.pcp_dcp_replicated_cache) + # the PCP+DCP case is equivalent to pure PCP, so the normal FA + # path is correct for both decode and prefill. The sharded-cache + # path (DCP decode-KV memory win: decode LSE-combine + prefill KV + # all-gather) lives in _forward_dcp_mrv2 and is reserved for a + # future opt-in (Phase 2b). window = ( attn_metadata.sliding_window if attn_metadata.sliding_window is not None @@ -1199,6 +1245,197 @@ def do_kv_cache_update( layer._v_scale, ) + def _forward_dcp_mrv2( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: FlashAttentionMetadata, + q_descale: torch.Tensor | None = None, + k_descale: torch.Tensor | None = None, + v_descale: torch.Tensor | None = None, + ) -> torch.Tensor: + """MRv2 PCP+DCP attention for GQA (``dcp_shares_pcp_ranks``). + + Mirrors the MLA split (``mla_attention.py::MLAAttentionBase. + forward_impl``) so this composes correctly with a PCP-partitioned + (DualChunkSwap) batch. ``_forward_with_dcp`` is the MRv1-era DCP path + and is intentionally not used here: its ``max_dcp_context_kv_len == 0`` + guard is rank-varying under PCP (rank 0's chunk starts at position 0 + with no context) and desyncs the NCCL collective. + + - Decode tokens: Q is replicated across the DCP ranks, so each rank + attends its local DCP KV shard and partial outputs combine via LSE + all-gather + all-reduce (``cp_lse_ag_out_ar``). No rank-varying + guard -- every rank takes the collective branch. + - Prefill/extend tokens: the prefill KV is replicated on every rank + (PCP all-gather in ``do_kv_cache_update``), so each rank's + DualChunkSwap chunk attends the full cache with standard FA -- no + collective, no LSE merge (different Q per rank is fine). + """ + assert self.vllm_flash_attn_version is not None, ( + "FlashAttention version not detected." + ) + + num_decode_tokens = attn_metadata.num_decode_tokens + num_prefill_tokens = attn_metadata.num_prefill_tokens + + if num_decode_tokens > 0 and num_prefill_tokens > 0: + raise NotImplementedError( + "MRV2 PCP+DCP mixed prefill+decode batch is not supported yet " + "(Phase 1.5). Run prefill and decode in separate steps.") + + cu_seqlens_q = attn_metadata.query_start_loc + max_seqlen_q = attn_metadata.max_query_len + block_table = attn_metadata.block_table + sliding_window_size = ( + list(self.sliding_window) if self.sliding_window is not None else None + ) + query = query.contiguous() + + if num_decode_tokens > 0: + self._forward_dcp_mrv2_decode( + query, key, value, key_cache, value_cache, output, + attn_metadata, cu_seqlens_q, max_seqlen_q, block_table, + sliding_window_size, q_descale, k_descale, v_descale, + ) + return output + + # Pure prefill/extend: standard FA against the full replicated cache. + # TODO(phase 2): if the prefill cache is DCP-sharded rather than + # replicated under pcp+dcp, switch to MLA-style KV all-gather. + flash_attn_varlen_func( + q=query, + k=key_cache, + v=value_cache, + out=output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + seqused_k=attn_metadata.seq_lens, + max_seqlen_k=attn_metadata.max_seq_len, + softmax_scale=self.scale, + causal=attn_metadata.causal, + alibi_slopes=self.alibi_slopes, + window_size=sliding_window_size, + block_table=block_table, + softcap=self.logits_soft_cap, + scheduler_metadata=attn_metadata.scheduler_metadata, + fa_version=self.vllm_flash_attn_version, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + num_splits=attn_metadata.max_num_splits, + ) + return output + + def _forward_dcp_mrv2_decode( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: FlashAttentionMetadata, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + block_table: torch.Tensor, + sliding_window_size: list[int] | None, + q_descale: torch.Tensor | None, + k_descale: torch.Tensor | None, + v_descale: torch.Tensor | None, + ) -> None: + """Pure-decode DCP combine for MRv2 PCP+DCP (mirrors MLA decode). + + Each rank attends its (replicated) decode Q against its local DCP KV + shard (cached prior), combines the partial outputs across DCP ranks via + ``cp_lse_ag_out_ar`` (LSE all-gather + all-reduce), then merges with the + per-token (new K/V) attention. Unlike ``_forward_with_dcp`` there is no + ``max_dcp_context_kv_len == 0`` early return: every rank runs the same + collective so the NCCL communicator stays synchronized. + """ + # dcp_shares_pcp_ranks => Q is replicated; attend local heads directly. + query_for_context = query + context_num_heads = self.num_heads + n = query_for_context.shape[0] + + dcp_context_out_tokens = max(n, self._dcp_max_num_tokens) + dcp_context_out_spec = ( + (dcp_context_out_tokens, context_num_heads, self.head_size), + self._dcp_dtype, + ) + (dcp_context_out_workspace,) = current_workspace_manager().get_simultaneous( + dcp_context_out_spec, + ) + dcp_context_out = dcp_context_out_workspace[:n] + + assert attn_metadata.dcp_context_kv_lens is not None + assert attn_metadata.max_dcp_context_kv_len is not None + context_attn_out, context_lse = flash_attn_varlen_func( + q=query_for_context, + k=key_cache, + v=value_cache, + out=dcp_context_out, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + seqused_k=attn_metadata.dcp_context_kv_lens, + max_seqlen_k=attn_metadata.max_dcp_context_kv_len, + softmax_scale=self.scale, + causal=False, + alibi_slopes=self.alibi_slopes, + window_size=sliding_window_size, + block_table=block_table, + softcap=self.logits_soft_cap, + return_softmax_lse=True, + fa_version=self.vllm_flash_attn_version, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + num_splits=attn_metadata.max_num_splits, + ) + # FA returns LSE as [H, B]; the DCP combine wants [B, H]. + context_attn_out_cor, context_lse_cor = self.dcp_combine( + context_attn_out, + context_lse.transpose(0, 1), + get_dcp_group(), + return_lse=True, + ) + context_lse_cor = context_lse_cor.transpose(0, 1).contiguous() + + query_attn_out, query_lse = flash_attn_varlen_func( + q=query, + k=key, + v=value, + out=output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + cu_seqlens_k=cu_seqlens_q, + max_seqlen_k=max_seqlen_q, + softmax_scale=self.scale, + causal=attn_metadata.causal, + alibi_slopes=self.alibi_slopes, + window_size=sliding_window_size, + softcap=self.logits_soft_cap, + return_softmax_lse=True, + fa_version=self.vllm_flash_attn_version, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + num_splits=attn_metadata.max_num_splits, + ) + assert context_attn_out_cor.shape == query_attn_out.shape + assert context_lse_cor.shape == query_lse.shape + merge_attn_states( + output, + context_attn_out_cor, + context_lse_cor, + query_attn_out, + query_lse, + ) + def _forward_with_dcp( self, query: torch.Tensor, diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index df8769c3f3a3..ca6060a72e71 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -470,8 +470,13 @@ def __init__( self.block_size = self.kv_cache_spec.block_size self.dcp_world_size = dcp_world_size self.pcp_world_size = pcp_world_size - if dcp_world_size > 1: - self.block_size *= dcp_world_size + # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache, no block inflation. + cache_dcp = ( + 1 if (dcp_world_size > 1 and dcp_world_size == pcp_world_size) + else dcp_world_size + ) + if cache_dcp > 1: + self.block_size *= cache_dcp # For models using only Mamba, block_size is set to max_model_len when # prefix caching is disabled, and hash_block_size validation is skipped. assert not enable_caching or (hash_block_size == self.block_size), ( diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index d4970d91b2d8..7feb29b76026 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -644,14 +644,18 @@ def resolve_kv_cache_block_sizes( """ cache_config = vllm_config.cache_config dcp = vllm_config.parallel_config.decode_context_parallel_size + pcp = vllm_config.parallel_config.prefill_context_parallel_size + # MRv2 PCP+DCP (dcp == pcp > 1): the KV cache is REPLICATED, so the + # scheduler/hash block size is NOT inflated by dcp (pure-DCP keeps it). + cache_dcp = 1 if (dcp > 1 and dcp == pcp) else dcp groups = kv_cache_config.kv_cache_groups if len(groups) <= 1: - bs = cache_config.block_size * dcp + bs = cache_config.block_size * cache_dcp return bs, bs group_block_sizes = [ - g.kv_cache_spec.block_size * dcp + g.kv_cache_spec.block_size * cache_dcp if isinstance(g.kv_cache_spec, AttentionSpec) else g.kv_cache_spec.block_size for g in groups diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index f8578c68a389..77c10a50da8e 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -75,8 +75,12 @@ def __init__( self.block_size = kv_cache_spec.block_size self.dcp_world_size = dcp_world_size self.pcp_world_size = pcp_world_size - if dcp_world_size > 1: - self.block_size *= dcp_world_size + # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache, so do not inflate the + # block size. Pure-DCP (dcp != pcp) keeps the sharded-KV inflation. + replicated = dcp_world_size > 1 and dcp_world_size == pcp_world_size + cache_dcp = 1 if replicated else dcp_world_size + if cache_dcp > 1: + self.block_size *= cache_dcp self.kv_cache_spec = kv_cache_spec self.block_pool = block_pool self.enable_caching = enable_caching diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 7a44d4e1d25e..f6f610e2377a 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -219,7 +219,10 @@ def real_page_size_bytes(self) -> int: def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: parallel_config = vllm_config.parallel_config - kv_shard_count = parallel_config.decode_context_parallel_size + dcp = parallel_config.decode_context_parallel_size + # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache, no sharding divisor. + pcp = parallel_config.prefill_context_parallel_size + kv_shard_count = 1 if (dcp > 1 and dcp == pcp) else dcp return cdiv(max_len, self.block_size * kv_shard_count) @@ -257,9 +260,11 @@ def __post_init__(self): def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_model_len = vllm_config.model_config.max_model_len - dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size - if dcp_world_size > 1: - max_model_len = cdiv(max_model_len, dcp_world_size) + dcp = vllm_config.parallel_config.decode_context_parallel_size + pcp = vllm_config.parallel_config.prefill_context_parallel_size + # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache, no 1/dcp saving. + if dcp > 1 and dcp != pcp: + max_model_len = cdiv(max_model_len, dcp) return cdiv(max_model_len, self.block_size) * self.page_size_bytes @classmethod diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 293b06ff97b9..7b35c23c31d5 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -177,6 +177,18 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0 self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size + # MRv2 PCP+DCP (dcp reuses the pcp ranks, i.e. dcp == pcp > 1): the KV + # cache is kept REPLICATED -- equivalent to pure PCP -- so the prefill + # PCP all-gather write and the FA reads share one non-sharded layout. + # This is correct but makes the dcp dimension cache-redundant (no + # decode-KV memory win); the sharded-cache path is Phase 2b. + self.pcp_size = self.parallel_config.prefill_context_parallel_size + self.pcp_dcp_replicated_cache = ( + self.use_dcp and self.pcp_size == self.dcp_size + ) + self.cache_dcp_size = 1 if self.pcp_dcp_replicated_cache else self.dcp_size + self.cache_dcp_rank = 0 if self.pcp_dcp_replicated_cache else self.dcp_rank + # Multimodal self.mm_registry = MULTIMODAL_REGISTRY self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( @@ -431,7 +443,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: # As a result, one block on the current rank covers `block_size * cp_size` # tokens in the full, global (unsharded) sequence. max_num_blocks = cdiv( - block_table_max_model_len, spec.block_size * self.dcp_size + block_table_max_model_len, spec.block_size * self.cache_dcp_size ) # Align to a multiple of (128 / block_size) as required by some attention # backends such as TRTLLM (#39324) @@ -455,8 +467,8 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: max_num_blocks_per_group=max_num_blocks_per_group, device=self.device, kernel_block_sizes=self.kernel_block_sizes, - cp_size=self.dcp_size, - cp_rank=self.dcp_rank, + cp_size=self.cache_dcp_size, + cp_rank=self.cache_dcp_rank, cp_interleave=self.cp_interleave, ) self.pcp_manager = pcp.maybe_build_pcp_manager( From 0fb6ece3d15c6aa7b80f40fdf5ec591491ef8c08 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Tue, 21 Jul 2026 09:30:15 +0000 Subject: [PATCH 05/18] update Signed-off-by: JaredforReal --- vllm/envs.py | 7 +++++++ vllm/v1/core/kv_cache_coordinator.py | 10 ++++++++-- vllm/v1/core/kv_cache_utils.py | 11 ++++++++--- vllm/v1/core/sched/scheduler.py | 2 +- vllm/v1/core/single_type_kv_cache_manager.py | 12 +++++++++--- vllm/v1/kv_cache_interface.py | 13 +++++++++---- vllm/v1/worker/gpu/model_runner.py | 4 +++- 7 files changed, 45 insertions(+), 14 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index 67d8fa5a2c65..f87d533a3ea5 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -272,6 +272,10 @@ VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024 VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary" VLLM_USE_V2_MODEL_RUNNER: bool | None = None + # MRv2 PCP+DCP: when set, use a DCP-SHARDED KV cache (each rank stores + # 1/dcp of positions) instead of the default replicated cache. Gives the + # decode-KV memory win at the cost of a sharded attention/write path. + VLLM_PCP_DCP_SHARDED_KV_CACHE: bool = False VLLM_LOG_MODEL_INSPECTION: bool = False VLLM_DEBUG_MFU_METRICS: bool = False VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False @@ -1922,6 +1926,9 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_USE_V2_MODEL_RUNNER": lambda: maybe_convert_bool( os.getenv("VLLM_USE_V2_MODEL_RUNNER", None) ), + "VLLM_PCP_DCP_SHARDED_KV_CACHE": lambda: maybe_convert_bool( + os.getenv("VLLM_PCP_DCP_SHARDED_KV_CACHE", "0") + ), # Log model inspection after loading. # If enabled, logs a transformers-style hierarchical view of the model # with quantization methods and attention backends. diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index ca6060a72e71..8281e388e462 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -470,9 +470,15 @@ def __init__( self.block_size = self.kv_cache_spec.block_size self.dcp_world_size = dcp_world_size self.pcp_world_size = pcp_world_size - # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache, no block inflation. + # MRv2 PCP+DCP (dcp == pcp > 1): replicated (default) -> no inflation; + # with VLLM_PCP_DCP_SHARDED_KV_CACHE the cache shards -> keep inflation. cache_dcp = ( - 1 if (dcp_world_size > 1 and dcp_world_size == pcp_world_size) + 1 + if ( + dcp_world_size > 1 + and dcp_world_size == pcp_world_size + and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE + ) else dcp_world_size ) if cache_dcp > 1: diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 7feb29b76026..b96d37a8ba9c 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -645,9 +645,14 @@ def resolve_kv_cache_block_sizes( cache_config = vllm_config.cache_config dcp = vllm_config.parallel_config.decode_context_parallel_size pcp = vllm_config.parallel_config.prefill_context_parallel_size - # MRv2 PCP+DCP (dcp == pcp > 1): the KV cache is REPLICATED, so the - # scheduler/hash block size is NOT inflated by dcp (pure-DCP keeps it). - cache_dcp = 1 if (dcp > 1 and dcp == pcp) else dcp + # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache (default) so the + # scheduler/hash block size is NOT inflated by dcp; with + # VLLM_PCP_DCP_SHARDED_KV_CACHE the cache shards and keeps the inflation. + cache_dcp = ( + 1 + if (dcp > 1 and dcp == pcp and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE) + else dcp + ) groups = kv_cache_config.kv_cache_groups if len(groups) <= 1: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 46a999ef7c50..de076e230e47 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -271,7 +271,7 @@ def __init__( log_stats=self.log_stats, enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=self.dcp_world_size, - pcp_world_size=1, + pcp_world_size=self.pcp_world_size, scheduler_block_size=self.block_size, hash_block_size=hash_block_size, metrics_collector=self.kv_metrics_collector, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 77c10a50da8e..675fd1288b89 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -6,6 +6,7 @@ from collections.abc import Sequence from typing import ClassVar +from vllm import envs from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import ( @@ -75,9 +76,14 @@ def __init__( self.block_size = kv_cache_spec.block_size self.dcp_world_size = dcp_world_size self.pcp_world_size = pcp_world_size - # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache, so do not inflate the - # block size. Pure-DCP (dcp != pcp) keeps the sharded-KV inflation. - replicated = dcp_world_size > 1 and dcp_world_size == pcp_world_size + # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache (default) so do not + # inflate; with VLLM_PCP_DCP_SHARDED_KV_CACHE the cache is sharded and + # keeps the inflation. Pure-DCP (dcp != pcp) always shards. + replicated = ( + dcp_world_size > 1 + and dcp_world_size == pcp_world_size + and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE + ) cache_dcp = 1 if replicated else dcp_world_size if cache_dcp > 1: self.block_size *= cache_dcp diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index f6f610e2377a..c1ede2029006 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -13,6 +13,7 @@ import torch from typing_extensions import Self +from vllm import envs from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_up from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim @@ -220,9 +221,11 @@ def real_page_size_bytes(self) -> int: def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: parallel_config = vllm_config.parallel_config dcp = parallel_config.decode_context_parallel_size - # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache, no sharding divisor. + # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache (default) -> no + # sharding divisor; with VLLM_PCP_DCP_SHARDED_KV_CACHE it shards. pcp = parallel_config.prefill_context_parallel_size - kv_shard_count = 1 if (dcp > 1 and dcp == pcp) else dcp + replicated = dcp > 1 and dcp == pcp and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE + kv_shard_count = 1 if replicated else dcp return cdiv(max_len, self.block_size * kv_shard_count) @@ -262,8 +265,10 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_model_len = vllm_config.model_config.max_model_len dcp = vllm_config.parallel_config.decode_context_parallel_size pcp = vllm_config.parallel_config.prefill_context_parallel_size - # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache, no 1/dcp saving. - if dcp > 1 and dcp != pcp: + # MRv2 PCP+DCP (dcp == pcp > 1): replicated (default) -> no 1/dcp saving; + # with VLLM_PCP_DCP_SHARDED_KV_CACHE it shards -> apply the 1/dcp saving. + replicated = dcp > 1 and dcp == pcp and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE + if dcp > 1 and not replicated: max_model_len = cdiv(max_model_len, dcp) return cdiv(max_model_len, self.block_size) * self.page_size_bytes diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 7b35c23c31d5..93ea9b2eb33b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -184,7 +184,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # decode-KV memory win); the sharded-cache path is Phase 2b. self.pcp_size = self.parallel_config.prefill_context_parallel_size self.pcp_dcp_replicated_cache = ( - self.use_dcp and self.pcp_size == self.dcp_size + self.use_dcp + and self.pcp_size == self.dcp_size + and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE ) self.cache_dcp_size = 1 if self.pcp_dcp_replicated_cache else self.dcp_size self.cache_dcp_rank = 0 if self.pcp_dcp_replicated_cache else self.dcp_rank From 94015a98a0d355e206a817b8e0acdf7d0f3924b7 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Tue, 21 Jul 2026 10:46:28 +0000 Subject: [PATCH 06/18] store Signed-off-by: JaredforReal --- vllm/v1/attention/backends/flash_attn.py | 157 ++++++++++++++++------- 1 file changed, 109 insertions(+), 48 deletions(-) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index fd925dfe4c44..fd637254e65f 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -534,7 +534,12 @@ def build( max_num_splits = 1 def schedule( - batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal, + batch_size, + cu_query_lens, + max_query_len, + seqlens, + max_seq_len, + causal, num_heads_q=None, ): cache_dtype = self.cache_config.cache_dtype @@ -579,7 +584,42 @@ def schedule( suffix_kv_lens = None prefix_scheduler_metadata = None - if self.dcp_world_size > 1 and not self.dcp_shares_pcp_ranks: + if ( + self.dcp_shares_pcp_ranks + and self.dcp_world_size > 1 + and envs.VLLM_PCP_DCP_SHARDED_KV_CACHE + ): + # MRv2 PCP+DCP with a SHARDED cache (VLLM_PCP_DCP_SHARDED_KV_CACHE). + # DualChunkSwap has already partitioned the batch per rank; DCP only + # shards the cache. Both decode and prefill use the context+query+merge + # path (_forward_dcp_mrv2_decode): the context attention reads each + # rank's local shard of the causal prefix (the full prefix is in the + # sharded cache -- do_kv_cache_update all-gathered+wrote it before + # forward) and LSE-combines; the query attention reads this rank's own + # new K/V. context_kv_lens = seq_lens - query_lens is the absolute + # position of the chunk start (= the prior prefix length), correct for + # both decode (query_len=1) and prefill. scheduler_metadata is unused. + scheduler_metadata = None + query_lens = query_start_loc[1:] - query_start_loc[:-1] + context_kv_lens = seq_lens - query_lens + local_context_kv_lens = get_dcp_local_seq_lens( + context_kv_lens, + self.dcp_world_size, + self.dcp_rank, + self.cp_kv_cache_interleave_size, + ) + self._dcp_context_kv_lens[:num_reqs] = local_context_kv_lens + self._dcp_context_kv_lens[num_reqs:] = 0 + dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs] + num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size + max_dcp_context_kv_len = ( + (max_seq_len + num_partitions - 1) // num_partitions + ) * self.cp_kv_cache_interleave_size + if max_query_len <= 1: + num_decode_tokens = num_actual_tokens + else: + num_prefill_tokens = num_actual_tokens + elif self.dcp_world_size > 1 and not self.dcp_shares_pcp_ranks: # Pure-DCP path (DCP reuses the TP ranks; no PCP). The MRv2 PCP+DCP # case (dcp_shares_pcp_ranks) keeps a REPLICATED cache # (model_runner.pcp_dcp_replicated_cache) and runs the normal FA @@ -1009,6 +1049,23 @@ def forward( v_descale=v_descale, ) return output + elif self.dcp_shares_pcp_ranks and envs.VLLM_PCP_DCP_SHARDED_KV_CACHE: + # MRv2 PCP+DCP with a SHARDED cache (VLLM_PCP_DCP_SHARDED_KV_CACHE): + # decode KV is sharded 1/dcp for the memory win. Decode attends + # its local shard + LSE-combine; prefill all-gathers the full KV. + self._forward_dcp_mrv2( + query[:num_actual_tokens], + key[:num_actual_tokens], + value[:num_actual_tokens], + key_cache, + value_cache, + output[:num_actual_tokens], + attn_metadata, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + return output else: # dcp_world_size <= 1, or dcp_shares_pcp_ranks (MRv2 PCP+DCP). # With the REPLICATED cache (model_runner.pcp_dcp_replicated_cache) @@ -1258,23 +1315,27 @@ def _forward_dcp_mrv2( k_descale: torch.Tensor | None = None, v_descale: torch.Tensor | None = None, ) -> torch.Tensor: - """MRv2 PCP+DCP attention for GQA (``dcp_shares_pcp_ranks``). - - Mirrors the MLA split (``mla_attention.py::MLAAttentionBase. - forward_impl``) so this composes correctly with a PCP-partitioned - (DualChunkSwap) batch. ``_forward_with_dcp`` is the MRv1-era DCP path - and is intentionally not used here: its ``max_dcp_context_kv_len == 0`` - guard is rank-varying under PCP (rank 0's chunk starts at position 0 - with no context) and desyncs the NCCL collective. - - - Decode tokens: Q is replicated across the DCP ranks, so each rank - attends its local DCP KV shard and partial outputs combine via LSE - all-gather + all-reduce (``cp_lse_ag_out_ar``). No rank-varying - guard -- every rank takes the collective branch. - - Prefill/extend tokens: the prefill KV is replicated on every rank - (PCP all-gather in ``do_kv_cache_update``), so each rank's - DualChunkSwap chunk attends the full cache with standard FA -- no - collective, no LSE merge (different Q per rank is fine). + """MRv2 PCP+DCP attention for GQA with a SHARDED cache. + + Both decode and prefill/extend use the same context+query+merge path + (``_forward_dcp_mrv2_decode``). This composes correctly with a + PCP-partitioned (DualChunkSwap) batch because the full causal prefix is + in the sharded cache by forward time (``do_kv_cache_update`` all-gathers + every rank's new K/V across PCP and writes it before forward runs): + + - Context attention: each rank attends its local DCP KV shard of the + prefix (``seqused_k = dcp_context_kv_lens``), non-causal, and the + partial outputs combine via LSE all-gather + all-reduce + (``cp_lse_ag_out_ar``). + - Query attention: each rank attends its own new K/V (the DualChunkSwap + chunk) with within-chunk causal. + - ``merge_attn_states`` fuses the two. + + ``_forward_with_dcp`` is the MRv1-era DCP path and is intentionally not + used: its ``max_dcp_context_kv_len == 0`` guard is rank-varying under + PCP and desyncs the NCCL collective. This path has no such guard -- + every rank takes the collective branch (chunk-0's empty context yields + an all-PAD ``seqused_k`` whose LSE combine contributes nothing). """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." @@ -1286,7 +1347,8 @@ def _forward_dcp_mrv2( if num_decode_tokens > 0 and num_prefill_tokens > 0: raise NotImplementedError( "MRV2 PCP+DCP mixed prefill+decode batch is not supported yet " - "(Phase 1.5). Run prefill and decode in separate steps.") + "(Phase 1.5). Run prefill and decode in separate steps." + ) cu_seqlens_q = attn_metadata.query_start_loc max_seqlen_q = attn_metadata.max_query_len @@ -1298,38 +1360,37 @@ def _forward_dcp_mrv2( if num_decode_tokens > 0: self._forward_dcp_mrv2_decode( - query, key, value, key_cache, value_cache, output, - attn_metadata, cu_seqlens_q, max_seqlen_q, block_table, - sliding_window_size, q_descale, k_descale, v_descale, + query, + key, + value, + key_cache, + value_cache, + output, + attn_metadata, + cu_seqlens_q, + max_seqlen_q, + block_table, + sliding_window_size, + q_descale, + k_descale, + v_descale, ) return output - # Pure prefill/extend: standard FA against the full replicated cache. - # TODO(phase 2): if the prefill cache is DCP-sharded rather than - # replicated under pcp+dcp, switch to MLA-style KV all-gather. - flash_attn_varlen_func( - q=query, - k=key_cache, - v=value_cache, - out=output, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - seqused_k=attn_metadata.seq_lens, - max_seqlen_k=attn_metadata.max_seq_len, - softmax_scale=self.scale, - causal=attn_metadata.causal, - alibi_slopes=self.alibi_slopes, - window_size=sliding_window_size, - block_table=block_table, - softcap=self.logits_soft_cap, - scheduler_metadata=attn_metadata.scheduler_metadata, - fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - num_splits=attn_metadata.max_num_splits, + # Prefill/extend under a SHARDED cache: DualChunkSwap PARTITIONS the Q + # (each rank holds a different chunk-Q), so the decode-style LSE-combine + # does NOT apply (it requires a replicated Q). Each rank's chunk-Q must + # attend the FULL causal KV, which lives across all DCP shards -- so the + # prefill must all-gather the full KV (new tokens from do_kv_cache_update + # + prior from the sharded cache for extend) and attend with the proven + # block_table mechanism against a workspace. NOT YET IMPLEMENTED. + raise NotImplementedError( + "MRV2 PCP+DCP SHARDED prefill is not implemented yet: DualChunkSwap " + "partitions Q so the LSE-combine can't replace a full-KV gather. " + "Prefill must all-gather the full KV and attend each chunk-Q against " + "it (block_table into a workspace). Use VLLM_PCP_DCP_SHARDED_KV_CACHE=0 " + "(Option A, replicated cache) until this lands." ) - return output def _forward_dcp_mrv2_decode( self, From 4cd1e73a2c6b79c031a3f7bc9b1e83f3585e831e Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Thu, 23 Jul 2026 07:15:59 +0000 Subject: [PATCH 07/18] fix pcp+dcp accuracy Signed-off-by: JaredforReal --- vllm/model_executor/layers/attention/pcp.py | 40 +++ vllm/v1/attention/backends/flash_attn.py | 338 ++++++++++++++------ vllm/v1/worker/gpu/model_runner.py | 2 + vllm/v1/worker/gpu/pcp_manager.py | 158 ++++++++- 4 files changed, 441 insertions(+), 97 deletions(-) diff --git a/vllm/model_executor/layers/attention/pcp.py b/vllm/model_executor/layers/attention/pcp.py index f12b5150106d..3b65cde96e86 100644 --- a/vllm/model_executor/layers/attention/pcp.py +++ b/vllm/model_executor/layers/attention/pcp.py @@ -110,6 +110,46 @@ def maybe_gather_indexer_k( return cache_k, cache_slot_mapping +def gather_prefill_qkv_global( + pcp_prefill_gather: tuple[torch.Tensor, torch.Tensor, torch.Tensor, int], + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """All-gather this rank's DualChunkSwap chunk Q/K/V across PCP and reorder + to global position order (real tokens only). + + ``pcp_prefill_gather`` is the (restore_idx, gather_idx, global_cu_seqlens, + padded_n) tuple from ``PCPManager.prefill_gather_indices()``. + """ + restore_idx, _, _, padded_n = pcp_prefill_gather + pcp_group = get_pcp_group() + + def _pad(t: torch.Tensor) -> torch.Tensor: + if t.shape[0] < padded_n: + pad = t.new_zeros((padded_n - t.shape[0],) + tuple(t.shape[1:])) + return torch.cat([t, pad], dim=0) + return t + + q_g = pcp_group.all_gather(_pad(query), dim=0)[restore_idx] + k_g = pcp_group.all_gather(_pad(key), dim=0)[restore_idx] + v_g = pcp_group.all_gather(_pad(value), dim=0)[restore_idx] + return q_g, k_g, v_g + + +def slice_prefill_output_local( + pcp_prefill_gather: tuple[torch.Tensor, torch.Tensor, torch.Tensor, int], + out_g: torch.Tensor, + num_actual: int, +) -> torch.Tensor: + """Slice a global-order output tensor back to this PCP rank's local chunk.""" + _, gather_idx, _, padded_n = pcp_prefill_gather + pcp_rank = get_pcp_group().rank_in_group + out_gathered = out_g[gather_idx] + local = out_gathered[pcp_rank * padded_n : (pcp_rank + 1) * padded_n] + return local[:num_actual] + + def finalize_mla_pcp_decode( output: torch.Tensor, num_heads: int, diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index fd637254e65f..eb19677b1681 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -10,6 +10,10 @@ import torch from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention.pcp import ( + gather_prefill_qkv_global, + slice_prefill_output_local, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import ( canonicalize_singleton_dim_strides, @@ -261,6 +265,10 @@ class FlashAttentionMetadata: max_dcp_context_kv_len: int | None = None dcp_context_kv_lens: torch.Tensor | None = None + # Per-segment is_prefilling flag (rank-local view). Used by the MRv2 + # sharded PCP+DCP path to extract decode tokens out of a mixed batch. + is_prefilling: torch.Tensor | None = None + # Split counts for FA2 DCP context attention. num_prefill_* tracks # context-bearing extend rows; pure prefills do not attend to DCP context. num_decode_reqs: int = 0 @@ -615,10 +623,18 @@ def schedule( max_dcp_context_kv_len = ( (max_seq_len + num_partitions - 1) // num_partitions ) * self.cp_kv_cache_interleave_size - if max_query_len <= 1: - num_decode_tokens = num_actual_tokens + # Detect prefill vs decode. Under DualChunkSwap a prefill chunk can be + # as small as 1 token, so max_query_len is unreliable -- use the + # per-segment is_prefilling flag instead. + is_prefilling = common_attn_metadata.is_prefilling + if is_prefilling is not None: + is_prefill_batch = bool(is_prefilling.any().item()) else: + is_prefill_batch = max_query_len > 1 + if is_prefill_batch: num_prefill_tokens = num_actual_tokens + else: + num_decode_tokens = num_actual_tokens elif self.dcp_world_size > 1 and not self.dcp_shares_pcp_ranks: # Pure-DCP path (DCP reuses the TP ranks; no PCP). The MRv2 PCP+DCP # case (dcp_shares_pcp_ranks) keeps a REPLICATED cache @@ -766,6 +782,7 @@ def schedule( slot_mapping=slot_mapping, max_dcp_context_kv_len=max_dcp_context_kv_len, dcp_context_kv_lens=dcp_context_kv_lens, + is_prefilling=common_attn_metadata.is_prefilling, num_decode_reqs=num_decode_reqs, num_prefill_reqs=num_prefill_reqs, num_decode_tokens=num_decode_tokens, @@ -1257,6 +1274,15 @@ def do_kv_cache_update( num_decode_tokens = ( attn_metadata.num_decode_tokens if attn_metadata is not None else 0 ) + # Rank-invariant gather decision: the per-rank num_decode_tokens can + # differ under DualChunkSwap (some ranks get zero prefill chunks), + # which would desync the PCP all-gather. If the global batch has any + # prefill, every rank treats the whole batch as prefill (gather all); + # only a pure-decode batch skips the gather. Used by both the sharded + # and replicated PCP+DCP paths. + gflags = get_forward_context().additional_kwargs.get("pcp_global_flags") + if gflags is not None and gflags[0]: + num_decode_tokens = 0 key_cache, value_cache = kv_cache.transpose(1, 2).split( self.head_size, dim=-1 ) @@ -1302,6 +1328,24 @@ def do_kv_cache_update( layer._v_scale, ) + def _fa_common_kwargs( + self, + q_descale: torch.Tensor | None, + k_descale: torch.Tensor | None, + v_descale: torch.Tensor | None, + ) -> dict: + sw = list(self.sliding_window) if self.sliding_window is not None else None + return dict( + softmax_scale=self.scale, + alibi_slopes=self.alibi_slopes, + window_size=sw, + softcap=self.logits_soft_cap, + fa_version=self.vllm_flash_attn_version, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + def _forward_dcp_mrv2( self, query: torch.Tensor, @@ -1315,50 +1359,34 @@ def _forward_dcp_mrv2( k_descale: torch.Tensor | None = None, v_descale: torch.Tensor | None = None, ) -> torch.Tensor: - """MRv2 PCP+DCP attention for GQA with a SHARDED cache. - - Both decode and prefill/extend use the same context+query+merge path - (``_forward_dcp_mrv2_decode``). This composes correctly with a - PCP-partitioned (DualChunkSwap) batch because the full causal prefix is - in the sharded cache by forward time (``do_kv_cache_update`` all-gathers - every rank's new K/V across PCP and writes it before forward runs): - - - Context attention: each rank attends its local DCP KV shard of the - prefix (``seqused_k = dcp_context_kv_lens``), non-causal, and the - partial outputs combine via LSE all-gather + all-reduce - (``cp_lse_ag_out_ar``). - - Query attention: each rank attends its own new K/V (the DualChunkSwap - chunk) with within-chunk causal. - - ``merge_attn_states`` fuses the two. - - ``_forward_with_dcp`` is the MRv1-era DCP path and is intentionally not - used: its ``max_dcp_context_kv_len == 0`` guard is rank-varying under - PCP and desyncs the NCCL collective. This path has no such guard -- - every rank takes the collective branch (chunk-0's empty context yields - an all-PAD ``seqused_k`` whose LSE combine contributes nothing). + """MRv2 PCP+DCP GQA attention over a SHARDED cache. + + Pure-decode batches use ``_forward_dcp_mrv2_decode`` (Q replicated, + attend the sharded prefix + LSE-combine). Prefill/extend/mixed batches + all-gather the DualChunkSwap-partitioned Q/K/V across PCP, run full + causal on the gathered tokens, optionally context-attend the cached + prefix (extend), and slice back; decode tokens in a mixed batch are then + overwritten via the decode path. """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." ) - - num_decode_tokens = attn_metadata.num_decode_tokens - num_prefill_tokens = attn_metadata.num_prefill_tokens - - if num_decode_tokens > 0 and num_prefill_tokens > 0: - raise NotImplementedError( - "MRV2 PCP+DCP mixed prefill+decode batch is not supported yet " - "(Phase 1.5). Run prefill and decode in separate steps." - ) + fc = get_forward_context() + # Rank-invariant (has_prefill, has_decode) from the global batch; per-rank + # num_decode_tokens is not safe to route on (DualChunkSwap can zero some + # ranks' prefill chunks and desync the collectives). + global_has_prefill, global_has_decode = fc.additional_kwargs.get( + "pcp_global_flags", (True, False) + ) cu_seqlens_q = attn_metadata.query_start_loc max_seqlen_q = attn_metadata.max_query_len block_table = attn_metadata.block_table - sliding_window_size = ( - list(self.sliding_window) if self.sliding_window is not None else None - ) query = query.contiguous() + fa_kw = self._fa_common_kwargs(q_descale, k_descale, v_descale) - if num_decode_tokens > 0: + # Pure-decode batch: Q is replicated, attend the sharded cache + LSE-combine. + if not global_has_prefill: self._forward_dcp_mrv2_decode( query, key, @@ -1370,27 +1398,166 @@ def _forward_dcp_mrv2( cu_seqlens_q, max_seqlen_q, block_table, - sliding_window_size, - q_descale, - k_descale, - v_descale, + fa_kw, + ) + return output + + # Prefill/extend/mixed: DualChunkSwap partitions the prefill Q, so each + # rank all-gathers Q/K/V across PCP, runs full causal on the gathered + # tokens, and slices its chunk back out. Decode tokens in a mixed batch + # self-attend here (no cached context) and are overwritten below. + info = fc.additional_kwargs.get("pcp_prefill_gather") + if info is None: + # Warmup / dummy run: no PCP partition -> no gather indices. Fall + # back to local causal attention (shape-correct is all warmup needs). + flash_attn_varlen_func( + q=query, + k=key, + v=value, + out=output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + cu_seqlens_k=cu_seqlens_q, + max_seqlen_k=max_seqlen_q, + causal=attn_metadata.causal, + num_splits=attn_metadata.max_num_splits, + **fa_kw, ) return output + global_cu_seqlens = info[2] + num_actual = query.shape[0] + q_g, k_g, v_g = gather_prefill_qkv_global(info, query, key, value) + cu_seqlens = global_cu_seqlens.to(device=query.device, dtype=torch.int32) + global_max_seqlen = int(global_cu_seqlens.max().item()) + out_g = torch.empty_like(q_g) + # Query attention: full causal over the gathered new tokens. + _, query_lse_g = flash_attn_varlen_func( + q=q_g, + k=k_g, + v=v_g, + out=out_g, + cu_seqlens_q=cu_seqlens, + max_seqlen_q=global_max_seqlen, + cu_seqlens_k=cu_seqlens, + max_seqlen_k=global_max_seqlen, + causal=attn_metadata.causal, + return_softmax_lse=True, + **fa_kw, + ) + # EXTEND (prefix caching): q_g (replicated) also attends the cached + # prefix -- each rank its DCP shard, LSE-combined. Skipped for pure + # first-prefills (rank-invariant has_extend = gctx[3]). + final_g = out_g + gctx = fc.additional_kwargs.get("pcp_global_ctx") + if gctx is not None and gctx[3] and key_cache is not None: + g_block_table, g_ctx_kv_lens, g_max_ctx_kv, _ = gctx + ctx_out_g, ctx_lse_g = flash_attn_varlen_func( + q=q_g, + k=key_cache, + v=value_cache, + cu_seqlens_q=cu_seqlens, + max_seqlen_q=global_max_seqlen, + seqused_k=g_ctx_kv_lens, + max_seqlen_k=g_max_ctx_kv, + causal=False, + block_table=g_block_table, + return_softmax_lse=True, + **fa_kw, + ) + ctx_out_cor, ctx_lse_cor = self.dcp_combine( + ctx_out_g, + ctx_lse_g.transpose(0, 1), + get_dcp_group(), + return_lse=True, + ) + ctx_lse_cor = ctx_lse_cor.transpose(0, 1).contiguous() + final_g = torch.empty_like(out_g) + merge_attn_states( + final_g, + ctx_out_cor, + ctx_lse_cor, + out_g, + query_lse_g, + ) + # global -> gathered (rank-order) -> this rank's local chunk. + output[:num_actual] = slice_prefill_output_local(info, final_g, num_actual) + + # MIXED batch: the prefill gather above gave decode tokens only + # self-attention (1-token segments, no cached context). Overwrite them + # with the correct decode computation. No-op for a pure-prefill batch. + if global_has_decode: + self._overwrite_mixed_decode_tokens( + query=query, + key=key, + value=value, + key_cache=key_cache, + value_cache=value_cache, + output=output, + attn_metadata=attn_metadata, + fa_kw=fa_kw, + ) + return output - # Prefill/extend under a SHARDED cache: DualChunkSwap PARTITIONS the Q - # (each rank holds a different chunk-Q), so the decode-style LSE-combine - # does NOT apply (it requires a replicated Q). Each rank's chunk-Q must - # attend the FULL causal KV, which lives across all DCP shards -- so the - # prefill must all-gather the full KV (new tokens from do_kv_cache_update - # + prior from the sharded cache for extend) and attend with the proven - # block_table mechanism against a workspace. NOT YET IMPLEMENTED. - raise NotImplementedError( - "MRV2 PCP+DCP SHARDED prefill is not implemented yet: DualChunkSwap " - "partitions Q so the LSE-combine can't replace a full-KV gather. " - "Prefill must all-gather the full KV and attend each chunk-Q against " - "it (block_table into a workspace). Use VLLM_PCP_DCP_SHARDED_KV_CACHE=0 " - "(Option A, replicated cache) until this lands." + def _overwrite_mixed_decode_tokens( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: FlashAttentionMetadata, + fa_kw: dict, + ) -> None: + """Recompute decode-token outputs the prefill-gather path self-attended. + + Decode tokens are replicated across PCP ranks; extract them, run the + decode path (sharded-context + LSE-combine) on the subset, and scatter + back. Zero-token segments (DualChunkSwap's dummy for empty ranks) are + excluded -- counting them would desync the DCP LSE-combine. + """ + is_pre = attn_metadata.is_prefilling + if is_pre is None: + return + qsl = attn_metadata.query_start_loc + is_pre = is_pre.to(device=qsl.device) + seg_query_lens = qsl[1:] - qsl[:-1] + decode_seg_mask = (~is_pre.bool()) & (seg_query_lens > 0) + num_decode = int(decode_seg_mask.sum().item()) + if num_decode == 0: + return + + token_mask = torch.repeat_interleave(decode_seg_mask, seg_query_lens) + token_mask = token_mask[: query.shape[0]] + dec_q = query[token_mask].contiguous() + dec_k = key[token_mask].contiguous() + dec_v = value[token_mask].contiguous() + + dec_query_lens = seg_query_lens[decode_seg_mask] + dec_cu_seqlens = torch.zeros(num_decode + 1, dtype=qsl.dtype, device=qsl.device) + torch.cumsum(dec_query_lens, dim=0, out=dec_cu_seqlens[1:]) + + ctx_kv_lens = attn_metadata.dcp_context_kv_lens + assert ctx_kv_lens is not None + dec_ctx_kv_lens = ctx_kv_lens[decode_seg_mask] + + dec_out = output[token_mask] + self._forward_dcp_mrv2_decode( + dec_q, + dec_k, + dec_v, + key_cache, + value_cache, + dec_out, + attn_metadata, + dec_cu_seqlens, + int(dec_query_lens.max().item()), + attn_metadata.block_table[decode_seg_mask], + fa_kw, + dcp_context_kv_lens=dec_ctx_kv_lens, + max_dcp_context_kv_len=int(dec_ctx_kv_lens.max().item()), ) + output[token_mask] = dec_out def _forward_dcp_mrv2_decode( self, @@ -1404,58 +1571,44 @@ def _forward_dcp_mrv2_decode( cu_seqlens_q: torch.Tensor, max_seqlen_q: int, block_table: torch.Tensor, - sliding_window_size: list[int] | None, - q_descale: torch.Tensor | None, - k_descale: torch.Tensor | None, - v_descale: torch.Tensor | None, + fa_kw: dict, + dcp_context_kv_lens: torch.Tensor | None = None, + max_dcp_context_kv_len: int | None = None, ) -> None: - """Pure-decode DCP combine for MRv2 PCP+DCP (mirrors MLA decode). - - Each rank attends its (replicated) decode Q against its local DCP KV - shard (cached prior), combines the partial outputs across DCP ranks via - ``cp_lse_ag_out_ar`` (LSE all-gather + all-reduce), then merges with the - per-token (new K/V) attention. Unlike ``_forward_with_dcp`` there is no - ``max_dcp_context_kv_len == 0`` early return: every rank runs the same - collective so the NCCL communicator stays synchronized. + """Decode (or decode-subset) over a sharded cache: attend the local DCP + shard of the cached prefix + LSE-combine, merge with the per-token new + K/V attention. Every rank runs the combine (no rank-varying guard) so + NCCL stays synchronized. Context lens default to the metadata but can be + overridden for a decode-only subset of a mixed batch. """ - # dcp_shares_pcp_ranks => Q is replicated; attend local heads directly. - query_for_context = query - context_num_heads = self.num_heads - n = query_for_context.shape[0] - - dcp_context_out_tokens = max(n, self._dcp_max_num_tokens) + if dcp_context_kv_lens is None: + dcp_context_kv_lens = attn_metadata.dcp_context_kv_lens + if max_dcp_context_kv_len is None: + max_dcp_context_kv_len = attn_metadata.max_dcp_context_kv_len + n = query.shape[0] dcp_context_out_spec = ( - (dcp_context_out_tokens, context_num_heads, self.head_size), + (max(n, self._dcp_max_num_tokens), self.num_heads, self.head_size), self._dcp_dtype, ) (dcp_context_out_workspace,) = current_workspace_manager().get_simultaneous( dcp_context_out_spec, ) - dcp_context_out = dcp_context_out_workspace[:n] - - assert attn_metadata.dcp_context_kv_lens is not None - assert attn_metadata.max_dcp_context_kv_len is not None + assert dcp_context_kv_lens is not None + assert max_dcp_context_kv_len is not None context_attn_out, context_lse = flash_attn_varlen_func( - q=query_for_context, + q=query, k=key_cache, v=value_cache, - out=dcp_context_out, + out=dcp_context_out_workspace[:n], cu_seqlens_q=cu_seqlens_q, max_seqlen_q=max_seqlen_q, - seqused_k=attn_metadata.dcp_context_kv_lens, - max_seqlen_k=attn_metadata.max_dcp_context_kv_len, - softmax_scale=self.scale, + seqused_k=dcp_context_kv_lens, + max_seqlen_k=max_dcp_context_kv_len, causal=False, - alibi_slopes=self.alibi_slopes, - window_size=sliding_window_size, block_table=block_table, - softcap=self.logits_soft_cap, return_softmax_lse=True, - fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, num_splits=attn_metadata.max_num_splits, + **fa_kw, ) # FA returns LSE as [H, B]; the DCP combine wants [B, H]. context_attn_out_cor, context_lse_cor = self.dcp_combine( @@ -1475,17 +1628,10 @@ def _forward_dcp_mrv2_decode( max_seqlen_q=max_seqlen_q, cu_seqlens_k=cu_seqlens_q, max_seqlen_k=max_seqlen_q, - softmax_scale=self.scale, causal=attn_metadata.causal, - alibi_slopes=self.alibi_slopes, - window_size=sliding_window_size, - softcap=self.logits_soft_cap, return_softmax_lse=True, - fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, num_splits=attn_metadata.max_num_splits, + **fa_kw, ) assert context_attn_out_cor.shape == query_attn_out.shape assert context_lse_cor.shape == query_lse.shape diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 93ea9b2eb33b..3ffc28b1c4b6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1361,6 +1361,8 @@ def execute_model( skip_compiled=skip_compiled, is_padding=input_batch.is_padding, ): + if self.pcp_manager is not None: + self.pcp_manager.populate_forward_context() self.kv_connector.pre_forward(scheduler_output) if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE: # Run the PIECEWISE graph (compiled PW cudagraph or breakable diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 2f157d9c234a..ac5df1ba373b 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -7,8 +7,9 @@ from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group +from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens @@ -69,8 +70,44 @@ def __init__( self._hidden_restore_idx: torch.Tensor | None = None self._padded_gather_idx: torch.Tensor | None = None self._gathered_kv_write_mask: torch.Tensor | None = None + # GLOBAL batch composition (rank-invariant: every PCP rank sees the same + # global batch, so these are identical across ranks). Used for + # rank-consistent mixed-batch detection in the sharded attention path -- + # per-rank is_prefilling can differ (DualChunkSwap leaves some ranks with + # zero prefill chunks), which would desync NCCL collectives. + self._global_has_prefill: bool = False + self._global_has_decode: bool = False self._pad_slot_id = torch.tensor(PAD_SLOT_ID, dtype=torch.int64, device=device) + # Global-context attention metadata for the sharded PCP+DCP prefill path + # (extend/prefix-caching support). Built once per step in + # build_global_context_metadata(); consumed by _forward_dcp_mrv2 to run + # context attention on the replicated gathered q_g against this rank's + # DCP shard of the cached prefix. + self._global_ctx_block_tables: tuple[torch.Tensor, ...] | None = ( + tuple( + table.new_zeros((max_num_reqs, table.shape[1])) + for table in block_tables.input_block_tables + ) + if block_tables is not None and max_num_reqs is not None + else None + ) + self._global_ctx_block_table_ptrs: torch.Tensor | None = ( + torch.tensor( + [t.data_ptr() for t in self._global_ctx_block_tables], + dtype=torch.uint64, + device=device, + ) + if self._global_ctx_block_tables is not None + else None + ) + self._global_ctx_kv_lens: torch.Tensor | None = ( + torch.zeros(max_num_reqs, dtype=torch.int32, device=device) + if max_num_reqs is not None + else None + ) + self._global_num_reqs: int = 0 + max_num_local_reqs = 2 * max_num_reqs if max_num_reqs is not None else None self._input_buffers = ( InputBuffers(max_num_local_reqs, max_num_tokens, device) @@ -332,6 +369,9 @@ def partition_batch(self, input_batch: InputBatch) -> InputBatch: num_scheduled_tokens = global_batch.num_scheduled_tokens num_computed_tokens = global_batch.num_computed_tokens_np is_prefilling = global_batch.is_prefilling_np + # Rank-invariant global composition (see __init__ comment). + self._global_has_prefill = bool(is_prefilling.any()) + self._global_has_decode = bool((~is_prefilling).any()) segments_by_rank, per_rank_num_tokens = self._build_batch_layout( num_scheduled_tokens, @@ -612,6 +652,122 @@ def restore_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: gathered = get_pcp_group().all_gather(hidden_states, dim=0) return gathered[self._hidden_restore_idx] + def global_batch_flags(self) -> tuple[bool, bool]: + """Return (has_prefill, has_decode) for the global (pre-partition) batch. + + Rank-invariant across PCP ranks (identical global batch). Used for + rank-consistent mixed-batch detection in the sharded attention path. + """ + return self._global_has_prefill, self._global_has_decode + + def populate_forward_context(self) -> None: + """Stash all PCP metadata the attention forward / kv-cache update need. + + Centralized here so the model runner stays a single call. Populates: + - ``pcp_prefill_gather``: gather/restore indices for the sharded + prefill-KV all-gather attention (None when no prefill ran). + - ``pcp_global_flags``: rank-invariant (has_prefill, has_decode) for + consistent mixed-batch routing. + - ``pcp_global_ctx``: global block table + per-rank DCP shard context + lengths for extend (prefix-caching) support in the prefill path. + """ + kwargs = get_forward_context().additional_kwargs + prefill_gather = self.prefill_gather_indices() + if prefill_gather is not None: + kwargs["pcp_prefill_gather"] = prefill_gather + kwargs["pcp_global_flags"] = self.global_batch_flags() + global_ctx = self.build_global_context_metadata() + if global_ctx is not None: + kwargs["pcp_global_ctx"] = global_ctx + + def build_global_context_metadata( + self, + ) -> tuple[torch.Tensor, torch.Tensor, int, bool] | None: + """Build per-global-request context-attention metadata for the extend + (prefix-caching) case in the sharded PCP+DCP prefill path. + + Returns (global_block_table, global_dcp_context_kv_lens, + max_global_dcp_context_kv_len, has_extend) for the CURRENT step, aligned + with the pre-partition global batch (so the gathered, replicated q_g can + attend this rank's DCP shard of each request's cached prefix). Returns + None when no prefill was partitioned this step. + + ``global_dcp_context_kv_lens`` is this rank's shard count of the prefix + (rank-specific), indexed by global request -- the LSE-combine then + merges the shards for the shared (replicated) q_g. ``has_extend`` is + rank-invariant (True iff some global request has a cached prefix) so the + prefill path can skip the context collective for pure first-prefills. + """ + gb = self._global_batch + if ( + gb is None + or self._global_ctx_block_tables is None + or self._global_ctx_kv_lens is None + or self._block_tables is None + ): + return None + block_tables = self._block_tables + global_ctx_kv_lens = self._global_ctx_kv_lens + num_reqs = gb.num_reqs + self._global_num_reqs = num_reqs + # Per-global-request FULL prefix length, then this rank's shard count. + qsl = gb.query_start_loc + query_lens = qsl[1 : num_reqs + 1] - qsl[:num_reqs] + context_kv_lens = gb.seq_lens[:num_reqs] - query_lens + has_extend = num_reqs > 0 and bool(context_kv_lens.max().item() > 0) + dcp_ctx_kv_lens = get_dcp_local_seq_lens( + context_kv_lens, + self.dcp_world_size, + self.dcp_rank, + self.cp_interleave, + ) + global_ctx_kv_lens[:num_reqs] = dcp_ctx_kv_lens + # Gather this rank's block-table rows for the global requests. + global_bt = block_tables.gather_block_tables( + gb.idx_mapping, + num_reqs, + out=self._global_ctx_block_tables, + out_ptrs=self._global_ctx_block_table_ptrs, + ) + num_partitions = self.dcp_world_size * self.cp_interleave + max_ctx_kv = ( + ((int(context_kv_lens.max().item()) + num_partitions - 1) // num_partitions) + * self.cp_interleave + if num_reqs > 0 + else 0 + ) + return global_bt[0], global_ctx_kv_lens[:num_reqs], max_ctx_kv, has_extend + + def prefill_gather_indices( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int] | None: + """Indices for the SHARDED-cache prefill KV-gather attention. + + Returns (restore_idx, gather_idx, global_cu_seqlens, padded_num_tokens) + for the current prefill step, or None for decode / when PCP did not + partition a prefill. ``restore_idx`` maps a global-position index to its + slot in the PCP all-gathered tensor (``global = gathered[restore_idx]``); + ``gather_idx`` is the inverse (``gathered = global[gather_idx]``); + ``global_cu_seqlens`` is the pre-partition (global) batch's per-request + cumulative query lengths. + """ + if self._hidden_restore_idx is None or self._global_batch is None: + return None + if self._padded_gather_idx is None: + return None + global_cu_seqlens = self._global_batch.query_start_loc + if not isinstance(global_cu_seqlens, torch.Tensor): + return None + # _padded_gather_idx has length padded_num_tokens * pcp_world_size for + # the current step (set in _build_batch_layout). + padded_num_tokens = self._padded_gather_idx.shape[0] // self.pcp_world_size + return ( + self._hidden_restore_idx, + self._padded_gather_idx, + global_cu_seqlens, + padded_num_tokens, + ) + def restore_for_sampling( self, hidden_states: torch.Tensor, From a8aaf8bb20cbb51f215a1046084fc062bedf9087 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Thu, 23 Jul 2026 09:44:35 +0000 Subject: [PATCH 08/18] clean replicated kv for dcp Signed-off-by: JaredforReal --- vllm/envs.py | 7 --- vllm/v1/attention/backends/flash_attn.py | 45 ++++++++------------ vllm/v1/core/kv_cache_coordinator.py | 15 +------ vllm/v1/core/kv_cache_utils.py | 13 +----- vllm/v1/core/single_type_kv_cache_manager.py | 14 +----- vllm/v1/kv_cache_interface.py | 18 ++------ vllm/v1/worker/gpu/model_runner.py | 20 ++------- 7 files changed, 30 insertions(+), 102 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index f87d533a3ea5..67d8fa5a2c65 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -272,10 +272,6 @@ VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024 VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary" VLLM_USE_V2_MODEL_RUNNER: bool | None = None - # MRv2 PCP+DCP: when set, use a DCP-SHARDED KV cache (each rank stores - # 1/dcp of positions) instead of the default replicated cache. Gives the - # decode-KV memory win at the cost of a sharded attention/write path. - VLLM_PCP_DCP_SHARDED_KV_CACHE: bool = False VLLM_LOG_MODEL_INSPECTION: bool = False VLLM_DEBUG_MFU_METRICS: bool = False VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False @@ -1926,9 +1922,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_USE_V2_MODEL_RUNNER": lambda: maybe_convert_bool( os.getenv("VLLM_USE_V2_MODEL_RUNNER", None) ), - "VLLM_PCP_DCP_SHARDED_KV_CACHE": lambda: maybe_convert_bool( - os.getenv("VLLM_PCP_DCP_SHARDED_KV_CACHE", "0") - ), # Log model inspection after loading. # If enabled, logs a transformers-style hierarchical view of the model # with quantization methods and attention backends. diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index eb19677b1681..36d849f36785 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -592,18 +592,15 @@ def schedule( suffix_kv_lens = None prefix_scheduler_metadata = None - if ( - self.dcp_shares_pcp_ranks - and self.dcp_world_size > 1 - and envs.VLLM_PCP_DCP_SHARDED_KV_CACHE - ): - # MRv2 PCP+DCP with a SHARDED cache (VLLM_PCP_DCP_SHARDED_KV_CACHE). - # DualChunkSwap has already partitioned the batch per rank; DCP only - # shards the cache. Both decode and prefill use the context+query+merge - # path (_forward_dcp_mrv2_decode): the context attention reads each - # rank's local shard of the causal prefix (the full prefix is in the - # sharded cache -- do_kv_cache_update all-gathered+wrote it before - # forward) and LSE-combines; the query attention reads this rank's own + if self.dcp_shares_pcp_ranks and self.dcp_world_size > 1: + # MRv2 PCP+DCP (dcp_shares_pcp_ranks): the cache is DCP-sharded + # (1/dcp KV per rank). DualChunkSwap has already partitioned the + # batch per rank; DCP only shards the cache. Both decode and prefill + # use the context+query+merge path (_forward_dcp_mrv2_decode): the + # context attention reads each rank's local shard of the causal + # prefix (the full prefix is in the sharded cache -- + # do_kv_cache_update all-gathered+wrote it before forward) and + # LSE-combines; the query attention reads this rank's own # new K/V. context_kv_lens = seq_lens - query_lens is the absolute # position of the chunk start (= the prior prefix length), correct for # both decode (query_len=1) and prefill. scheduler_metadata is unused. @@ -637,11 +634,8 @@ def schedule( num_decode_tokens = num_actual_tokens elif self.dcp_world_size > 1 and not self.dcp_shares_pcp_ranks: # Pure-DCP path (DCP reuses the TP ranks; no PCP). The MRv2 PCP+DCP - # case (dcp_shares_pcp_ranks) keeps a REPLICATED cache - # (model_runner.pcp_dcp_replicated_cache) and runs the normal FA - # path, so it must skip this DCP block entirely -- otherwise - # split_dcp_context_queries sets a DCP-split num_decode_tokens that - # corrupts do_kv_cache_update's PCP gather and garbles the cache. + # case (dcp_shares_pcp_ranks) is handled by the block above and runs + # the sharded _forward_dcp_mrv2 path, so it never reaches here. query_lens = query_start_loc[1:] - query_start_loc[:-1] context_kv_lens = seq_lens - query_lens local_context_kv_lens = get_dcp_local_seq_lens( @@ -1066,10 +1060,10 @@ def forward( v_descale=v_descale, ) return output - elif self.dcp_shares_pcp_ranks and envs.VLLM_PCP_DCP_SHARDED_KV_CACHE: - # MRv2 PCP+DCP with a SHARDED cache (VLLM_PCP_DCP_SHARDED_KV_CACHE): - # decode KV is sharded 1/dcp for the memory win. Decode attends - # its local shard + LSE-combine; prefill all-gathers the full KV. + elif self.dcp_shares_pcp_ranks: + # MRv2 PCP+DCP (dcp_shares_pcp_ranks): the cache is DCP-sharded + # (1/dcp KV per rank) for the memory win. Decode attends its + # local shard + LSE-combine; prefill all-gathers the full KV. self._forward_dcp_mrv2( query[:num_actual_tokens], key[:num_actual_tokens], @@ -1084,13 +1078,8 @@ def forward( ) return output else: - # dcp_world_size <= 1, or dcp_shares_pcp_ranks (MRv2 PCP+DCP). - # With the REPLICATED cache (model_runner.pcp_dcp_replicated_cache) - # the PCP+DCP case is equivalent to pure PCP, so the normal FA - # path is correct for both decode and prefill. The sharded-cache - # path (DCP decode-KV memory win: decode LSE-combine + prefill KV - # all-gather) lives in _forward_dcp_mrv2 and is reserved for a - # future opt-in (Phase 2b). + # dcp_world_size <= 1 (no DCP): pure PCP or unparallelized -- + # the normal FA path is correct for both decode and prefill. window = ( attn_metadata.sliding_window if attn_metadata.sliding_window is not None diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 8281e388e462..df8769c3f3a3 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -470,19 +470,8 @@ def __init__( self.block_size = self.kv_cache_spec.block_size self.dcp_world_size = dcp_world_size self.pcp_world_size = pcp_world_size - # MRv2 PCP+DCP (dcp == pcp > 1): replicated (default) -> no inflation; - # with VLLM_PCP_DCP_SHARDED_KV_CACHE the cache shards -> keep inflation. - cache_dcp = ( - 1 - if ( - dcp_world_size > 1 - and dcp_world_size == pcp_world_size - and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE - ) - else dcp_world_size - ) - if cache_dcp > 1: - self.block_size *= cache_dcp + if dcp_world_size > 1: + self.block_size *= dcp_world_size # For models using only Mamba, block_size is set to max_model_len when # prefix caching is disabled, and hash_block_size validation is skipped. assert not enable_caching or (hash_block_size == self.block_size), ( diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index b96d37a8ba9c..d4970d91b2d8 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -644,23 +644,14 @@ def resolve_kv_cache_block_sizes( """ cache_config = vllm_config.cache_config dcp = vllm_config.parallel_config.decode_context_parallel_size - pcp = vllm_config.parallel_config.prefill_context_parallel_size - # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache (default) so the - # scheduler/hash block size is NOT inflated by dcp; with - # VLLM_PCP_DCP_SHARDED_KV_CACHE the cache shards and keeps the inflation. - cache_dcp = ( - 1 - if (dcp > 1 and dcp == pcp and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE) - else dcp - ) groups = kv_cache_config.kv_cache_groups if len(groups) <= 1: - bs = cache_config.block_size * cache_dcp + bs = cache_config.block_size * dcp return bs, bs group_block_sizes = [ - g.kv_cache_spec.block_size * cache_dcp + g.kv_cache_spec.block_size * dcp if isinstance(g.kv_cache_spec, AttentionSpec) else g.kv_cache_spec.block_size for g in groups diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 675fd1288b89..f8578c68a389 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -6,7 +6,6 @@ from collections.abc import Sequence from typing import ClassVar -from vllm import envs from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import ( @@ -76,17 +75,8 @@ def __init__( self.block_size = kv_cache_spec.block_size self.dcp_world_size = dcp_world_size self.pcp_world_size = pcp_world_size - # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache (default) so do not - # inflate; with VLLM_PCP_DCP_SHARDED_KV_CACHE the cache is sharded and - # keeps the inflation. Pure-DCP (dcp != pcp) always shards. - replicated = ( - dcp_world_size > 1 - and dcp_world_size == pcp_world_size - and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE - ) - cache_dcp = 1 if replicated else dcp_world_size - if cache_dcp > 1: - self.block_size *= cache_dcp + if dcp_world_size > 1: + self.block_size *= dcp_world_size self.kv_cache_spec = kv_cache_spec self.block_pool = block_pool self.enable_caching = enable_caching diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index c1ede2029006..7a44d4e1d25e 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -13,7 +13,6 @@ import torch from typing_extensions import Self -from vllm import envs from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_up from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim @@ -220,12 +219,7 @@ def real_page_size_bytes(self) -> int: def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: parallel_config = vllm_config.parallel_config - dcp = parallel_config.decode_context_parallel_size - # MRv2 PCP+DCP (dcp == pcp > 1): replicated cache (default) -> no - # sharding divisor; with VLLM_PCP_DCP_SHARDED_KV_CACHE it shards. - pcp = parallel_config.prefill_context_parallel_size - replicated = dcp > 1 and dcp == pcp and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE - kv_shard_count = 1 if replicated else dcp + kv_shard_count = parallel_config.decode_context_parallel_size return cdiv(max_len, self.block_size * kv_shard_count) @@ -263,13 +257,9 @@ def __post_init__(self): def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_model_len = vllm_config.model_config.max_model_len - dcp = vllm_config.parallel_config.decode_context_parallel_size - pcp = vllm_config.parallel_config.prefill_context_parallel_size - # MRv2 PCP+DCP (dcp == pcp > 1): replicated (default) -> no 1/dcp saving; - # with VLLM_PCP_DCP_SHARDED_KV_CACHE it shards -> apply the 1/dcp saving. - replicated = dcp > 1 and dcp == pcp and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE - if dcp > 1 and not replicated: - max_model_len = cdiv(max_model_len, dcp) + dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size + if dcp_world_size > 1: + max_model_len = cdiv(max_model_len, dcp_world_size) return cdiv(max_model_len, self.block_size) * self.page_size_bytes @classmethod diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3ffc28b1c4b6..6db6cad80d24 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -177,20 +177,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0 self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size - # MRv2 PCP+DCP (dcp reuses the pcp ranks, i.e. dcp == pcp > 1): the KV - # cache is kept REPLICATED -- equivalent to pure PCP -- so the prefill - # PCP all-gather write and the FA reads share one non-sharded layout. - # This is correct but makes the dcp dimension cache-redundant (no - # decode-KV memory win); the sharded-cache path is Phase 2b. - self.pcp_size = self.parallel_config.prefill_context_parallel_size - self.pcp_dcp_replicated_cache = ( - self.use_dcp - and self.pcp_size == self.dcp_size - and not envs.VLLM_PCP_DCP_SHARDED_KV_CACHE - ) - self.cache_dcp_size = 1 if self.pcp_dcp_replicated_cache else self.dcp_size - self.cache_dcp_rank = 0 if self.pcp_dcp_replicated_cache else self.dcp_rank - # Multimodal self.mm_registry = MULTIMODAL_REGISTRY self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( @@ -445,7 +431,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: # As a result, one block on the current rank covers `block_size * cp_size` # tokens in the full, global (unsharded) sequence. max_num_blocks = cdiv( - block_table_max_model_len, spec.block_size * self.cache_dcp_size + block_table_max_model_len, spec.block_size * self.dcp_size ) # Align to a multiple of (128 / block_size) as required by some attention # backends such as TRTLLM (#39324) @@ -469,8 +455,8 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: max_num_blocks_per_group=max_num_blocks_per_group, device=self.device, kernel_block_sizes=self.kernel_block_sizes, - cp_size=self.cache_dcp_size, - cp_rank=self.cache_dcp_rank, + cp_size=self.dcp_size, + cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, ) self.pcp_manager = pcp.maybe_build_pcp_manager( From 59d2b9172ee1a023ac92dc8a5fb5e88906f682e1 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Fri, 24 Jul 2026 02:15:47 +0000 Subject: [PATCH 09/18] clean more Signed-off-by: JaredforReal --- vllm/model_executor/layers/attention/pcp.py | 58 +++++ vllm/v1/attention/backends/flash_attn.py | 253 +++++++++----------- 2 files changed, 174 insertions(+), 137 deletions(-) diff --git a/vllm/model_executor/layers/attention/pcp.py b/vllm/model_executor/layers/attention/pcp.py index 3b65cde96e86..d82055d352f3 100644 --- a/vllm/model_executor/layers/attention/pcp.py +++ b/vllm/model_executor/layers/attention/pcp.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import NamedTuple + import torch from vllm.distributed.parallel_state import ( @@ -150,6 +152,62 @@ def slice_prefill_output_local( return local[:num_actual] +class DecodeSubset(NamedTuple): + """Decode-token subset of a mixed prefill+decode batch.""" + + token_mask: torch.Tensor + q: torch.Tensor + k: torch.Tensor + v: torch.Tensor + cu_seqlens: torch.Tensor + max_seqlen: int + ctx_kv_lens: torch.Tensor + block_table: torch.Tensor + + +def build_mixed_decode_subset( + is_prefilling: torch.Tensor | None, + query_start_loc: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + dcp_context_kv_lens: torch.Tensor, + block_table: torch.Tensor, +) -> DecodeSubset | None: + """Extract the decode-token subset of a mixed prefill+decode batch. + + DualChunkSwap gives ranks with no prefill chunks a 0-token dummy decode + segment mirroring global req 0; such segments are excluded (counting them + would desync the DCP LSE-combine). Returns None for a pure-prefill batch. + """ + if is_prefilling is None: + return None + is_pre = is_prefilling.to(device=query_start_loc.device) + seg_lens = query_start_loc[1:] - query_start_loc[:-1] + decode_seg = (~is_pre.bool()) & (seg_lens > 0) + num_decode = int(decode_seg.sum().item()) + if num_decode == 0: + return None + token_mask = torch.repeat_interleave(decode_seg, seg_lens)[: query.shape[0]] + dec_lens = seg_lens[decode_seg] + dec_cu = torch.zeros( + num_decode + 1, + dtype=query_start_loc.dtype, + device=query_start_loc.device, + ) + torch.cumsum(dec_lens, dim=0, out=dec_cu[1:]) + return DecodeSubset( + token_mask=token_mask, + q=query[token_mask].contiguous(), + k=key[token_mask].contiguous(), + v=value[token_mask].contiguous(), + cu_seqlens=dec_cu, + max_seqlen=int(dec_lens.max().item()), + ctx_kv_lens=dcp_context_kv_lens[decode_seg], + block_table=block_table[decode_seg], + ) + + def finalize_mla_pcp_decode( output: torch.Tensor, num_heads: int, diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 36d849f36785..0f79fa5e0969 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -11,6 +11,7 @@ from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention.pcp import ( + build_mixed_decode_subset, gather_prefill_qkv_global, slice_prefill_output_local, ) @@ -593,17 +594,10 @@ def schedule( prefix_scheduler_metadata = None if self.dcp_shares_pcp_ranks and self.dcp_world_size > 1: - # MRv2 PCP+DCP (dcp_shares_pcp_ranks): the cache is DCP-sharded - # (1/dcp KV per rank). DualChunkSwap has already partitioned the - # batch per rank; DCP only shards the cache. Both decode and prefill - # use the context+query+merge path (_forward_dcp_mrv2_decode): the - # context attention reads each rank's local shard of the causal - # prefix (the full prefix is in the sharded cache -- - # do_kv_cache_update all-gathered+wrote it before forward) and - # LSE-combines; the query attention reads this rank's own - # new K/V. context_kv_lens = seq_lens - query_lens is the absolute - # position of the chunk start (= the prior prefix length), correct for - # both decode (query_len=1) and prefill. scheduler_metadata is unused. + # MRv2 PCP+DCP: cache is DCP-sharded (1/dcp KV/rank). DualChunkSwap + # already partitioned the batch per rank; both decode and prefill use + # the context+query+merge path. context_kv_lens = seq_lens - + # query_lens is the chunk-start position (= prior prefix length). scheduler_metadata = None query_lens = query_start_loc[1:] - query_start_loc[:-1] context_kv_lens = seq_lens - query_lens @@ -1248,13 +1242,9 @@ def do_kv_cache_update( return if self.use_pcp: - # Under PCP each rank holds only its DualChunkSwap prefill chunks in - # the rank-local batch. All-gather the prefill K/V across PCP ranks - # (keeping decode writes local) so every rank's cache receives the - # full prefill KV (replicated when dcp=1; the DCP-local shard when - # pcp+dcp). maybe_gather_kv_cache_inputs gathers K/V separately -> - # contiguous (cache kernel head-stride contract), builds the gathered - # slot mapping, and handles the decode/prefill split + empty-prefill. + # All-gather the prefill K/V across PCP ranks (decode writes stay + # local) so every rank's cache gets the full prefill KV. K/V are + # gathered separately to stay contiguous (cache kernel head-stride). from vllm.model_executor.layers.attention.pcp import ( maybe_gather_kv_cache_inputs, ) @@ -1263,12 +1253,9 @@ def do_kv_cache_update( num_decode_tokens = ( attn_metadata.num_decode_tokens if attn_metadata is not None else 0 ) - # Rank-invariant gather decision: the per-rank num_decode_tokens can - # differ under DualChunkSwap (some ranks get zero prefill chunks), - # which would desync the PCP all-gather. If the global batch has any - # prefill, every rank treats the whole batch as prefill (gather all); - # only a pure-decode batch skips the gather. Used by both the sharded - # and replicated PCP+DCP paths. + # Rank-invariant gather decision: per-rank num_decode_tokens can differ + # under DualChunkSwap and desync the all-gather, so if the global batch + # has any prefill every rank gathers the whole batch. gflags = get_forward_context().additional_kwargs.get("pcp_global_flags") if gflags is not None and gflags[0]: num_decode_tokens = 0 @@ -1335,6 +1322,48 @@ def _fa_common_kwargs( v_descale=v_descale, ) + def _sharded_context_attention( + self, + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + ctx_kv_lens: torch.Tensor, + max_ctx_kv: int, + block_table: torch.Tensor, + fa_kw: dict, + out: torch.Tensor | None = None, + num_splits: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Attend ``q`` to the sharded cached prefix, then LSE-combine across DCP. + + Returns the combined output and corrected LSE ([B, H] layout). ``out`` + and ``num_splits`` are passed through only when given. + """ + kw = dict( + q=q, + k=key_cache, + v=value_cache, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + seqused_k=ctx_kv_lens, + max_seqlen_k=max_ctx_kv, + causal=False, + block_table=block_table, + return_softmax_lse=True, + **fa_kw, + ) + if out is not None: + kw["out"] = out + if num_splits is not None: + kw["num_splits"] = num_splits + ctx_out, ctx_lse = flash_attn_varlen_func(**kw) + ctx_out_cor, ctx_lse_cor = self.dcp_combine( + ctx_out, ctx_lse.transpose(0, 1), get_dcp_group(), return_lse=True + ) + return ctx_out_cor, ctx_lse_cor.transpose(0, 1).contiguous() + def _forward_dcp_mrv2( self, query: torch.Tensor, @@ -1348,22 +1377,20 @@ def _forward_dcp_mrv2( k_descale: torch.Tensor | None = None, v_descale: torch.Tensor | None = None, ) -> torch.Tensor: - """MRv2 PCP+DCP GQA attention over a SHARDED cache. - - Pure-decode batches use ``_forward_dcp_mrv2_decode`` (Q replicated, - attend the sharded prefix + LSE-combine). Prefill/extend/mixed batches - all-gather the DualChunkSwap-partitioned Q/K/V across PCP, run full - causal on the gathered tokens, optionally context-attend the cached - prefix (extend), and slice back; decode tokens in a mixed batch are then - overwritten via the decode path. + """MRv2 PCP+DCP GQA attention over a sharded (1/dcp) KV cache. + + Pure-decode uses ``_forward_dcp_mrv2_decode`` (Q replicated, attend the + sharded prefix + LSE-combine). Prefill/extend/mixed all-gather the + DualChunkSwap-partitioned Q/K/V, run full causal, optionally + context-attend the cached prefix (extend), and slice back; decode tokens + in a mixed batch are then overwritten via the decode path. """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." ) fc = get_forward_context() - # Rank-invariant (has_prefill, has_decode) from the global batch; per-rank - # num_decode_tokens is not safe to route on (DualChunkSwap can zero some - # ranks' prefill chunks and desync the collectives). + # Route on rank-invariant global flags, not per-rank num_decode_tokens: + # DualChunkSwap can zero some ranks' prefill chunks and desync collectives. global_has_prefill, global_has_decode = fc.additional_kwargs.get( "pcp_global_flags", (True, False) ) @@ -1391,14 +1418,12 @@ def _forward_dcp_mrv2( ) return output - # Prefill/extend/mixed: DualChunkSwap partitions the prefill Q, so each - # rank all-gathers Q/K/V across PCP, runs full causal on the gathered - # tokens, and slices its chunk back out. Decode tokens in a mixed batch - # self-attend here (no cached context) and are overwritten below. + # Prefill/extend/mixed: gather the DualChunkSwap-partitioned Q/K/V across + # PCP, run full causal on the gathered tokens, slice back. Decode tokens + # in a mixed batch self-attend here (no cached context) and are fixed below. info = fc.additional_kwargs.get("pcp_prefill_gather") if info is None: - # Warmup / dummy run: no PCP partition -> no gather indices. Fall - # back to local causal attention (shape-correct is all warmup needs). + # Warmup/dummy: no PCP partition -> local causal (shape-correct only). flash_attn_varlen_func( q=query, k=key, @@ -1433,47 +1458,30 @@ def _forward_dcp_mrv2( return_softmax_lse=True, **fa_kw, ) - # EXTEND (prefix caching): q_g (replicated) also attends the cached - # prefix -- each rank its DCP shard, LSE-combined. Skipped for pure - # first-prefills (rank-invariant has_extend = gctx[3]). + # EXTEND (prefix caching): q_g also attends the cached prefix, sharded + + # LSE-combined. Skipped for pure first-prefills (rank-invariant has_extend). final_g = out_g gctx = fc.additional_kwargs.get("pcp_global_ctx") if gctx is not None and gctx[3] and key_cache is not None: g_block_table, g_ctx_kv_lens, g_max_ctx_kv, _ = gctx - ctx_out_g, ctx_lse_g = flash_attn_varlen_func( - q=q_g, - k=key_cache, - v=value_cache, - cu_seqlens_q=cu_seqlens, - max_seqlen_q=global_max_seqlen, - seqused_k=g_ctx_kv_lens, - max_seqlen_k=g_max_ctx_kv, - causal=False, - block_table=g_block_table, - return_softmax_lse=True, - **fa_kw, - ) - ctx_out_cor, ctx_lse_cor = self.dcp_combine( - ctx_out_g, - ctx_lse_g.transpose(0, 1), - get_dcp_group(), - return_lse=True, + ctx_out_cor, ctx_lse_cor = self._sharded_context_attention( + q_g, + key_cache, + value_cache, + cu_seqlens, + global_max_seqlen, + g_ctx_kv_lens, + g_max_ctx_kv, + g_block_table, + fa_kw, ) - ctx_lse_cor = ctx_lse_cor.transpose(0, 1).contiguous() final_g = torch.empty_like(out_g) - merge_attn_states( - final_g, - ctx_out_cor, - ctx_lse_cor, - out_g, - query_lse_g, - ) + merge_attn_states(final_g, ctx_out_cor, ctx_lse_cor, out_g, query_lse_g) # global -> gathered (rank-order) -> this rank's local chunk. output[:num_actual] = slice_prefill_output_local(info, final_g, num_actual) - # MIXED batch: the prefill gather above gave decode tokens only - # self-attention (1-token segments, no cached context). Overwrite them - # with the correct decode computation. No-op for a pure-prefill batch. + # MIXED batch: the gather above self-attended decode tokens (no cached + # context); recompute them. No-op for a pure-prefill batch. if global_has_decode: self._overwrite_mixed_decode_tokens( query=query, @@ -1498,55 +1506,37 @@ def _overwrite_mixed_decode_tokens( attn_metadata: FlashAttentionMetadata, fa_kw: dict, ) -> None: - """Recompute decode-token outputs the prefill-gather path self-attended. - - Decode tokens are replicated across PCP ranks; extract them, run the - decode path (sharded-context + LSE-combine) on the subset, and scatter - back. Zero-token segments (DualChunkSwap's dummy for empty ranks) are - excluded -- counting them would desync the DCP LSE-combine. - """ - is_pre = attn_metadata.is_prefilling - if is_pre is None: - return - qsl = attn_metadata.query_start_loc - is_pre = is_pre.to(device=qsl.device) - seg_query_lens = qsl[1:] - qsl[:-1] - decode_seg_mask = (~is_pre.bool()) & (seg_query_lens > 0) - num_decode = int(decode_seg_mask.sum().item()) - if num_decode == 0: - return - - token_mask = torch.repeat_interleave(decode_seg_mask, seg_query_lens) - token_mask = token_mask[: query.shape[0]] - dec_q = query[token_mask].contiguous() - dec_k = key[token_mask].contiguous() - dec_v = value[token_mask].contiguous() - - dec_query_lens = seg_query_lens[decode_seg_mask] - dec_cu_seqlens = torch.zeros(num_decode + 1, dtype=qsl.dtype, device=qsl.device) - torch.cumsum(dec_query_lens, dim=0, out=dec_cu_seqlens[1:]) - + """Recompute decode-token outputs the prefill-gather path self-attended.""" ctx_kv_lens = attn_metadata.dcp_context_kv_lens assert ctx_kv_lens is not None - dec_ctx_kv_lens = ctx_kv_lens[decode_seg_mask] - - dec_out = output[token_mask] + sub = build_mixed_decode_subset( + attn_metadata.is_prefilling, + attn_metadata.query_start_loc, + query, + key, + value, + ctx_kv_lens, + attn_metadata.block_table, + ) + if sub is None: + return + dec_out = output[sub.token_mask] self._forward_dcp_mrv2_decode( - dec_q, - dec_k, - dec_v, + sub.q, + sub.k, + sub.v, key_cache, value_cache, dec_out, attn_metadata, - dec_cu_seqlens, - int(dec_query_lens.max().item()), - attn_metadata.block_table[decode_seg_mask], + sub.cu_seqlens, + sub.max_seqlen, + sub.block_table, fa_kw, - dcp_context_kv_lens=dec_ctx_kv_lens, - max_dcp_context_kv_len=int(dec_ctx_kv_lens.max().item()), + dcp_context_kv_lens=sub.ctx_kv_lens, + max_dcp_context_kv_len=int(sub.ctx_kv_lens.max().item()), ) - output[token_mask] = dec_out + output[sub.token_mask] = dec_out def _forward_dcp_mrv2_decode( self, @@ -1565,10 +1555,9 @@ def _forward_dcp_mrv2_decode( max_dcp_context_kv_len: int | None = None, ) -> None: """Decode (or decode-subset) over a sharded cache: attend the local DCP - shard of the cached prefix + LSE-combine, merge with the per-token new - K/V attention. Every rank runs the combine (no rank-varying guard) so - NCCL stays synchronized. Context lens default to the metadata but can be - overridden for a decode-only subset of a mixed batch. + shard of the prefix + LSE-combine, merge with per-token new K/V. Every + rank runs the combine (no rank-varying guard). Context lens default to + the metadata but can be overridden for a mixed-batch decode subset. """ if dcp_context_kv_lens is None: dcp_context_kv_lens = attn_metadata.dcp_context_kv_lens @@ -1584,29 +1573,19 @@ def _forward_dcp_mrv2_decode( ) assert dcp_context_kv_lens is not None assert max_dcp_context_kv_len is not None - context_attn_out, context_lse = flash_attn_varlen_func( - q=query, - k=key_cache, - v=value_cache, + context_attn_out_cor, context_lse_cor = self._sharded_context_attention( + query, + key_cache, + value_cache, + cu_seqlens_q, + max_seqlen_q, + dcp_context_kv_lens, + max_dcp_context_kv_len, + block_table, + fa_kw, out=dcp_context_out_workspace[:n], - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - seqused_k=dcp_context_kv_lens, - max_seqlen_k=max_dcp_context_kv_len, - causal=False, - block_table=block_table, - return_softmax_lse=True, num_splits=attn_metadata.max_num_splits, - **fa_kw, ) - # FA returns LSE as [H, B]; the DCP combine wants [B, H]. - context_attn_out_cor, context_lse_cor = self.dcp_combine( - context_attn_out, - context_lse.transpose(0, 1), - get_dcp_group(), - return_lse=True, - ) - context_lse_cor = context_lse_cor.transpose(0, 1).contiguous() query_attn_out, query_lse = flash_attn_varlen_func( q=query, From 019fcdbebec0d688f0c6ccf91094c5042667332e Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Wed, 29 Jul 2026 07:35:59 +0000 Subject: [PATCH 10/18] apply mla & gqa sharing PCP/DCP query-gather and LSE-combine Signed-off-by: JaredforReal Co-authored-by: Lucas Wilkinson --- vllm/config/model.py | 18 +-- .../layers/attention/mla_attention.py | 64 +++------- vllm/model_executor/layers/attention/pcp.py | 97 ++++++++++++++- vllm/v1/attention/backends/flash_attn.py | 115 ++++++------------ 4 files changed, 155 insertions(+), 139 deletions(-) diff --git a/vllm/config/model.py b/vllm/config/model.py index e0b160f6f8d8..54996eee77c9 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1251,21 +1251,13 @@ def verify_with_parallel_config( ) decode_context_parallel_size = parallel_config.decode_context_parallel_size - # When DCP spans the PCP axis (dcp == pcp, pcp > 1) it does NOT shard - # the TP-local KV heads: Q is replicated across PCP ranks and the - # partial attentions combine via all-reduce (see FlashAttention - # dcp_shares_pcp_ranks). So the TP-head DCP constraints below apply - # only when DCP reuses the TP ranks (pcp == 1) or spans the full - # TP x PCP axis (dcp == tp*pcp). - dcp_shares_pcp_ranks = ( - parallel_config.prefill_context_parallel_size > 1 - and decode_context_parallel_size - == parallel_config.prefill_context_parallel_size - ) + # DCP groups span the PCP axis before TP, so DCP only splits query heads + # when it reaches past PCP (see pcp.maybe_all_gather_q_for_dcp). When + # dcp == pcp the group holds replicated Q and none of these TP-head + # constraints apply. if ( - decode_context_parallel_size > 1 + decode_context_parallel_size > parallel_config.prefill_context_parallel_size and not self.use_mla - and not dcp_shares_pcp_ranks ): total_num_kv_heads = self.get_total_num_kv_heads() assert tensor_parallel_size > total_num_kv_heads, ( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 884c4ca4e687..83239e6e00b2 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -211,7 +211,6 @@ from vllm.config.cache import CacheDType from vllm.distributed.parallel_state import ( get_dcp_group, - get_tp_group, is_global_first_rank, ) from vllm.forward_context import ForwardContext, get_forward_context @@ -227,8 +226,10 @@ maybe_transfer_kv_layer, ) from vllm.model_executor.layers.attention.pcp import ( - finalize_mla_pcp_decode, + cp_reconcile_heads, + maybe_all_gather_q_for_dcp, maybe_gather_mla_latent_cache_inputs, + resolve_dcp_combine_fn, ) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.linear import ( @@ -274,8 +275,6 @@ get_dcp_local_seq_lens, split_decodes_and_prefills, ) -from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs -from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.attention.ops.triton_merge_attn_states import mask_empty_context from vllm.v1.attention.selector import get_attn_backend @@ -543,12 +542,7 @@ def __init__( self.use_sparse = use_sparse - _vllm_config = get_current_vllm_config_or_none() - self.dcp_a2a = ( - _vllm_config is not None - and _vllm_config.parallel_config.decode_context_parallel_size > 1 - and _vllm_config.parallel_config.dcp_comm_backend == "a2a" - ) + self.dcp_combine = resolve_dcp_combine_fn(get_current_vllm_config_or_none()) # Initialize q/k/v range constants. self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32) @@ -873,18 +867,15 @@ def forward_impl( mqa_q = (mqa_ql_nope, mqa_q_pe) # concatenate nope + pe -> (B, N, L + P) (fp8 op above may have fused) if self.impl.dcp_world_size > 1: - if self.use_pcp: - if self.impl.dcp_world_size > self.impl.pcp_world_size: - if isinstance(mqa_q, tuple): - mqa_q = torch.cat(mqa_q, dim=-1) - mqa_q = get_tp_group().all_gather(mqa_q, dim=1) - else: - if isinstance(mqa_q, tuple): - # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) - mqa_q = torch.cat(mqa_q, dim=-1) - if not qrep_decode: - # mqa_q do allgather in head dim. - mqa_q = get_dcp_group().all_gather(mqa_q, dim=1) + if not self.use_pcp and isinstance(mqa_q, tuple): + # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) + mqa_q = torch.cat(mqa_q, dim=-1) + mqa_q = maybe_all_gather_q_for_dcp( + mqa_q, + self.impl.dcp_world_size, + self.impl.pcp_world_size, + already_replicated=qrep_decode and not self.use_pcp, + ) # call decode attn if not self.impl.is_sparse: @@ -894,29 +885,14 @@ def forward_impl( # correct dcp attn_out with lse. if self.impl.dcp_world_size > 1: assert lse is not None - if self.dcp_a2a: - attn_out = dcp_a2a_lse_reduce( - attn_out, - lse, - get_dcp_group(), - is_lse_base_on_e=self.impl.lse_base_on_e, - ) - elif self.use_pcp: - attn_out = cp_lse_ag_out_ar( - attn_out, - lse, - get_dcp_group(), - is_lse_base_on_e=self.impl.lse_base_on_e, - ) - else: - attn_out = cp_lse_ag_out_rs( - attn_out, - lse, - get_dcp_group(), - is_lse_base_on_e=self.impl.lse_base_on_e, - ) + attn_out = self.dcp_combine( + attn_out, + lse, + get_dcp_group(), + is_lse_base_on_e=self.impl.lse_base_on_e, + ) if self.use_pcp: - attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads) + attn_out = cp_reconcile_heads(attn_out, self.num_heads) # v_up projection self._v_up_proj(attn_out, out=mqa_output_slice) diff --git a/vllm/model_executor/layers/attention/pcp.py b/vllm/model_executor/layers/attention/pcp.py index d82055d352f3..97d09e5aafec 100644 --- a/vllm/model_executor/layers/attention/pcp.py +++ b/vllm/model_executor/layers/attention/pcp.py @@ -1,13 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import NamedTuple +from typing import TYPE_CHECKING, NamedTuple, TypeVar import torch from vllm.distributed.parallel_state import ( + GroupCoordinator, + get_dcp_group, get_pcp_group, get_tp_group, ) +from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs +from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +# Query as handed to the DCP context attention: a tensor, or MLA's split +# (nope, pe) pair which is only concatenated if a gather happens. +_QueryT = TypeVar("_QueryT", torch.Tensor, tuple[torch.Tensor, ...]) def _gather_prefill_cache_inputs( @@ -208,7 +219,89 @@ def build_mixed_decode_subset( ) -def finalize_mla_pcp_decode( +def _dcp_q_gather_group( + dcp_world_size: int, + pcp_world_size: int, +) -> GroupCoordinator | None: + """Group holding the query-head shards this rank must gather, or None. + + DCP groups span the PCP axis before TP (see ``init_model_parallel_groups``), + so how far the group reaches decides which ranks hold distinct query heads: + + - ``pcp == 1``: the DCP group is a TP subgroup that splits Q heads, so + gather over the DCP group. + - ``dcp > pcp``: the group spans TP x PCP. Q is head-sharded across TP and + replicated across PCP, so gather over the TP group -- gathering over the + DCP group here would duplicate the PCP-replicated heads. + - otherwise: the DCP group is the PCP group and already holds every head. + """ + if pcp_world_size == 1: + return get_dcp_group() if dcp_world_size > 1 else None + if dcp_world_size > pcp_world_size: + return get_tp_group() + return None + + +def dcp_q_gather_size(dcp_world_size: int, pcp_world_size: int) -> int: + """Head-shard count ``maybe_all_gather_q_for_dcp`` will produce. + + For sizing work that has to be planned before the gather happens, such as + FlashAttention's AOT scheduler metadata. 1 means no gather. + """ + group = _dcp_q_gather_group(dcp_world_size, pcp_world_size) + return 1 if group is None else group.world_size + + +def maybe_all_gather_q_for_dcp( + q: _QueryT, + dcp_world_size: int, + pcp_world_size: int, + already_replicated: bool = False, +) -> _QueryT | torch.Tensor: + """All-gather query heads for attention against the DCP-sharded cache. + + Returns ``q`` untouched when this rank already holds every head the kernel + needs; ``already_replicated`` lets a caller declare that up front. A split + ``q`` -- MLA's ``(nope, pe)`` pair -- is concatenated only when a gather + actually happens, so callers that can consume the split form keep it. + + Pairs with ``resolve_dcp_combine_fn``: changing one without the other breaks + the head layout the combine expects. + """ + if already_replicated: + return q + group = _dcp_q_gather_group(dcp_world_size, pcp_world_size) + if group is None: + return q + gathered = torch.cat(q, dim=-1) if isinstance(q, tuple) else q + return group.all_gather(gathered, dim=1) + + +def resolve_dcp_combine_fn(vllm_config: "VllmConfig | None"): + """LSE-combine to fold per-rank partial attentions over the DCP group. + + Under PCP the partials cover the full head set on every rank, so they + all-reduce and each rank takes its own heads back out with + ``cp_reconcile_heads``. Without PCP the DCP group is a TP subgroup, so the + reduce-scatter lands each rank's heads directly. + + Pairs with ``maybe_all_gather_q_for_dcp``. ``vllm_config`` may be None when + there is no config in scope (unit tests), which implies no CP. + """ + if vllm_config is None: + return cp_lse_ag_out_rs + parallel_config = vllm_config.parallel_config + if ( + parallel_config.decode_context_parallel_size > 1 + and parallel_config.dcp_comm_backend == "a2a" + ): + return dcp_a2a_lse_reduce + if parallel_config.prefill_context_parallel_size > 1: + return cp_lse_ag_out_ar + return cp_lse_ag_out_rs + + +def cp_reconcile_heads( output: torch.Tensor, num_heads: int, ) -> torch.Tensor: diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index adb76afd0ce8..3579a3ab6b86 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -12,7 +12,10 @@ from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention.pcp import ( build_mixed_decode_subset, + dcp_q_gather_size, gather_prefill_qkv_global, + maybe_all_gather_q_for_dcp, + resolve_dcp_combine_fn, slice_prefill_output_local, ) from vllm.platforms import current_platform @@ -34,8 +37,6 @@ is_flash_attn_varlen_func_available, ) from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens -from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs -from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.worker.workspace import current_workspace_manager @@ -433,16 +434,7 @@ def __init__( self.parallel_config.cp_kv_cache_interleave_size ) - # DCP shares the PCP ranks (dcp == pcp): the MRv2 PCP+DCP path - # (_forward_dcp_mrv2) is used instead of _forward_with_dcp. Q is - # replicated across these ranks (no DCP head gather), so the - # scheduler_metadata for the prefill path must use num_heads_q without - # the DCP multiply. - self.dcp_shares_pcp_ranks = ( - self.dcp_world_size > 1 - and self.parallel_config.decode_context_parallel_size - == self.parallel_config.prefill_context_parallel_size - ) + self.use_pcp = self.pcp_world_size > 1 self.use_full_cuda_graph = ( self.compilation_config.cudagraph_mode.has_full_cudagraphs() @@ -560,17 +552,16 @@ def schedule( seqlens, max_seq_len, causal, - num_heads_q=None, ): cache_dtype = self.cache_config.cache_dtype if is_quantized_kv_cache(cache_dtype): qkv_dtype = current_platform.fp8_dtype() else: qkv_dtype = self.kv_cache_dtype - if num_heads_q is None: - # DCP that reuses the TP ranks gathers Q across DCP heads before - # the kernel, so the scheduler sees num_heads_q * dcp heads. - num_heads_q = self.num_heads_q * self.dcp_world_size + # The kernel sees whatever maybe_all_gather_q_for_dcp() hands it. + num_heads_q = self.num_heads_q * dcp_q_gather_size( + self.dcp_world_size, self.pcp_world_size + ) if aot_schedule: return get_scheduler_metadata( batch_size=batch_size, @@ -604,7 +595,7 @@ def schedule( suffix_kv_lens = None prefix_scheduler_metadata = None - if self.dcp_shares_pcp_ranks and self.dcp_world_size > 1: + if self.use_pcp and self.dcp_world_size > 1: # MRv2 PCP+DCP: cache is DCP-sharded (1/dcp KV/rank). DualChunkSwap # already partitioned the batch per rank; both decode and prefill use # the context+query+merge path. context_kv_lens = seq_lens - @@ -637,10 +628,10 @@ def schedule( num_prefill_tokens = num_actual_tokens else: num_decode_tokens = num_actual_tokens - elif self.dcp_world_size > 1 and not self.dcp_shares_pcp_ranks: - # Pure-DCP path (DCP reuses the TP ranks; no PCP). The MRv2 PCP+DCP - # case (dcp_shares_pcp_ranks) is handled by the block above and runs - # the sharded _forward_dcp_mrv2 path, so it never reaches here. + elif self.dcp_world_size > 1: + # Pure-DCP path (DCP reuses the TP ranks; no PCP). The PCP+DCP case + # is handled by the block above and runs the sharded + # _forward_dcp_mrv2 path, so it never reaches here. query_lens = query_start_loc[1:] - query_start_loc[:-1] context_kv_lens = seq_lens - query_lens local_context_kv_lens = get_dcp_local_seq_lens( @@ -693,29 +684,14 @@ def schedule( (max_seq_len + num_partitions - 1) // num_partitions ) * self.cp_kv_cache_interleave_size - if self.dcp_shares_pcp_ranks: - # MRv2 PCP+DCP: prefill attends the full replicated cache - # with standard FA (_forward_dcp_mrv2); decode does not use - # scheduler_metadata. Build the standard scheduler_metadata - # (full seq_lens; Q is replicated so no DCP head multiply). - scheduler_metadata = schedule( - batch_size=num_reqs, - cu_query_lens=query_start_loc, - max_query_len=max_query_len, - seqlens=seq_lens, - max_seq_len=max_seq_len, - causal=causal, - num_heads_q=self.num_heads_q, - ) - else: - scheduler_metadata = schedule( - batch_size=num_reqs, - cu_query_lens=query_start_loc, - max_query_len=max_query_len, - seqlens=dcp_context_kv_lens, - max_seq_len=max_dcp_context_kv_len, - causal=False, - ) + scheduler_metadata = schedule( + batch_size=num_reqs, + cu_query_lens=query_start_loc, + max_query_len=max_query_len, + seqlens=dcp_context_kv_lens, + max_seq_len=max_dcp_context_kv_len, + causal=False, + ) elif use_cascade: cu_prefix_query_lens = torch.tensor( [0, num_actual_tokens], dtype=torch.int32, device=self.device @@ -923,25 +899,11 @@ def __init__( self.supports_quant_query_input = flash_attn_supports_quant_query_input() vllm_config = get_current_vllm_config_or_none() - dcp_a2a = ( - vllm_config is not None - and vllm_config.parallel_config.decode_context_parallel_size > 1 - and vllm_config.parallel_config.dcp_comm_backend == "a2a" - ) - # When DCP shares the PCP ranks (dcp == pcp), Q is replicated across the - # DCP ranks, so the partial attentions combine with an all-reduce (every - # rank ends with the full output) instead of the gathered-Q + - # reduce-scatter used when DCP reuses the TP ranks. - self.dcp_shares_pcp_ranks = ( - vllm_config is not None - and vllm_config.parallel_config.decode_context_parallel_size > 1 - and vllm_config.parallel_config.decode_context_parallel_size - == vllm_config.parallel_config.prefill_context_parallel_size - ) - if self.dcp_shares_pcp_ranks: - self.dcp_combine = cp_lse_ag_out_ar - else: - self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs + # self.pcp_world_size is auto-populated by AttentionImplBase.__new__. + self.use_pcp = self.pcp_world_size > 1 + # How Q and the partials move across the DCP group is decided by the + # PCP/DCP topology alone; both backends share the rule (see pcp.py). + self.dcp_combine = resolve_dcp_combine_fn(vllm_config) self._dcp_dtype: torch.dtype | None = None self._dcp_max_num_tokens: int = 0 @@ -1064,7 +1026,7 @@ def forward( k_descale = layer._k_scale.expand(descale_shape) v_descale = layer._v_scale.expand(descale_shape) - if self.dcp_world_size > 1 and not self.dcp_shares_pcp_ranks: + if self.dcp_world_size > 1 and not self.use_pcp: # Pure DCP (DCP reuses the TP ranks; no PCP). MRv1-era sharded # path -- the cache is DCP-sharded and writes are rank-local. self._forward_with_dcp( @@ -1080,8 +1042,8 @@ def forward( v_descale=v_descale, ) return output - elif self.dcp_shares_pcp_ranks: - # MRv2 PCP+DCP (dcp_shares_pcp_ranks): the cache is DCP-sharded + elif self.dcp_world_size > 1: + # MRv2 PCP+DCP (dcp == pcp): the cache is DCP-sharded # (1/dcp KV per rank) for the memory win. Decode attends its # local shard + LSE-combine; prefill all-gathers the full KV. self._forward_dcp_mrv2( @@ -1685,19 +1647,13 @@ def _forward_with_dcp( ) return output - # When DCP shares the PCP ranks (dcp == pcp), Q is replicated across the - # DCP ranks, so attend directly with local heads and combine via - # all-reduce. Otherwise gather Q across heads and reduce-scatter. - if self.dcp_shares_pcp_ranks: - query_for_context = query - context_num_heads = self.num_heads - else: - query_for_context = get_dcp_group().all_gather(query, dim=1) - context_num_heads = self.num_heads * self.dcp_world_size + query_for_context = maybe_all_gather_q_for_dcp( + query, self.dcp_world_size, self.pcp_world_size + ) sliding_window_size = ( list(self.sliding_window) if self.sliding_window is not None else None ) - n = query_for_context.shape[0] + n, context_num_heads = query_for_context.shape[:2] num_reqs = cu_seqlens_q.shape[0] - 1 num_decodes = attn_metadata.num_decode_reqs num_context_prefills = attn_metadata.num_prefill_reqs @@ -1727,9 +1683,8 @@ def _forward_with_dcp( if split_dcp_context: # TODO: Remove this DCP + FA2 mixed decode/prefill workaround once # FA4 supports this Qwen3.5 shape. - assert not self.dcp_shares_pcp_ranks, ( - "FA2 split-DCP context path does not support dcp_shares_pcp_ranks" - " (pcp+dcp); use FA3/FA4." + assert not self.use_pcp, ( + "FA2 split-DCP context path does not support PCP+DCP; use FA3/FA4." ) assert attn_metadata.dcp_context_kv_lens is not None assert attn_metadata.max_dcp_context_kv_len is not None From 481d21aef4bbec25aee00f80cc1494910481785d Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Thu, 30 Jul 2026 05:27:05 +0000 Subject: [PATCH 11/18] refactor pcp x dcp Signed-off-by: JaredforReal Co-authored-by: Lucas Wilkinson --- vllm/model_executor/layers/attention/pcp.py | 98 +----- vllm/v1/attention/backends/flash_attn.py | 320 +++++++++--------- vllm/v1/worker/gpu/pcp_manager.py | 349 ++++++++++++++------ 3 files changed, 399 insertions(+), 368 deletions(-) diff --git a/vllm/model_executor/layers/attention/pcp.py b/vllm/model_executor/layers/attention/pcp.py index 97d09e5aafec..a4bc4ec6c4b1 100644 --- a/vllm/model_executor/layers/attention/pcp.py +++ b/vllm/model_executor/layers/attention/pcp.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING, NamedTuple, TypeVar +from typing import TYPE_CHECKING, TypeVar import torch @@ -123,102 +123,6 @@ def maybe_gather_indexer_k( return cache_k, cache_slot_mapping -def gather_prefill_qkv_global( - pcp_prefill_gather: tuple[torch.Tensor, torch.Tensor, torch.Tensor, int], - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """All-gather this rank's DualChunkSwap chunk Q/K/V across PCP and reorder - to global position order (real tokens only). - - ``pcp_prefill_gather`` is the (restore_idx, gather_idx, global_cu_seqlens, - padded_n) tuple from ``PCPManager.prefill_gather_indices()``. - """ - restore_idx, _, _, padded_n = pcp_prefill_gather - pcp_group = get_pcp_group() - - def _pad(t: torch.Tensor) -> torch.Tensor: - if t.shape[0] < padded_n: - pad = t.new_zeros((padded_n - t.shape[0],) + tuple(t.shape[1:])) - return torch.cat([t, pad], dim=0) - return t - - q_g = pcp_group.all_gather(_pad(query), dim=0)[restore_idx] - k_g = pcp_group.all_gather(_pad(key), dim=0)[restore_idx] - v_g = pcp_group.all_gather(_pad(value), dim=0)[restore_idx] - return q_g, k_g, v_g - - -def slice_prefill_output_local( - pcp_prefill_gather: tuple[torch.Tensor, torch.Tensor, torch.Tensor, int], - out_g: torch.Tensor, - num_actual: int, -) -> torch.Tensor: - """Slice a global-order output tensor back to this PCP rank's local chunk.""" - _, gather_idx, _, padded_n = pcp_prefill_gather - pcp_rank = get_pcp_group().rank_in_group - out_gathered = out_g[gather_idx] - local = out_gathered[pcp_rank * padded_n : (pcp_rank + 1) * padded_n] - return local[:num_actual] - - -class DecodeSubset(NamedTuple): - """Decode-token subset of a mixed prefill+decode batch.""" - - token_mask: torch.Tensor - q: torch.Tensor - k: torch.Tensor - v: torch.Tensor - cu_seqlens: torch.Tensor - max_seqlen: int - ctx_kv_lens: torch.Tensor - block_table: torch.Tensor - - -def build_mixed_decode_subset( - is_prefilling: torch.Tensor | None, - query_start_loc: torch.Tensor, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - dcp_context_kv_lens: torch.Tensor, - block_table: torch.Tensor, -) -> DecodeSubset | None: - """Extract the decode-token subset of a mixed prefill+decode batch. - - DualChunkSwap gives ranks with no prefill chunks a 0-token dummy decode - segment mirroring global req 0; such segments are excluded (counting them - would desync the DCP LSE-combine). Returns None for a pure-prefill batch. - """ - if is_prefilling is None: - return None - is_pre = is_prefilling.to(device=query_start_loc.device) - seg_lens = query_start_loc[1:] - query_start_loc[:-1] - decode_seg = (~is_pre.bool()) & (seg_lens > 0) - num_decode = int(decode_seg.sum().item()) - if num_decode == 0: - return None - token_mask = torch.repeat_interleave(decode_seg, seg_lens)[: query.shape[0]] - dec_lens = seg_lens[decode_seg] - dec_cu = torch.zeros( - num_decode + 1, - dtype=query_start_loc.dtype, - device=query_start_loc.device, - ) - torch.cumsum(dec_lens, dim=0, out=dec_cu[1:]) - return DecodeSubset( - token_mask=token_mask, - q=query[token_mask].contiguous(), - k=key[token_mask].contiguous(), - v=value[token_mask].contiguous(), - cu_seqlens=dec_cu, - max_seqlen=int(dec_lens.max().item()), - ctx_kv_lens=dcp_context_kv_lens[decode_seg], - block_table=block_table[decode_seg], - ) - - def _dcp_q_gather_group( dcp_world_size: int, pcp_world_size: int, diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 31f31ccab06e..10011ee0371d 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -4,19 +4,19 @@ import copy from dataclasses import dataclass -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar import numpy as np import torch +if TYPE_CHECKING: + from vllm.v1.worker.gpu.pcp_manager import PCPRowPlan + from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention.pcp import ( - build_mixed_decode_subset, dcp_q_gather_size, - gather_prefill_qkv_global, maybe_all_gather_q_for_dcp, resolve_dcp_combine_fn, - slice_prefill_output_local, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import ( @@ -54,7 +54,7 @@ get_layers_from_vllm_config, ) from vllm.config.cache import CacheDType -from vllm.distributed.parallel_state import get_dcp_group +from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group from vllm.forward_context import ( get_forward_context, is_forward_context_available, @@ -278,8 +278,9 @@ class FlashAttentionMetadata: max_dcp_context_kv_len: int | None = None dcp_context_kv_lens: torch.Tensor | None = None - # Per-segment is_prefilling flag (rank-local view). Used by the MRv2 - # sharded PCP+DCP path to extract decode tokens out of a mixed batch. + # Per-segment is_prefilling flag (rank-local view). Used by the builder to + # detect prefill vs decode batches under DualChunkSwap (a prefill chunk can + # be a single token, so max_query_len is unreliable). is_prefilling: torch.Tensor | None = None # Split counts for FA2 DCP context attention. num_prefill_* tracks @@ -596,10 +597,10 @@ def schedule( prefix_scheduler_metadata = None if self.use_pcp and self.dcp_world_size > 1: - # MRv2 PCP+DCP: cache is DCP-sharded (1/dcp KV/rank). DualChunkSwap - # already partitioned the batch per rank; both decode and prefill use - # the context+query+merge path. context_kv_lens = seq_lens - - # query_lens is the chunk-start position (= prior prefix length). + # MRv2 PCP+DCP: cache is DCP-sharded (1/dcp KV/rank). This local + # metadata (context lens of the rank-local DualChunkSwap rows) feeds + # the pure-decode path (_pcp_dcp_decode_rows); prefill/mixed steps + # run on the PCPRowPlan from the forward context instead. scheduler_metadata = None query_lens = query_start_loc[1:] - query_start_loc[:-1] context_kv_lens = seq_lens - query_lens @@ -631,7 +632,7 @@ def schedule( elif self.dcp_world_size > 1: # Pure-DCP path (DCP reuses the TP ranks; no PCP). The PCP+DCP case # is handled by the block above and runs the sharded - # _forward_dcp_mrv2 path, so it never reaches here. + # _forward_pcp_dcp path, so it never reaches here. query_lens = query_start_loc[1:] - query_start_loc[:-1] context_kv_lens = seq_lens - query_lens local_context_kv_lens = get_dcp_local_seq_lens( @@ -900,11 +901,18 @@ def __init__( self.supports_quant_query_input = flash_attn_supports_quant_query_input() vllm_config = get_current_vllm_config_or_none() - # self.pcp_world_size is auto-populated by AttentionImplBase.__new__. + # self.pcp_world_size / self.pcp_rank are auto-populated by + # AttentionImplBase.__new__ from get_pcp_group(). self.use_pcp = self.pcp_world_size > 1 # How Q and the partials move across the DCP group is decided by the # PCP/DCP topology alone; both backends share the rule (see pcp.py). self.dcp_combine = resolve_dcp_combine_fn(vllm_config) + if self.use_pcp and self.dcp_world_size > 1: + assert self.dcp_world_size == self.pcp_world_size, ( + "FlashAttention MRv2 PCP+DCP requires dcp == pcp (the DCP " + "group spans exactly the PCP ranks), got " + f"dcp={self.dcp_world_size}, pcp={self.pcp_world_size}." + ) self._dcp_dtype: torch.dtype | None = None self._dcp_max_num_tokens: int = 0 @@ -914,10 +922,6 @@ def __init__( vllm_config.scheduler_config.max_num_batched_tokens ) - # self.pcp_world_size / self.pcp_rank are auto-populated by - # AttentionImplBase.__new__ from get_pcp_group(). - self.use_pcp = self.pcp_world_size > 1 - def forward( self, layer: torch.nn.Module, @@ -1045,9 +1049,11 @@ def forward( return output elif self.dcp_world_size > 1: # MRv2 PCP+DCP (dcp == pcp): the cache is DCP-sharded - # (1/dcp KV per rank) for the memory win. Decode attends its - # local shard + LSE-combine; prefill all-gathers the full KV. - self._forward_dcp_mrv2( + # (1/dcp KV per rank) for the memory win. Every row is a + # (cached prefix, new-token span) pair: the suffix attends the + # write-gathered new K/V, the prefix attends the local cache + # shard + LSE-combine. + self._forward_pcp_dcp( query[:num_actual_tokens], key[:num_actual_tokens], value[:num_actual_tokens], @@ -1055,6 +1061,7 @@ def forward( value_cache, output[:num_actual_tokens], attn_metadata, + layer, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, @@ -1258,6 +1265,14 @@ def do_kv_cache_update( num_decode_tokens, self.use_pcp, ) + # Stash the write-gathered K/V for the suffix attention of the + # sharded PCP+DCP path (a gather only happens on prefill steps). + # Keyed by layer and scoped to this step's forward context, so a + # layer whose cache update is skipped (kv sharing) cannot pick up + # stale tensors. + get_forward_context().additional_kwargs[ + f"pcp_gathered_kv:{layer.layer_name}" + ] = (cache_key, cache_value) if num_decode_tokens == 0 else None reshape_and_cache_flash( cache_key, cache_value, @@ -1353,7 +1368,7 @@ def _sharded_context_attention( ) return ctx_out_cor, ctx_lse_cor.transpose(0, 1).contiguous() - def _forward_dcp_mrv2( + def _forward_pcp_dcp( self, query: torch.Tensor, key: torch.Tensor, @@ -1362,37 +1377,42 @@ def _forward_dcp_mrv2( value_cache: torch.Tensor, output: torch.Tensor, attn_metadata: FlashAttentionMetadata, + layer: torch.nn.Module, q_descale: torch.Tensor | None = None, k_descale: torch.Tensor | None = None, v_descale: torch.Tensor | None = None, ) -> torch.Tensor: """MRv2 PCP+DCP GQA attention over a sharded (1/dcp) KV cache. - Pure-decode uses ``_forward_dcp_mrv2_decode`` (Q replicated, attend the - sharded prefix + LSE-combine). Prefill/extend/mixed all-gather the - DualChunkSwap-partitioned Q/K/V, run full causal, optionally - context-attend the cached prefix (extend), and slice back; decode tokens - in a mixed batch are then overwritten via the decode path. + Every rank-local row (decode or DualChunkSwap prefill chunk) is a + (cached prefix, new-token span) pair and is evaluated as: + + - suffix: causal attention of the row against the write-gathered new + K/V -- the same all-gather the cache write already performs, so no + extra collective, and every rank computes only its own rows; + - prefix: non-causal attention of the PCP-gathered prefix-row + queries against this rank's DCP cache shard, then an LSE + all-reduce combine across the group (same-queries requirement), + sliced back to this rank's rows; + - merge_attn_states folds the two. + + A pure-decode batch skips the query gather entirely: decode rows are + replicated across ranks, so local metadata suffices. """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." ) fc = get_forward_context() - # Route on rank-invariant global flags, not per-rank num_decode_tokens: + # Route on rank-invariant global flags, not per-rank row composition: # DualChunkSwap can zero some ranks' prefill chunks and desync collectives. - global_has_prefill, global_has_decode = fc.additional_kwargs.get( + global_has_prefill, _ = fc.additional_kwargs.get( "pcp_global_flags", (True, False) ) - - cu_seqlens_q = attn_metadata.query_start_loc - max_seqlen_q = attn_metadata.max_query_len - block_table = attn_metadata.block_table query = query.contiguous() fa_kw = self._fa_common_kwargs(q_descale, k_descale, v_descale) - # Pure-decode batch: Q is replicated, attend the sharded cache + LSE-combine. if not global_has_prefill: - self._forward_dcp_mrv2_decode( + self._pcp_dcp_decode_rows( query, key, value, @@ -1400,134 +1420,110 @@ def _forward_dcp_mrv2( value_cache, output, attn_metadata, - cu_seqlens_q, - max_seqlen_q, - block_table, fa_kw, ) return output - # Prefill/extend/mixed: gather the DualChunkSwap-partitioned Q/K/V across - # PCP, run full causal on the gathered tokens, slice back. Decode tokens - # in a mixed batch self-attend here (no cached context) and are fixed below. - info = fc.additional_kwargs.get("pcp_prefill_gather") - if info is None: + plan: PCPRowPlan | None = fc.additional_kwargs.get("pcp_row_plan") + if plan is None: # Warmup/dummy: no PCP partition -> local causal (shape-correct only). flash_attn_varlen_func( q=query, k=key, v=value, out=output, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - cu_seqlens_k=cu_seqlens_q, - max_seqlen_k=max_seqlen_q, + cu_seqlens_q=attn_metadata.query_start_loc, + max_seqlen_q=attn_metadata.max_query_len, + cu_seqlens_k=attn_metadata.query_start_loc, + max_seqlen_k=attn_metadata.max_query_len, causal=attn_metadata.causal, num_splits=attn_metadata.max_num_splits, **fa_kw, ) return output - global_cu_seqlens = info[2] - num_actual = query.shape[0] - q_g, k_g, v_g = gather_prefill_qkv_global(info, query, key, value) - cu_seqlens = global_cu_seqlens.to(device=query.device, dtype=torch.int32) - global_max_seqlen = int(global_cu_seqlens.max().item()) - out_g = torch.empty_like(q_g) - # Query attention: full causal over the gathered new tokens. - _, query_lse_g = flash_attn_varlen_func( - q=q_g, - k=k_g, - v=v_g, - out=out_g, - cu_seqlens_q=cu_seqlens, - max_seqlen_q=global_max_seqlen, - cu_seqlens_k=cu_seqlens, - max_seqlen_k=global_max_seqlen, - causal=attn_metadata.causal, - return_softmax_lse=True, - **fa_kw, - ) - # EXTEND (prefix caching): q_g also attends the cached prefix, sharded + - # LSE-combined. Skipped for pure first-prefills (rank-invariant has_extend). - final_g = out_g - gctx = fc.additional_kwargs.get("pcp_global_ctx") - if gctx is not None and gctx[3] and key_cache is not None: - g_block_table, g_ctx_kv_lens, g_max_ctx_kv, _ = gctx - ctx_out_cor, ctx_lse_cor = self._sharded_context_attention( - q_g, - key_cache, - value_cache, - cu_seqlens, - global_max_seqlen, - g_ctx_kv_lens, - g_max_ctx_kv, - g_block_table, - fa_kw, + + n = query.shape[0] + suffix_lse = None + if n > 0: + # Suffix: local rows vs the write-gathered new K/V. The gather is + # the one do_kv_cache_update already did for the cache write. + gathered_kv = fc.additional_kwargs.get( + f"pcp_gathered_kv:{layer.layer_name}" ) - final_g = torch.empty_like(out_g) - merge_attn_states(final_g, ctx_out_cor, ctx_lse_cor, out_g, query_lse_g) - # global -> gathered (rank-order) -> this rank's local chunk. - output[:num_actual] = slice_prefill_output_local(info, final_g, num_actual) - - # MIXED batch: the gather above self-attended decode tokens (no cached - # context); recompute them. No-op for a pure-prefill batch. - if global_has_decode: - self._overwrite_mixed_decode_tokens( - query=query, - key=key, - value=value, - key_cache=key_cache, - value_cache=value_cache, - output=output, - attn_metadata=attn_metadata, - fa_kw=fa_kw, + assert gathered_kv is not None, ( + "PCP+DCP prefill step without write-gathered K/V: " + "do_kv_cache_update did not run for this layer (kv sharing?)." ) - return output + k_g, v_g = gathered_kv + _, suffix_lse = flash_attn_varlen_func( + q=query, + k=k_g[plan.suffix_kv_idx], + v=v_g[plan.suffix_kv_idx], + out=output, + cu_seqlens_q=attn_metadata.query_start_loc, + max_seqlen_q=attn_metadata.max_query_len, + cu_seqlens_k=plan.suffix_cu_k, + max_seqlen_k=plan.suffix_max_k, + causal=attn_metadata.causal, + return_softmax_lse=True, + num_splits=attn_metadata.max_num_splits, + **fa_kw, + ) + if not plan.has_prefix: + return output - def _overwrite_mixed_decode_tokens( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - output: torch.Tensor, - attn_metadata: FlashAttentionMetadata, - fa_kw: dict, - ) -> None: - """Recompute decode-token outputs the prefill-gather path self-attended.""" - ctx_kv_lens = attn_metadata.dcp_context_kv_lens - assert ctx_kv_lens is not None - sub = build_mixed_decode_subset( - attn_metadata.is_prefilling, - attn_metadata.query_start_loc, - query, - key, - value, - ctx_kv_lens, - attn_metadata.block_table, + # Prefix: gather the prefix-row queries (extend chunks + decode rows) + # so every rank evaluates the same rows against its DCP cache shard, + # then LSE all-reduce the partials. + num_prefix_rows = plan.prefix_cu_q.shape[0] - 1 + pfx_descale_shape = (num_prefix_rows, self.num_kv_heads) + pfx_fa_kw = self._fa_common_kwargs( + layer._q_scale.expand(pfx_descale_shape) + if self.supports_quant_query_input + else None, + layer._k_scale.expand(pfx_descale_shape), + layer._v_scale.expand(pfx_descale_shape), ) - if sub is None: - return - dec_out = output[sub.token_mask] - self._forward_dcp_mrv2_decode( - sub.q, - sub.k, - sub.v, - key_cache, - value_cache, - dec_out, - attn_metadata, - sub.cu_seqlens, - sub.max_seqlen, - sub.block_table, - fa_kw, - dcp_context_kv_lens=sub.ctx_kv_lens, - max_dcp_context_kv_len=int(sub.ctx_kv_lens.max().item()), + if n > 0: + q_local = query[plan.prefix_q_local_idx] + else: + # This rank holds no rows; still must join the collectives. + q_local = query.new_zeros( + (plan.padded_num_prefix_tokens, self.num_heads, self.head_size) + ) + q_g = get_pcp_group().all_gather(q_local, dim=0)[plan.prefix_q_restore_idx] + ctx_out, ctx_lse = flash_attn_varlen_func( + q=q_g, + k=key_cache, + v=value_cache, + cu_seqlens_q=plan.prefix_cu_q, + max_seqlen_q=plan.prefix_max_q, + seqused_k=plan.prefix_dcp_ctx_lens, + max_seqlen_k=plan.prefix_max_ctx, + causal=False, + block_table=plan.prefix_block_table, + return_softmax_lse=True, + num_splits=attn_metadata.max_num_splits, + **pfx_fa_kw, ) - output[sub.token_mask] = dec_out + ctx_out_cor, ctx_lse_cor = self.dcp_combine( + ctx_out, ctx_lse.transpose(0, 1), get_dcp_group(), return_lse=True + ) + if n == 0: + return output + assert suffix_lse is not None + ctx_lse_cor = ctx_lse_cor.transpose(0, 1).contiguous() + pfx_out = output.new_zeros((n, self.num_heads, self.head_size)) + pfx_lse = suffix_lse.new_full((suffix_lse.shape[0], n), float("-inf")) + pfx_out[plan.prefix_local_token_idx] = ctx_out_cor[plan.prefix_local_out_idx] + pfx_lse[:, plan.prefix_local_token_idx] = ctx_lse_cor[ + :, plan.prefix_local_out_idx + ] + # Rows without a prefix row keep their suffix output (lse -inf -> 0 weight). + merge_attn_states(output, pfx_out, pfx_lse, output, suffix_lse) + return output - def _forward_dcp_mrv2_decode( + def _pcp_dcp_decode_rows( self, query: torch.Tensor, key: torch.Tensor, @@ -1536,23 +1532,16 @@ def _forward_dcp_mrv2_decode( value_cache: torch.Tensor, output: torch.Tensor, attn_metadata: FlashAttentionMetadata, - cu_seqlens_q: torch.Tensor, - max_seqlen_q: int, - block_table: torch.Tensor, fa_kw: dict, - dcp_context_kv_lens: torch.Tensor | None = None, - max_dcp_context_kv_len: int | None = None, ) -> None: - """Decode (or decode-subset) over a sharded cache: attend the local DCP - shard of the prefix + LSE-combine, merge with per-token new K/V. Every - rank runs the combine (no rank-varying guard). Context lens default to - the metadata but can be overridden for a mixed-batch decode subset. - """ - if dcp_context_kv_lens is None: - dcp_context_kv_lens = attn_metadata.dcp_context_kv_lens - if max_dcp_context_kv_len is None: - max_dcp_context_kv_len = attn_metadata.max_dcp_context_kv_len + """Pure-decode batch over the sharded cache: replicated rows attend the + local DCP shard of the prefix + LSE-combine, then merge with per-token + new K/V.""" n = query.shape[0] + if n == 0: + return + assert attn_metadata.dcp_context_kv_lens is not None + assert attn_metadata.max_dcp_context_kv_len is not None dcp_context_out_spec = ( (max(n, self._dcp_max_num_tokens), self.num_heads, self.head_size), self._dcp_dtype, @@ -1560,31 +1549,28 @@ def _forward_dcp_mrv2_decode( (dcp_context_out_workspace,) = current_workspace_manager().get_simultaneous( dcp_context_out_spec, ) - assert dcp_context_kv_lens is not None - assert max_dcp_context_kv_len is not None context_attn_out_cor, context_lse_cor = self._sharded_context_attention( query, key_cache, value_cache, - cu_seqlens_q, - max_seqlen_q, - dcp_context_kv_lens, - max_dcp_context_kv_len, - block_table, + attn_metadata.query_start_loc, + attn_metadata.max_query_len, + attn_metadata.dcp_context_kv_lens, + attn_metadata.max_dcp_context_kv_len, + attn_metadata.block_table, fa_kw, out=dcp_context_out_workspace[:n], num_splits=attn_metadata.max_num_splits, ) - query_attn_out, query_lse = flash_attn_varlen_func( q=query, k=key, v=value, out=output, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - cu_seqlens_k=cu_seqlens_q, - max_seqlen_k=max_seqlen_q, + cu_seqlens_q=attn_metadata.query_start_loc, + max_seqlen_q=attn_metadata.max_query_len, + cu_seqlens_k=attn_metadata.query_start_loc, + max_seqlen_k=attn_metadata.max_query_len, causal=attn_metadata.causal, return_softmax_lse=True, num_splits=attn_metadata.max_num_splits, diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index ac5df1ba373b..c20b8e54294e 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass, replace +from typing import NamedTuple import numpy as np import torch @@ -35,6 +36,36 @@ def num_tokens(self) -> int: return self.global_batch_slice.stop - self.global_batch_slice.start +class PCPRowPlan(NamedTuple): + """Per-step row-level attention plan for the sharded PCP+DCP path. + + Every rank-local row is a (cached prefix, new-token span) pair. The suffix + (new tokens) attends the write-gathered K/V -- the same all-gather the + cache write already performs -- so it needs no collective. The prefix + (cached context) attends this rank's DCP cache shard with the + PCP-gathered prefix-row queries and LSE-combines (all-reduce) across the + group, which requires every rank to hold the same prefix-row queries. + """ + + # Suffix attention (rank-local rows vs the stashed write-gathered K/V). + suffix_kv_idx: torch.Tensor # [total_suffix_k] int64, into stashed K/V + suffix_cu_k: torch.Tensor # [num_local_rows + 1] int32 + suffix_max_k: int + # Prefix attention (gathered rows vs the local DCP cache shard). + has_prefix: bool # rank-invariant: some global row has cached context + num_local_prefix_tokens: int + padded_num_prefix_tokens: int # all-gather slab size (rank-invariant) + prefix_q_local_idx: torch.Tensor # [padded_num_prefix_tokens] int64 + prefix_q_restore_idx: torch.Tensor # [total_global_prefix_tokens] int64 + prefix_cu_q: torch.Tensor # [num_global_prefix_rows + 1] int32 + prefix_max_q: int + prefix_dcp_ctx_lens: torch.Tensor # [num_global_prefix_rows] int32 + prefix_max_ctx: int + prefix_block_table: torch.Tensor # [num_global_prefix_rows, max_num_blocks] + prefix_local_token_idx: torch.Tensor # [num_local_prefix_tokens] int64 + prefix_local_out_idx: torch.Tensor # [num_local_prefix_tokens] int64 + + class PCPManager: """MRV2 PC batch manager. @@ -79,34 +110,35 @@ def __init__( self._global_has_decode: bool = False self._pad_slot_id = torch.tensor(PAD_SLOT_ID, dtype=torch.int64, device=device) - # Global-context attention metadata for the sharded PCP+DCP prefill path - # (extend/prefix-caching support). Built once per step in - # build_global_context_metadata(); consumed by _forward_dcp_mrv2 to run - # context attention on the replicated gathered q_g against this rank's - # DCP shard of the cached prefix. - self._global_ctx_block_tables: tuple[torch.Tensor, ...] | None = ( + # Per-step row-plan state for the sharded PCP+DCP attention path. + # ``_segments_by_rank``/``_hidden_restore_idx_np`` are captured at + # partition time; ``build_cp_row_plan`` derives the plan consumed by + # FlashAttentionImpl._forward_pcp_dcp (see PCPRowPlan). + self._segments_by_rank: list[list[RankSegment]] | None = None + self._hidden_restore_idx_np: np.ndarray | None = None + # Global prefix rows are decode rows plus every extend chunk row, so a + # request contributes up to 2*pcp rows. Size the gather buffers for the + # worst case. + max_num_prefix_rows = ( + 2 * pcp_world_size * max_num_reqs if max_num_reqs is not None else None + ) + self._prefix_block_tables: tuple[torch.Tensor, ...] | None = ( tuple( - table.new_zeros((max_num_reqs, table.shape[1])) + table.new_zeros((max_num_prefix_rows, table.shape[1])) for table in block_tables.input_block_tables ) - if block_tables is not None and max_num_reqs is not None + if block_tables is not None and max_num_prefix_rows is not None else None ) - self._global_ctx_block_table_ptrs: torch.Tensor | None = ( + self._prefix_block_table_ptrs: torch.Tensor | None = ( torch.tensor( - [t.data_ptr() for t in self._global_ctx_block_tables], + [t.data_ptr() for t in self._prefix_block_tables], dtype=torch.uint64, device=device, ) - if self._global_ctx_block_tables is not None + if self._prefix_block_tables is not None else None ) - self._global_ctx_kv_lens: torch.Tensor | None = ( - torch.zeros(max_num_reqs, dtype=torch.int32, device=device) - if max_num_reqs is not None - else None - ) - self._global_num_reqs: int = 0 max_num_local_reqs = 2 * max_num_reqs if max_num_reqs is not None else None self._input_buffers = ( @@ -344,6 +376,7 @@ def _build_batch_layout( dtype=np.int64, ) + self._hidden_restore_idx_np = hidden_restore_idx self._hidden_restore_idx = async_copy_to_gpu( hidden_restore_idx, device=self.device ) @@ -379,6 +412,7 @@ def partition_batch(self, input_batch: InputBatch) -> InputBatch: is_prefilling, global_batch.query_start_loc_np, ) + self._segments_by_rank = segments_by_rank local_segments = segments_by_rank[self.pcp_rank] if not local_segments: @@ -664,108 +698,215 @@ def populate_forward_context(self) -> None: """Stash all PCP metadata the attention forward / kv-cache update need. Centralized here so the model runner stays a single call. Populates: - - ``pcp_prefill_gather``: gather/restore indices for the sharded - prefill-KV all-gather attention (None when no prefill ran). - ``pcp_global_flags``: rank-invariant (has_prefill, has_decode) for - consistent mixed-batch routing. - - ``pcp_global_ctx``: global block table + per-rank DCP shard context - lengths for extend (prefix-caching) support in the prefill path. + consistent mixed-batch routing and cache-write gather decisions. + - ``pcp_row_plan``: per-step :class:`PCPRowPlan` for the sharded + PCP+DCP attention path (None for warmup or pure-decode steps). """ kwargs = get_forward_context().additional_kwargs - prefill_gather = self.prefill_gather_indices() - if prefill_gather is not None: - kwargs["pcp_prefill_gather"] = prefill_gather kwargs["pcp_global_flags"] = self.global_batch_flags() - global_ctx = self.build_global_context_metadata() - if global_ctx is not None: - kwargs["pcp_global_ctx"] = global_ctx + kwargs["pcp_row_plan"] = self.build_cp_row_plan() - def build_global_context_metadata( - self, - ) -> tuple[torch.Tensor, torch.Tensor, int, bool] | None: - """Build per-global-request context-attention metadata for the extend - (prefix-caching) case in the sharded PCP+DCP prefill path. - - Returns (global_block_table, global_dcp_context_kv_lens, - max_global_dcp_context_kv_len, has_extend) for the CURRENT step, aligned - with the pre-partition global batch (so the gathered, replicated q_g can - attend this rank's DCP shard of each request's cached prefix). Returns - None when no prefill was partitioned this step. - - ``global_dcp_context_kv_lens`` is this rank's shard count of the prefix - (rank-specific), indexed by global request -- the LSE-combine then - merges the shards for the shared (replicated) q_g. ``has_extend`` is - rank-invariant (True iff some global request has a cached prefix) so the - prefill path can skip the context collective for pure first-prefills. + def build_cp_row_plan(self) -> "PCPRowPlan | None": + """Build the per-step row plan for the sharded PCP+DCP attention path. + + Returns None when the plan is not needed: no global batch (warmup), + the cache is not DCP-sharded, or the global batch has no prefill + (decode rows are replicated across ranks, so the local-metadata path + handles them without any gather). Derived entirely from partition-time + CPU state; rank-invariant except the DCP-sharded context lengths. """ gb = self._global_batch + segments_by_rank = self._segments_by_rank + hidden_restore_idx = self._hidden_restore_idx_np if ( gb is None - or self._global_ctx_block_tables is None - or self._global_ctx_kv_lens is None - or self._block_tables is None + or segments_by_rank is None + or hidden_restore_idx is None + or self.dcp_world_size <= 1 + or not self._global_has_prefill ): return None - block_tables = self._block_tables - global_ctx_kv_lens = self._global_ctx_kv_lens + assert self._block_tables is not None + assert self._prefix_block_tables is not None + assert self._prefix_block_table_ptrs is not None + + pcp = self.pcp_world_size + num_chunks = 2 * pcp + rank = self.pcp_rank + num_computed = gb.num_computed_tokens_np + is_prefilling = gb.is_prefilling_np + query_start_loc = gb.query_start_loc_np + num_scheduled = gb.num_scheduled_tokens num_reqs = gb.num_reqs - self._global_num_reqs = num_reqs - # Per-global-request FULL prefix length, then this rank's shard count. - qsl = gb.query_start_loc - query_lens = qsl[1 : num_reqs + 1] - qsl[:num_reqs] - context_kv_lens = gb.seq_lens[:num_reqs] - query_lens - has_extend = num_reqs > 0 and bool(context_kv_lens.max().item() > 0) - dcp_ctx_kv_lens = get_dcp_local_seq_lens( - context_kv_lens, + + def chunk_size(j: int) -> int: + return -(-int(num_scheduled[j]) // num_chunks) + + def seg_chunk_idx(seg: RankSegment) -> int: + j = seg.global_batch_req_idx + if not bool(is_prefilling[j]): + return -1 + chunk_offset = seg.global_batch_slice.start - int(query_start_loc[j]) + return chunk_offset // chunk_size(j) + + def is_prefix_seg(seg: RankSegment) -> bool: + # Prefix rows are exactly the rows with a cached context. + return ( + seg.num_tokens > 0 and int(num_computed[seg.global_batch_req_idx]) > 0 + ) + + local_segments = segments_by_rank[rank] or [ + RankSegment(0, slice(0, 0), slice(0, 0)) + ] + + # -- Suffix: per local row, the request's new tokens up to the row end, + # addressed into the stashed write-gathered K/V (padded-slab layout). + suffix_idx_parts = [] + suffix_lens = [] + for seg in local_segments: + j = seg.global_batch_req_idx + req_start = int(query_start_loc[j]) + row_end = seg.global_batch_slice.stop + suffix_lens.append(row_end - req_start) + suffix_idx_parts.append(hidden_restore_idx[req_start:row_end]) + suffix_kv_idx_np = ( + np.concatenate(suffix_idx_parts) + if suffix_idx_parts + else np.empty(0, dtype=np.int64) + ) + suffix_cu_np = np.zeros(len(local_segments) + 1, dtype=np.int32) + np.cumsum(np.asarray(suffix_lens, dtype=np.int32), out=suffix_cu_np[1:]) + suffix_max_k = max(suffix_lens, default=0) + + # -- Prefix rows: slab offsets of every rank's gather contribution. + # Each rank gathers its prefix-row tokens in local-row order; slab r + # starts at r * padded_num_prefix_tokens in the gathered buffer. + slab_offsets: list[dict[tuple[int, int], int]] = [] + per_rank_prefix_tokens = [] + for r in range(pcp): + offsets = {} + offset = 0 + for seg in segments_by_rank[r]: + if not is_prefix_seg(seg): + continue + offsets[(seg.global_batch_req_idx, seg_chunk_idx(seg))] = offset + offset += seg.num_tokens + slab_offsets.append(offsets) + per_rank_prefix_tokens.append(offset) + padded_num_prefix = max(per_rank_prefix_tokens, default=0) + + # Canonical global prefix-row order: decode rows (request order), then + # extend chunk rows (request order, chunk index order). + global_prefix_rows: list[tuple[int, int, int]] = [] # (req, chunk, len) + for j in range(num_reqs): + if not bool(is_prefilling[j]) and int(num_computed[j]) > 0: + global_prefix_rows.append((j, -1, int(num_scheduled[j]))) + for j in range(num_reqs): + if bool(is_prefilling[j]) and int(num_computed[j]) > 0: + cs = chunk_size(j) + for c in range(num_chunks): + chunk_len = min(cs, int(num_scheduled[j]) - c * cs) + if chunk_len > 0: + global_prefix_rows.append((j, c, chunk_len)) + has_prefix = len(global_prefix_rows) > 0 + + restore_parts = [] + prefix_cu_np = np.zeros(len(global_prefix_rows) + 1, dtype=np.int32) + row_req_np = np.zeros(len(global_prefix_rows), dtype=np.int64) + row_ctx_np = np.zeros(len(global_prefix_rows), dtype=np.int32) + for i, (j, c, row_len) in enumerate(global_prefix_rows): + # Chunk c lives on the rank r with r == c or r == num_chunks-1-c. + holding_rank = 0 if c < 0 else min(c, num_chunks - 1 - c) + slab_start = slab_offsets[holding_rank][(j, c)] + flat_start = holding_rank * padded_num_prefix + slab_start + restore_parts.append( + np.arange(flat_start, flat_start + row_len, dtype=np.int64) + ) + prefix_cu_np[i + 1] = prefix_cu_np[i] + row_len + row_req_np[i] = j + row_ctx_np[i] = num_computed[j] + prefix_q_restore_np = ( + np.concatenate(restore_parts) + if restore_parts + else np.empty(0, dtype=np.int64) + ) + prefix_max_q = max((row_len for _, _, row_len in global_prefix_rows), default=0) + dcp_ctx_np = get_dcp_local_seq_lens( + torch.from_numpy(row_ctx_np), self.dcp_world_size, self.dcp_rank, self.cp_interleave, - ) - global_ctx_kv_lens[:num_reqs] = dcp_ctx_kv_lens - # Gather this rank's block-table rows for the global requests. - global_bt = block_tables.gather_block_tables( - gb.idx_mapping, - num_reqs, - out=self._global_ctx_block_tables, - out_ptrs=self._global_ctx_block_table_ptrs, - ) - num_partitions = self.dcp_world_size * self.cp_interleave - max_ctx_kv = ( - ((int(context_kv_lens.max().item()) + num_partitions - 1) // num_partitions) - * self.cp_interleave - if num_reqs > 0 - else 0 - ) - return global_bt[0], global_ctx_kv_lens[:num_reqs], max_ctx_kv, has_extend - - def prefill_gather_indices( - self, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int] | None: - """Indices for the SHARDED-cache prefill KV-gather attention. - - Returns (restore_idx, gather_idx, global_cu_seqlens, padded_num_tokens) - for the current prefill step, or None for decode / when PCP did not - partition a prefill. ``restore_idx`` maps a global-position index to its - slot in the PCP all-gathered tensor (``global = gathered[restore_idx]``); - ``gather_idx`` is the inverse (``gathered = global[gather_idx]``); - ``global_cu_seqlens`` is the pre-partition (global) batch's per-request - cumulative query lengths. - """ - if self._hidden_restore_idx is None or self._global_batch is None: - return None - if self._padded_gather_idx is None: - return None - global_cu_seqlens = self._global_batch.query_start_loc - if not isinstance(global_cu_seqlens, torch.Tensor): - return None - # _padded_gather_idx has length padded_num_tokens * pcp_world_size for - # the current step (set in _build_batch_layout). - padded_num_tokens = self._padded_gather_idx.shape[0] // self.pcp_world_size - return ( - self._hidden_restore_idx, - self._padded_gather_idx, - global_cu_seqlens, - padded_num_tokens, + ).numpy() + prefix_max_ctx = int(dcp_ctx_np.max()) if len(dcp_ctx_np) else 0 + + # Local prefix rows: where their tokens sit in the local batch and in + # the combined (global prefix-token order) prefix attention output. + row_out_start = { + (j, c): int(prefix_cu_np[i]) + for i, (j, c, _) in enumerate(global_prefix_rows) + } + local_token_parts = [] + local_out_parts = [] + for seg in local_segments: + if not is_prefix_seg(seg): + continue + out_start = row_out_start[(seg.global_batch_req_idx, seg_chunk_idx(seg))] + local_token_parts.append( + np.arange( + seg.rank_local_batch_slice.start, + seg.rank_local_batch_slice.stop, + dtype=np.int64, + ) + ) + local_out_parts.append( + np.arange(out_start, out_start + seg.num_tokens, dtype=np.int64) + ) + prefix_local_token_np = ( + np.concatenate(local_token_parts) + if local_token_parts + else np.empty(0, dtype=np.int64) + ) + prefix_local_out_np = ( + np.concatenate(local_out_parts) + if local_out_parts + else np.empty(0, dtype=np.int64) + ) + num_local_prefix = len(prefix_local_token_np) + prefix_q_local_np = np.zeros(padded_num_prefix, dtype=np.int64) + prefix_q_local_np[:num_local_prefix] = prefix_local_token_np + + row_req_gpu = async_copy_to_gpu(row_req_np, device=self.device) + num_prefix_rows = len(global_prefix_rows) + prefix_block_table = self._block_tables.gather_block_tables( + torch.index_select(gb.idx_mapping, 0, row_req_gpu), + num_prefix_rows, + out=self._prefix_block_tables, + out_ptrs=self._prefix_block_table_ptrs, + )[0] + + return PCPRowPlan( + suffix_kv_idx=async_copy_to_gpu(suffix_kv_idx_np, device=self.device), + suffix_cu_k=async_copy_to_gpu(suffix_cu_np, device=self.device), + suffix_max_k=suffix_max_k, + has_prefix=has_prefix, + num_local_prefix_tokens=num_local_prefix, + padded_num_prefix_tokens=padded_num_prefix, + prefix_q_local_idx=async_copy_to_gpu(prefix_q_local_np, device=self.device), + prefix_q_restore_idx=async_copy_to_gpu( + prefix_q_restore_np, device=self.device + ), + prefix_cu_q=async_copy_to_gpu(prefix_cu_np, device=self.device), + prefix_max_q=prefix_max_q, + prefix_dcp_ctx_lens=async_copy_to_gpu(dcp_ctx_np, device=self.device), + prefix_max_ctx=prefix_max_ctx, + prefix_block_table=prefix_block_table, + prefix_local_token_idx=async_copy_to_gpu( + prefix_local_token_np, device=self.device + ), + prefix_local_out_idx=async_copy_to_gpu( + prefix_local_out_np, device=self.device + ), ) def restore_for_sampling( From 1e9b233e5086e77efaa1723b7580b12af7f767a2 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Thu, 30 Jul 2026 06:37:56 +0000 Subject: [PATCH 12/18] clean Signed-off-by: JaredforReal --- .../layers/attention/mla_attention.py | 24 +- vllm/model_executor/layers/attention/pcp.py | 67 +-- .../layers/sparse_attn_indexer.py | 6 +- vllm/v1/attention/backends/flash_attn.py | 405 +++++++----------- 4 files changed, 181 insertions(+), 321 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index ebe490e5537f..415f41274fd1 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -228,7 +228,7 @@ from vllm.model_executor.layers.attention.pcp import ( cp_reconcile_heads, maybe_all_gather_q_for_dcp, - maybe_gather_mla_latent_cache_inputs, + maybe_gather_cache_inputs, resolve_dcp_combine_fn, ) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -342,6 +342,24 @@ def _canonicalize_sparse_mla_kv_cache_dtype( return kv_cache_dtype +def _gather_latent_cache_inputs( + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + slot_mapping: torch.Tensor | None, + num_decode_tokens: int | None, + use_pcp: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """PCP latent-cache write gather; flattens k_pe for the all-gather.""" + if not use_pcp or num_decode_tokens is None: + return kv_c_normed, k_pe, slot_mapping + assert slot_mapping is not None + k_pe_flat = k_pe.reshape(kv_c_normed.shape[0], -1) + (cache_kv_c, cache_k_pe_flat), cache_slot_mapping = maybe_gather_cache_inputs( + (kv_c_normed, k_pe_flat), slot_mapping, num_decode_tokens, use_pcp + ) + return cache_kv_c, cache_k_pe_flat.view(-1, *k_pe.shape[1:]), cache_slot_mapping + + class MLAAttention(nn.Module, AttentionLayerBase): """Multi-Head Latent Attention layer. @@ -625,7 +643,7 @@ def forward( ) layer_slot_mapping = slot_mapping.get(self.layer_name) kv_for_cache, kpe_for_cache, layer_slot_mapping = ( - maybe_gather_mla_latent_cache_inputs( + _gather_latent_cache_inputs( kv_c_normed, k_pe, layer_slot_mapping, @@ -1142,7 +1160,7 @@ def unified_mla_kv_cache_update( layer_name ) if layer_slot_mapping is not None: - kv_c_normed, k_pe, layer_slot_mapping = maybe_gather_mla_latent_cache_inputs( + kv_c_normed, k_pe, layer_slot_mapping = _gather_latent_cache_inputs( kv_c_normed, k_pe, layer_slot_mapping, diff --git a/vllm/model_executor/layers/attention/pcp.py b/vllm/model_executor/layers/attention/pcp.py index a4bc4ec6c4b1..467b20f6097b 100644 --- a/vllm/model_executor/layers/attention/pcp.py +++ b/vllm/model_executor/layers/attention/pcp.py @@ -58,69 +58,28 @@ def _gather_prefill_cache_inputs( return cache_inputs, cache_slot_mapping -def maybe_gather_mla_latent_cache_inputs( - kv_c_normed: torch.Tensor, - k_pe: torch.Tensor, - slot_mapping: torch.Tensor | None, - num_decode_tokens: int | None, - use_pcp: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - if not use_pcp or num_decode_tokens is None: - return kv_c_normed, k_pe, slot_mapping - assert slot_mapping is not None - num_tokens = kv_c_normed.shape[0] - k_pe_flat = k_pe.reshape(num_tokens, -1) - (cache_kv_c, cache_k_pe_flat), cache_slot_mapping = _gather_prefill_cache_inputs( - (kv_c_normed, k_pe_flat), - slot_mapping, - num_decode_tokens, - ) - cache_k_pe = cache_k_pe_flat.view(-1, *k_pe.shape[1:]) - return cache_kv_c, cache_k_pe, cache_slot_mapping - - -def maybe_gather_kv_cache_inputs( - key: torch.Tensor, - value: torch.Tensor, +def maybe_gather_cache_inputs( + tensors: tuple[torch.Tensor, ...], slot_mapping: torch.Tensor | None, num_decode_tokens: int | None, use_pcp: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """GQA/MHA PCP KV-cache gather. +) -> tuple[tuple[torch.Tensor, ...], torch.Tensor | None]: + """PCP cache-write gather. - All-gather the prefill portion of K/V across PCP ranks so every rank can - write the full prefill KV to its (replicated, ``dcp=1``) cache, while - keeping decode writes local. Returns contiguous K, V and the gathered - cache slot mapping ready for ``reshape_and_cache_flash``. No-op when PCP - is off. + All-gather the prefill portion of ``tensors`` (K/V for GQA, latent KV and + k_pe for MLA, k for the sparse indexer) across PCP ranks so every rank can + write the full prefill contents to its cache, while keeping decode writes + local. Returns the gathered tensors and the gathered cache slot mapping. + No-op when PCP is off. - Gathering K and V separately (rather than ``cat``-then-``split``) keeps - each output contiguous, so the cache kernel's ``head stride == head_size`` + Gathering each tensor separately (rather than ``cat``-then-``split``) + keeps each output contiguous, so the cache kernel's head-stride assumption holds. """ if not use_pcp or num_decode_tokens is None: - return key, value, slot_mapping + return tensors, slot_mapping assert slot_mapping is not None - (cache_key, cache_value), cache_slot_mapping = _gather_prefill_cache_inputs( - (key, value), - slot_mapping, - num_decode_tokens, - ) - return cache_key, cache_value, cache_slot_mapping - - -def maybe_gather_indexer_k( - k: torch.Tensor, - slot_mapping: torch.Tensor, - num_decode_tokens: int, - use_pcp: bool, -) -> tuple[torch.Tensor, torch.Tensor]: - if not use_pcp: - return k, slot_mapping - (cache_k,), cache_slot_mapping = _gather_prefill_cache_inputs( - (k,), slot_mapping, num_decode_tokens - ) - return cache_k, cache_slot_mapping + return _gather_prefill_cache_inputs(tensors, slot_mapping, num_decode_tokens) def _dcp_q_gather_group( diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 9a671be55639..fa493f3e217a 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -13,7 +13,7 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp -from vllm.model_executor.layers.attention.pcp import maybe_gather_indexer_k +from vllm.model_executor.layers.attention.pcp import maybe_gather_cache_inputs from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) @@ -388,8 +388,8 @@ def sparse_attn_indexer( if not skip_k_cache_insert: assert k is not None - k, slot_mapping_for_cache = maybe_gather_indexer_k( - k, + (k,), slot_mapping_for_cache = maybe_gather_cache_inputs( + (k,), slot_mapping, num_decode_tokens, use_pcp, diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 10011ee0371d..f97b7dedaf84 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -16,6 +16,7 @@ from vllm.model_executor.layers.attention.pcp import ( dcp_q_gather_size, maybe_all_gather_q_for_dcp, + maybe_gather_cache_inputs, resolve_dcp_combine_fn, ) from vllm.platforms import current_platform @@ -489,6 +490,37 @@ def __init__( [self.rswa_window], dtype=torch.int32, device=self.device ) + def _build_dcp_context_lens( + self, + num_reqs: int, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + max_seq_len: int, + ) -> tuple[torch.Tensor, int]: + """This rank's DCP shard of each row's cached-context length. + + The cached context (seq_lens - query_lens) is interleave-sharded + across DCP ranks; each rank attends only its shard. The returned max + is the aligned per-rank upper bound, ceil(L / (N * I)) * I with + L = max_seq_len, N = dcp_world_size, I = interleave size, which avoids + a GPU->CPU sync while minimizing workspace over-allocation. + """ + query_lens = query_start_loc[1:] - query_start_loc[:-1] + context_kv_lens = seq_lens - query_lens + local_context_kv_lens = get_dcp_local_seq_lens( + context_kv_lens, + self.dcp_world_size, + self.dcp_rank, + self.cp_kv_cache_interleave_size, + ) + self._dcp_context_kv_lens[:num_reqs] = local_context_kv_lens + self._dcp_context_kv_lens[num_reqs:] = 0 + num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size + max_dcp_context_kv_len = ( + (max_seq_len + num_partitions - 1) // num_partitions + ) * self.cp_kv_cache_interleave_size + return self._dcp_context_kv_lens[:num_reqs], max_dcp_context_kv_len + def build( self, common_prefix_len: int, @@ -596,103 +628,68 @@ def schedule( suffix_kv_lens = None prefix_scheduler_metadata = None - if self.use_pcp and self.dcp_world_size > 1: - # MRv2 PCP+DCP: cache is DCP-sharded (1/dcp KV/rank). This local - # metadata (context lens of the rank-local DualChunkSwap rows) feeds - # the pure-decode path (_pcp_dcp_decode_rows); prefill/mixed steps - # run on the PCPRowPlan from the forward context instead. - scheduler_metadata = None - query_lens = query_start_loc[1:] - query_start_loc[:-1] - context_kv_lens = seq_lens - query_lens - local_context_kv_lens = get_dcp_local_seq_lens( - context_kv_lens, - self.dcp_world_size, - self.dcp_rank, - self.cp_kv_cache_interleave_size, + if self.dcp_world_size > 1: + # Sharded (1/dcp) KV cache: every row's cached context + # (seq_lens - query_lens) is attended via this rank's interleaved + # shard. Pure-DCP batches and PCP+DCP decode batches both use + # this metadata; PCP+DCP prefill/mixed steps run on the + # PCPRowPlan from the forward context instead. + dcp_context_kv_lens, max_dcp_context_kv_len = self._build_dcp_context_lens( + num_reqs, query_start_loc, seq_lens, max_seq_len ) - self._dcp_context_kv_lens[:num_reqs] = local_context_kv_lens - self._dcp_context_kv_lens[num_reqs:] = 0 - dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs] - num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size - max_dcp_context_kv_len = ( - (max_seq_len + num_partitions - 1) // num_partitions - ) * self.cp_kv_cache_interleave_size - # Detect prefill vs decode. Under DualChunkSwap a prefill chunk can be - # as small as 1 token, so max_query_len is unreliable -- use the - # per-segment is_prefilling flag instead. - is_prefilling = common_attn_metadata.is_prefilling - if is_prefilling is not None: - is_prefill_batch = bool(is_prefilling.any().item()) - else: - is_prefill_batch = max_query_len > 1 - if is_prefill_batch: - num_prefill_tokens = num_actual_tokens + scheduler_metadata = None + if self.use_pcp: + # Detect prefill vs decode for the cache-write gather. Under + # DualChunkSwap a prefill chunk can be as small as 1 token, so + # max_query_len is unreliable -- use the is_prefilling flag. + is_prefilling = common_attn_metadata.is_prefilling + if is_prefilling is not None: + is_prefill_batch = bool(is_prefilling.any().item()) + else: + is_prefill_batch = max_query_len > 1 + if is_prefill_batch: + num_prefill_tokens = num_actual_tokens + else: + num_decode_tokens = num_actual_tokens else: - num_decode_tokens = num_actual_tokens - elif self.dcp_world_size > 1: - # Pure-DCP path (DCP reuses the TP ranks; no PCP). The PCP+DCP case - # is handled by the block above and runs the sharded - # _forward_pcp_dcp path, so it never reaches here. - query_lens = query_start_loc[1:] - query_start_loc[:-1] - context_kv_lens = seq_lens - query_lens - local_context_kv_lens = get_dcp_local_seq_lens( - context_kv_lens, - self.dcp_world_size, - self.dcp_rank, - self.cp_kv_cache_interleave_size, - ) - self._dcp_context_kv_lens[:num_reqs] = local_context_kv_lens - self._dcp_context_kv_lens[num_reqs:] = 0 - dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs] - - skip_dcp_context_attention = False - if common_attn_metadata.seq_lens_cpu_upper_bound is not None: - query_lens_cpu = ( - common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1] - - common_attn_metadata.query_start_loc_cpu[:num_reqs] - ) - context_kv_lens_cpu = ( - common_attn_metadata.seq_lens_cpu_upper_bound[:num_reqs] - - query_lens_cpu - ) - skip_dcp_context_attention = should_skip_dcp_context_attention( - context_kv_lens_cpu - ) + skip_dcp_context_attention = False + if common_attn_metadata.seq_lens_cpu_upper_bound is not None: + query_lens_cpu = ( + common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1] + - common_attn_metadata.query_start_loc_cpu[:num_reqs] + ) + context_kv_lens_cpu = ( + common_attn_metadata.seq_lens_cpu_upper_bound[:num_reqs] + - query_lens_cpu + ) + skip_dcp_context_attention = should_skip_dcp_context_attention( + context_kv_lens_cpu + ) - if max_query_len > 1: - ( - num_decode_reqs, - num_prefill_reqs, - num_decode_tokens, - num_prefill_tokens, - ) = split_dcp_context_queries( - common_attn_metadata.query_start_loc_cpu, - common_attn_metadata.seq_lens_cpu_upper_bound, - max_query_len, - num_actual_tokens, - ) + if max_query_len > 1: + ( + num_decode_reqs, + num_prefill_reqs, + num_decode_tokens, + num_prefill_tokens, + ) = split_dcp_context_queries( + common_attn_metadata.query_start_loc_cpu, + common_attn_metadata.seq_lens_cpu_upper_bound, + max_query_len, + num_actual_tokens, + ) - # After DCP distribution, the maximum number of tokens for any rank is - # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size, - # and I is cp_kv_cache_interleave_size. - # This eliminates GPU->CPU sync while minimizing workspace over-allocation. - if skip_dcp_context_attention: - max_dcp_context_kv_len = 0 - scheduler_metadata = None - else: - num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size - max_dcp_context_kv_len = ( - (max_seq_len + num_partitions - 1) // num_partitions - ) * self.cp_kv_cache_interleave_size - - scheduler_metadata = schedule( - batch_size=num_reqs, - cu_query_lens=query_start_loc, - max_query_len=max_query_len, - seqlens=dcp_context_kv_lens, - max_seq_len=max_dcp_context_kv_len, - causal=False, - ) + if skip_dcp_context_attention: + max_dcp_context_kv_len = 0 + else: + scheduler_metadata = schedule( + batch_size=num_reqs, + cu_query_lens=query_start_loc, + max_query_len=max_query_len, + seqlens=dcp_context_kv_lens, + max_seq_len=max_dcp_context_kv_len, + causal=False, + ) elif use_cascade: cu_prefix_query_lens = torch.tensor( [0, num_actual_tokens], dtype=torch.int32, device=self.device @@ -1031,41 +1028,48 @@ def forward( k_descale = layer._k_scale.expand(descale_shape) v_descale = layer._v_scale.expand(descale_shape) - if self.dcp_world_size > 1 and not self.use_pcp: - # Pure DCP (DCP reuses the TP ranks; no PCP). MRv1-era sharded - # path -- the cache is DCP-sharded and writes are rank-local. - self._forward_with_dcp( - query[:num_actual_tokens], - key[:num_actual_tokens], - value[:num_actual_tokens], - key_cache, - value_cache, - output[:num_actual_tokens], - attn_metadata, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - ) - return output - elif self.dcp_world_size > 1: - # MRv2 PCP+DCP (dcp == pcp): the cache is DCP-sharded - # (1/dcp KV per rank) for the memory win. Every row is a - # (cached prefix, new-token span) pair: the suffix attends the - # write-gathered new K/V, the prefix attends the local cache - # shard + LSE-combine. - self._forward_pcp_dcp( - query[:num_actual_tokens], - key[:num_actual_tokens], - value[:num_actual_tokens], - key_cache, - value_cache, - output[:num_actual_tokens], - attn_metadata, - layer, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - ) + if self.dcp_world_size > 1: + # Sharded (1/dcp) KV cache. One path covers pure DCP and + # PCP+DCP decode batches (rows replicated across ranks, so the + # topology helpers make the Q gather a no-op and the combine an + # all-reduce); only a PCP+DCP prefill/mixed step -- where + # DualChunkSwap partitions the queries -- needs the row plan. + pcp_prefill_step = False + if self.use_pcp: + pcp_prefill_step, _ = get_forward_context().additional_kwargs.get( + "pcp_global_flags", (True, False) + ) + if pcp_prefill_step: + # MRv2 PCP+DCP prefill (dcp == pcp): every row is a + # (cached prefix, new-token span) pair -- the suffix + # attends the write-gathered new K/V, the prefix attends + # the local cache shard + LSE-combine. + self._forward_pcp_dcp( + query[:num_actual_tokens], + key[:num_actual_tokens], + value[:num_actual_tokens], + key_cache, + value_cache, + output[:num_actual_tokens], + attn_metadata, + layer, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + else: + self._forward_with_dcp( + query[:num_actual_tokens], + key[:num_actual_tokens], + value[:num_actual_tokens], + key_cache, + value_cache, + output[:num_actual_tokens], + attn_metadata, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) return output else: # dcp_world_size <= 1 (no DCP): pure PCP or unparallelized -- @@ -1241,10 +1245,6 @@ def do_kv_cache_update( # All-gather the prefill K/V across PCP ranks (decode writes stay # local) so every rank's cache gets the full prefill KV. K/V are # gathered separately to stay contiguous (cache kernel head-stride). - from vllm.model_executor.layers.attention.pcp import ( - maybe_gather_kv_cache_inputs, - ) - attn_metadata = self._get_attn_metadata_for_layer(layer) num_decode_tokens = ( attn_metadata.num_decode_tokens if attn_metadata is not None else 0 @@ -1258,9 +1258,8 @@ def do_kv_cache_update( key_cache, value_cache = kv_cache.transpose(1, 2).split( self.head_size, dim=-1 ) - cache_key, cache_value, cache_slot_mapping = maybe_gather_kv_cache_inputs( - key, - value, + (cache_key, cache_value), cache_slot_mapping = maybe_gather_cache_inputs( + (key, value), slot_mapping, num_decode_tokens, self.use_pcp, @@ -1326,48 +1325,6 @@ def _fa_common_kwargs( v_descale=v_descale, ) - def _sharded_context_attention( - self, - q: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - cu_seqlens_q: torch.Tensor, - max_seqlen_q: int, - ctx_kv_lens: torch.Tensor, - max_ctx_kv: int, - block_table: torch.Tensor, - fa_kw: dict, - out: torch.Tensor | None = None, - num_splits: int | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Attend ``q`` to the sharded cached prefix, then LSE-combine across DCP. - - Returns the combined output and corrected LSE ([B, H] layout). ``out`` - and ``num_splits`` are passed through only when given. - """ - kw = dict( - q=q, - k=key_cache, - v=value_cache, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - seqused_k=ctx_kv_lens, - max_seqlen_k=max_ctx_kv, - causal=False, - block_table=block_table, - return_softmax_lse=True, - **fa_kw, - ) - if out is not None: - kw["out"] = out - if num_splits is not None: - kw["num_splits"] = num_splits - ctx_out, ctx_lse = flash_attn_varlen_func(**kw) - ctx_out_cor, ctx_lse_cor = self.dcp_combine( - ctx_out, ctx_lse.transpose(0, 1), get_dcp_group(), return_lse=True - ) - return ctx_out_cor, ctx_lse_cor.transpose(0, 1).contiguous() - def _forward_pcp_dcp( self, query: torch.Tensor, @@ -1382,7 +1339,7 @@ def _forward_pcp_dcp( k_descale: torch.Tensor | None = None, v_descale: torch.Tensor | None = None, ) -> torch.Tensor: - """MRv2 PCP+DCP GQA attention over a sharded (1/dcp) KV cache. + """MRv2 PCP+DCP GQA prefill/mixed attention over a sharded KV cache. Every rank-local row (decode or DualChunkSwap prefill chunk) is a (cached prefix, new-token span) pair and is evaluated as: @@ -1396,34 +1353,16 @@ def _forward_pcp_dcp( sliced back to this rank's rows; - merge_attn_states folds the two. - A pure-decode batch skips the query gather entirely: decode rows are - replicated across ranks, so local metadata suffices. + Pure-decode batches do not reach here: their rows are replicated + across ranks, so forward() routes them to _forward_with_dcp. """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." ) fc = get_forward_context() - # Route on rank-invariant global flags, not per-rank row composition: - # DualChunkSwap can zero some ranks' prefill chunks and desync collectives. - global_has_prefill, _ = fc.additional_kwargs.get( - "pcp_global_flags", (True, False) - ) query = query.contiguous() fa_kw = self._fa_common_kwargs(q_descale, k_descale, v_descale) - if not global_has_prefill: - self._pcp_dcp_decode_rows( - query, - key, - value, - key_cache, - value_cache, - output, - attn_metadata, - fa_kw, - ) - return output - plan: PCPRowPlan | None = fc.additional_kwargs.get("pcp_row_plan") if plan is None: # Warmup/dummy: no PCP partition -> local causal (shape-correct only). @@ -1523,69 +1462,6 @@ def _forward_pcp_dcp( merge_attn_states(output, pfx_out, pfx_lse, output, suffix_lse) return output - def _pcp_dcp_decode_rows( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - output: torch.Tensor, - attn_metadata: FlashAttentionMetadata, - fa_kw: dict, - ) -> None: - """Pure-decode batch over the sharded cache: replicated rows attend the - local DCP shard of the prefix + LSE-combine, then merge with per-token - new K/V.""" - n = query.shape[0] - if n == 0: - return - assert attn_metadata.dcp_context_kv_lens is not None - assert attn_metadata.max_dcp_context_kv_len is not None - dcp_context_out_spec = ( - (max(n, self._dcp_max_num_tokens), self.num_heads, self.head_size), - self._dcp_dtype, - ) - (dcp_context_out_workspace,) = current_workspace_manager().get_simultaneous( - dcp_context_out_spec, - ) - context_attn_out_cor, context_lse_cor = self._sharded_context_attention( - query, - key_cache, - value_cache, - attn_metadata.query_start_loc, - attn_metadata.max_query_len, - attn_metadata.dcp_context_kv_lens, - attn_metadata.max_dcp_context_kv_len, - attn_metadata.block_table, - fa_kw, - out=dcp_context_out_workspace[:n], - num_splits=attn_metadata.max_num_splits, - ) - query_attn_out, query_lse = flash_attn_varlen_func( - q=query, - k=key, - v=value, - out=output, - cu_seqlens_q=attn_metadata.query_start_loc, - max_seqlen_q=attn_metadata.max_query_len, - cu_seqlens_k=attn_metadata.query_start_loc, - max_seqlen_k=attn_metadata.max_query_len, - causal=attn_metadata.causal, - return_softmax_lse=True, - num_splits=attn_metadata.max_num_splits, - **fa_kw, - ) - assert context_attn_out_cor.shape == query_attn_out.shape - assert context_lse_cor.shape == query_lse.shape - merge_attn_states( - output, - context_attn_out_cor, - context_lse_cor, - query_attn_out, - query_lse, - ) - def _forward_with_dcp( self, query: torch.Tensor, @@ -1599,6 +1475,13 @@ def _forward_with_dcp( k_descale: torch.Tensor | None = None, v_descale: torch.Tensor | None = None, ) -> torch.Tensor: + """Sharded-cache attention for pure DCP and PCP+DCP decode batches. + + Attend the DCP-sharded cached context with the topology-appropriate + query gather (no-op when DCP spans exactly the PCP ranks, i.e. decode + rows are replicated), LSE-combine across the group, then merge with + causal attention over the new tokens. + """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." ) From 20f3d1ac17057f53fd9ea858735be26d1f840ef8 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Thu, 30 Jul 2026 08:41:29 +0000 Subject: [PATCH 13/18] split PCP Row Plan Signed-off-by: JaredforReal --- vllm/v1/attention/backends/flash_attn.py | 69 ++++------ vllm/v1/worker/gpu/pcp_manager.py | 164 ++++++++++++++--------- 2 files changed, 124 insertions(+), 109 deletions(-) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index f97b7dedaf84..d7f72a5a499f 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -1033,13 +1033,14 @@ def forward( # PCP+DCP decode batches (rows replicated across ranks, so the # topology helpers make the Q gather a no-op and the combine an # all-reduce); only a PCP+DCP prefill/mixed step -- where - # DualChunkSwap partitions the queries -- needs the row plan. - pcp_prefill_step = False - if self.use_pcp: - pcp_prefill_step, _ = get_forward_context().additional_kwargs.get( - "pcp_global_flags", (True, False) - ) - if pcp_prefill_step: + # DualChunkSwap partitions the queries and the manager builds + # a row plan -- goes to _forward_pcp_dcp. + plan = ( + get_forward_context().additional_kwargs.get("pcp_row_plan") + if self.use_pcp + else None + ) + if plan is not None: # MRv2 PCP+DCP prefill (dcp == pcp): every row is a # (cached prefix, new-token span) pair -- the suffix # attends the write-gathered new K/V, the prefix attends @@ -1053,6 +1054,7 @@ def forward( output[:num_actual_tokens], attn_metadata, layer, + plan, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, @@ -1252,8 +1254,7 @@ def do_kv_cache_update( # Rank-invariant gather decision: per-rank num_decode_tokens can differ # under DualChunkSwap and desync the all-gather, so if the global batch # has any prefill every rank gathers the whole batch. - gflags = get_forward_context().additional_kwargs.get("pcp_global_flags") - if gflags is not None and gflags[0]: + if get_forward_context().additional_kwargs.get("pcp_has_prefill"): num_decode_tokens = 0 key_cache, value_cache = kv_cache.transpose(1, 2).split( self.head_size, dim=-1 @@ -1335,6 +1336,7 @@ def _forward_pcp_dcp( output: torch.Tensor, attn_metadata: FlashAttentionMetadata, layer: torch.nn.Module, + plan: "PCPRowPlan", q_descale: torch.Tensor | None = None, k_descale: torch.Tensor | None = None, v_descale: torch.Tensor | None = None, @@ -1353,8 +1355,8 @@ def _forward_pcp_dcp( sliced back to this rank's rows; - merge_attn_states folds the two. - Pure-decode batches do not reach here: their rows are replicated - across ranks, so forward() routes them to _forward_with_dcp. + Pure-decode batches do not reach here: no prefill means no row plan, + and forward() routes those steps to _forward_with_dcp. """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." @@ -1363,24 +1365,6 @@ def _forward_pcp_dcp( query = query.contiguous() fa_kw = self._fa_common_kwargs(q_descale, k_descale, v_descale) - plan: PCPRowPlan | None = fc.additional_kwargs.get("pcp_row_plan") - if plan is None: - # Warmup/dummy: no PCP partition -> local causal (shape-correct only). - flash_attn_varlen_func( - q=query, - k=key, - v=value, - out=output, - cu_seqlens_q=attn_metadata.query_start_loc, - max_seqlen_q=attn_metadata.max_query_len, - cu_seqlens_k=attn_metadata.query_start_loc, - max_seqlen_k=attn_metadata.max_query_len, - causal=attn_metadata.causal, - num_splits=attn_metadata.max_num_splits, - **fa_kw, - ) - return output - n = query.shape[0] suffix_lse = None if n > 0: @@ -1408,13 +1392,14 @@ def _forward_pcp_dcp( num_splits=attn_metadata.max_num_splits, **fa_kw, ) - if not plan.has_prefix: + prefix = plan.prefix + if prefix is None: return output # Prefix: gather the prefix-row queries (extend chunks + decode rows) # so every rank evaluates the same rows against its DCP cache shard, # then LSE all-reduce the partials. - num_prefix_rows = plan.prefix_cu_q.shape[0] - 1 + num_prefix_rows = prefix.cu_q.shape[0] - 1 pfx_descale_shape = (num_prefix_rows, self.num_kv_heads) pfx_fa_kw = self._fa_common_kwargs( layer._q_scale.expand(pfx_descale_shape) @@ -1424,23 +1409,23 @@ def _forward_pcp_dcp( layer._v_scale.expand(pfx_descale_shape), ) if n > 0: - q_local = query[plan.prefix_q_local_idx] + q_local = query[prefix.q_local_idx] else: # This rank holds no rows; still must join the collectives. q_local = query.new_zeros( - (plan.padded_num_prefix_tokens, self.num_heads, self.head_size) + (prefix.padded_num_tokens, self.num_heads, self.head_size) ) - q_g = get_pcp_group().all_gather(q_local, dim=0)[plan.prefix_q_restore_idx] + q_g = get_pcp_group().all_gather(q_local, dim=0)[prefix.q_restore_idx] ctx_out, ctx_lse = flash_attn_varlen_func( q=q_g, k=key_cache, v=value_cache, - cu_seqlens_q=plan.prefix_cu_q, - max_seqlen_q=plan.prefix_max_q, - seqused_k=plan.prefix_dcp_ctx_lens, - max_seqlen_k=plan.prefix_max_ctx, + cu_seqlens_q=prefix.cu_q, + max_seqlen_q=prefix.max_q, + seqused_k=prefix.dcp_ctx_lens, + max_seqlen_k=prefix.max_ctx, causal=False, - block_table=plan.prefix_block_table, + block_table=prefix.block_table, return_softmax_lse=True, num_splits=attn_metadata.max_num_splits, **pfx_fa_kw, @@ -1454,10 +1439,8 @@ def _forward_pcp_dcp( ctx_lse_cor = ctx_lse_cor.transpose(0, 1).contiguous() pfx_out = output.new_zeros((n, self.num_heads, self.head_size)) pfx_lse = suffix_lse.new_full((suffix_lse.shape[0], n), float("-inf")) - pfx_out[plan.prefix_local_token_idx] = ctx_out_cor[plan.prefix_local_out_idx] - pfx_lse[:, plan.prefix_local_token_idx] = ctx_lse_cor[ - :, plan.prefix_local_out_idx - ] + pfx_out[prefix.local_token_idx] = ctx_out_cor[prefix.local_out_idx] + pfx_lse[:, prefix.local_token_idx] = ctx_lse_cor[:, prefix.local_out_idx] # Rows without a prefix row keep their suffix output (lse -inf -> 0 weight). merge_attn_states(output, pfx_out, pfx_lse, output, suffix_lse) return output diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index c20b8e54294e..74bf28befa0c 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -36,34 +36,43 @@ def num_tokens(self) -> int: return self.global_batch_slice.stop - self.global_batch_slice.start +class PCPPrefixPlan(NamedTuple): + """Prefix (cached-context) attention plan for one PCP+DCP prefill step. + + Every rank gathers the same prefix-row queries (extend chunks + decode + rows), attends its DCP shard of the cached prefix, and LSE-combines + (all-reduce) across the group; the combine requires the same queries on + every rank. Index tensors are views into one packed buffer uploaded with + a single H2D copy per step. + """ + + padded_num_tokens: int # all-gather slab size (rank-invariant) + q_local_idx: torch.Tensor # [padded_num_tokens] int64, local token coords + q_restore_idx: torch.Tensor # [total_global_prefix_tokens] int64 + cu_q: torch.Tensor # [num_global_prefix_rows + 1] int32 + max_q: int + dcp_ctx_lens: torch.Tensor # [num_global_prefix_rows] int32, this DCP rank + max_ctx: int + block_table: torch.Tensor # [num_global_prefix_rows, max_num_blocks] + local_token_idx: torch.Tensor # [num_local_prefix_tokens] int64 + local_out_idx: torch.Tensor # [num_local_prefix_tokens] int64 + + class PCPRowPlan(NamedTuple): - """Per-step row-level attention plan for the sharded PCP+DCP path. + """Per-step row plan for the sharded PCP+DCP prefill path. Every rank-local row is a (cached prefix, new-token span) pair. The suffix (new tokens) attends the write-gathered K/V -- the same all-gather the - cache write already performs -- so it needs no collective. The prefix - (cached context) attends this rank's DCP cache shard with the - PCP-gathered prefix-row queries and LSE-combines (all-reduce) across the - group, which requires every rank to hold the same prefix-row queries. + cache write already performs -- so it needs no collective. ``prefix`` is + None when no global row has a cached context this step (pure + fresh-prefill batch), letting the prefix attention and its query gather + be skipped on every rank. """ - # Suffix attention (rank-local rows vs the stashed write-gathered K/V). suffix_kv_idx: torch.Tensor # [total_suffix_k] int64, into stashed K/V suffix_cu_k: torch.Tensor # [num_local_rows + 1] int32 suffix_max_k: int - # Prefix attention (gathered rows vs the local DCP cache shard). - has_prefix: bool # rank-invariant: some global row has cached context - num_local_prefix_tokens: int - padded_num_prefix_tokens: int # all-gather slab size (rank-invariant) - prefix_q_local_idx: torch.Tensor # [padded_num_prefix_tokens] int64 - prefix_q_restore_idx: torch.Tensor # [total_global_prefix_tokens] int64 - prefix_cu_q: torch.Tensor # [num_global_prefix_rows + 1] int32 - prefix_max_q: int - prefix_dcp_ctx_lens: torch.Tensor # [num_global_prefix_rows] int32 - prefix_max_ctx: int - prefix_block_table: torch.Tensor # [num_global_prefix_rows, max_num_blocks] - prefix_local_token_idx: torch.Tensor # [num_local_prefix_tokens] int64 - prefix_local_out_idx: torch.Tensor # [num_local_prefix_tokens] int64 + prefix: PCPPrefixPlan | None class PCPManager: @@ -102,12 +111,11 @@ def __init__( self._padded_gather_idx: torch.Tensor | None = None self._gathered_kv_write_mask: torch.Tensor | None = None # GLOBAL batch composition (rank-invariant: every PCP rank sees the same - # global batch, so these are identical across ranks). Used for - # rank-consistent mixed-batch detection in the sharded attention path -- - # per-rank is_prefilling can differ (DualChunkSwap leaves some ranks with - # zero prefill chunks), which would desync NCCL collectives. + # global batch, so this is identical across ranks). Gates the + # cache-write all-gather and the row-plan build -- per-rank + # is_prefilling can differ (DualChunkSwap leaves some ranks with zero + # prefill chunks), which would desync NCCL collectives. self._global_has_prefill: bool = False - self._global_has_decode: bool = False self._pad_slot_id = torch.tensor(PAD_SLOT_ID, dtype=torch.int64, device=device) # Per-step row-plan state for the sharded PCP+DCP attention path. @@ -404,7 +412,6 @@ def partition_batch(self, input_batch: InputBatch) -> InputBatch: is_prefilling = global_batch.is_prefilling_np # Rank-invariant global composition (see __init__ comment). self._global_has_prefill = bool(is_prefilling.any()) - self._global_has_decode = bool((~is_prefilling).any()) segments_by_rank, per_rank_num_tokens = self._build_batch_layout( num_scheduled_tokens, @@ -686,25 +693,19 @@ def restore_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: gathered = get_pcp_group().all_gather(hidden_states, dim=0) return gathered[self._hidden_restore_idx] - def global_batch_flags(self) -> tuple[bool, bool]: - """Return (has_prefill, has_decode) for the global (pre-partition) batch. - - Rank-invariant across PCP ranks (identical global batch). Used for - rank-consistent mixed-batch detection in the sharded attention path. - """ - return self._global_has_prefill, self._global_has_decode - def populate_forward_context(self) -> None: - """Stash all PCP metadata the attention forward / kv-cache update need. + """Stash the PCP metadata the attention forward / kv-cache update need. Centralized here so the model runner stays a single call. Populates: - - ``pcp_global_flags``: rank-invariant (has_prefill, has_decode) for - consistent mixed-batch routing and cache-write gather decisions. + - ``pcp_has_prefill``: rank-invariant bool gating the cache-write + all-gather (when the global batch has any prefill, every rank + gathers the whole batch; otherwise writes stay local). - ``pcp_row_plan``: per-step :class:`PCPRowPlan` for the sharded - PCP+DCP attention path (None for warmup or pure-decode steps). + PCP+DCP prefill path (None for warmup or pure-decode steps, which + is also how forward() tells the two paths apart). """ kwargs = get_forward_context().additional_kwargs - kwargs["pcp_global_flags"] = self.global_batch_flags() + kwargs["pcp_has_prefill"] = self._global_has_prefill kwargs["pcp_row_plan"] = self.build_cp_row_plan() def build_cp_row_plan(self) -> "PCPRowPlan | None": @@ -876,37 +877,68 @@ def is_prefix_seg(seg: RankSegment) -> bool: prefix_q_local_np = np.zeros(padded_num_prefix, dtype=np.int64) prefix_q_local_np[:num_local_prefix] = prefix_local_token_np - row_req_gpu = async_copy_to_gpu(row_req_np, device=self.device) + # Pack the index arrays into one int64 and one int32 buffer -- one + # H2D copy each; a pin_memory() per small tensor is the dominant cost + # here. The plan hands out views into the packed buffers. num_prefix_rows = len(global_prefix_rows) - prefix_block_table = self._block_tables.gather_block_tables( - torch.index_select(gb.idx_mapping, 0, row_req_gpu), - num_prefix_rows, - out=self._prefix_block_tables, - out_ptrs=self._prefix_block_table_ptrs, - )[0] - + blob_a = async_copy_to_gpu( + np.concatenate( + ( + suffix_kv_idx_np, + prefix_q_local_np, + prefix_q_restore_np, + prefix_local_token_np, + prefix_local_out_np, + row_req_np, + ) + ), + device=self.device, + ) + blob_b = async_copy_to_gpu( + np.concatenate((suffix_cu_np, prefix_cu_np, dcp_ctx_np)), + device=self.device, + ) + n_sfx = len(suffix_kv_idx_np) + n_restore = len(prefix_q_restore_np) + suffix_kv_idx = blob_a[:n_sfx] + prefix_q_local_idx = blob_a[n_sfx : n_sfx + padded_num_prefix] + q_restore_off = n_sfx + padded_num_prefix + prefix_q_restore_idx = blob_a[q_restore_off : q_restore_off + n_restore] + loc_tok_off = q_restore_off + n_restore + prefix_local_token_idx = blob_a[loc_tok_off : loc_tok_off + num_local_prefix] + loc_out_off = loc_tok_off + num_local_prefix + prefix_local_out_idx = blob_a[loc_out_off : loc_out_off + num_local_prefix] + row_req_gpu = blob_a[loc_out_off + num_local_prefix :] + n_cu_k = len(suffix_cu_np) + suffix_cu_k = blob_b[:n_cu_k] + prefix_cu_q = blob_b[n_cu_k : n_cu_k + num_prefix_rows + 1] + prefix_dcp_ctx_lens = blob_b[n_cu_k + num_prefix_rows + 1 :] + + prefix_plan = None + if has_prefix: + prefix_block_table = self._block_tables.gather_block_tables( + torch.index_select(gb.idx_mapping, 0, row_req_gpu), + num_prefix_rows, + out=self._prefix_block_tables, + out_ptrs=self._prefix_block_table_ptrs, + )[0] + prefix_plan = PCPPrefixPlan( + padded_num_tokens=padded_num_prefix, + q_local_idx=prefix_q_local_idx, + q_restore_idx=prefix_q_restore_idx, + cu_q=prefix_cu_q, + max_q=prefix_max_q, + dcp_ctx_lens=prefix_dcp_ctx_lens, + max_ctx=prefix_max_ctx, + block_table=prefix_block_table, + local_token_idx=prefix_local_token_idx, + local_out_idx=prefix_local_out_idx, + ) return PCPRowPlan( - suffix_kv_idx=async_copy_to_gpu(suffix_kv_idx_np, device=self.device), - suffix_cu_k=async_copy_to_gpu(suffix_cu_np, device=self.device), + suffix_kv_idx=suffix_kv_idx, + suffix_cu_k=suffix_cu_k, suffix_max_k=suffix_max_k, - has_prefix=has_prefix, - num_local_prefix_tokens=num_local_prefix, - padded_num_prefix_tokens=padded_num_prefix, - prefix_q_local_idx=async_copy_to_gpu(prefix_q_local_np, device=self.device), - prefix_q_restore_idx=async_copy_to_gpu( - prefix_q_restore_np, device=self.device - ), - prefix_cu_q=async_copy_to_gpu(prefix_cu_np, device=self.device), - prefix_max_q=prefix_max_q, - prefix_dcp_ctx_lens=async_copy_to_gpu(dcp_ctx_np, device=self.device), - prefix_max_ctx=prefix_max_ctx, - prefix_block_table=prefix_block_table, - prefix_local_token_idx=async_copy_to_gpu( - prefix_local_token_np, device=self.device - ), - prefix_local_out_idx=async_copy_to_gpu( - prefix_local_out_np, device=self.device - ), + prefix=prefix_plan, ) def restore_for_sampling( From c65076b0eb25fdae9bda905d750932e61d4f23ad Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Fri, 31 Jul 2026 03:42:15 +0000 Subject: [PATCH 14/18] remove diff at attn forward context Signed-off-by: JaredforReal --- vllm/v1/attention/backends/flash_attn.py | 31 +++++++++++++----------- vllm/v1/worker/gpu/model_runner.py | 2 +- vllm/v1/worker/gpu/pcp_manager.py | 21 ++++++++++------ 3 files changed, 32 insertions(+), 22 deletions(-) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index d7f72a5a499f..f59c0708411d 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -279,6 +279,12 @@ class FlashAttentionMetadata: max_dcp_context_kv_len: int | None = None dcp_context_kv_lens: torch.Tensor | None = None + # PCP per-step state, stamped by PCPManager.populate_attn_metadata (not via + # forward_context). pcp_row_plan selects _forward_pcp_dcp vs _forward_with_dcp; + # pcp_has_prefill gates the rank-invariant cache-write all-gather. + pcp_has_prefill: bool = False + pcp_row_plan: "PCPRowPlan | None" = None + # Per-segment is_prefilling flag (rank-local view). Used by the builder to # detect prefill vs decode batches under DualChunkSwap (a prefill chunk can # be a single token, so max_query_len is unreliable). @@ -901,6 +907,10 @@ def __init__( # self.pcp_world_size / self.pcp_rank are auto-populated by # AttentionImplBase.__new__ from get_pcp_group(). self.use_pcp = self.pcp_world_size > 1 + # Per-layer write-gathered K/V from do_kv_cache_update, reused by the + # suffix attention in _forward_pcp_dcp. Keyed by layer_name because this + # impl is shared across the layers of an attention group. + self._pcp_gathered_kv: dict[str, tuple[torch.Tensor, torch.Tensor] | None] = {} # How Q and the partials move across the DCP group is decided by the # PCP/DCP topology alone; both backends share the rule (see pcp.py). self.dcp_combine = resolve_dcp_combine_fn(vllm_config) @@ -1035,11 +1045,7 @@ def forward( # all-reduce); only a PCP+DCP prefill/mixed step -- where # DualChunkSwap partitions the queries and the manager builds # a row plan -- goes to _forward_pcp_dcp. - plan = ( - get_forward_context().additional_kwargs.get("pcp_row_plan") - if self.use_pcp - else None - ) + plan = attn_metadata.pcp_row_plan if self.use_pcp else None if plan is not None: # MRv2 PCP+DCP prefill (dcp == pcp): every row is a # (cached prefix, new-token span) pair -- the suffix @@ -1254,7 +1260,7 @@ def do_kv_cache_update( # Rank-invariant gather decision: per-rank num_decode_tokens can differ # under DualChunkSwap and desync the all-gather, so if the global batch # has any prefill every rank gathers the whole batch. - if get_forward_context().additional_kwargs.get("pcp_has_prefill"): + if attn_metadata is not None and attn_metadata.pcp_has_prefill: num_decode_tokens = 0 key_cache, value_cache = kv_cache.transpose(1, 2).split( self.head_size, dim=-1 @@ -1267,12 +1273,12 @@ def do_kv_cache_update( ) # Stash the write-gathered K/V for the suffix attention of the # sharded PCP+DCP path (a gather only happens on prefill steps). - # Keyed by layer and scoped to this step's forward context, so a + # Keyed by layer on the impl (shared across an attention group), so a # layer whose cache update is skipped (kv sharing) cannot pick up # stale tensors. - get_forward_context().additional_kwargs[ - f"pcp_gathered_kv:{layer.layer_name}" - ] = (cache_key, cache_value) if num_decode_tokens == 0 else None + self._pcp_gathered_kv[layer.layer_name] = ( + (cache_key, cache_value) if num_decode_tokens == 0 else None + ) reshape_and_cache_flash( cache_key, cache_value, @@ -1361,7 +1367,6 @@ def _forward_pcp_dcp( assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." ) - fc = get_forward_context() query = query.contiguous() fa_kw = self._fa_common_kwargs(q_descale, k_descale, v_descale) @@ -1370,9 +1375,7 @@ def _forward_pcp_dcp( if n > 0: # Suffix: local rows vs the write-gathered new K/V. The gather is # the one do_kv_cache_update already did for the cache write. - gathered_kv = fc.additional_kwargs.get( - f"pcp_gathered_kv:{layer.layer_name}" - ) + gathered_kv = self._pcp_gathered_kv.get(layer.layer_name) assert gathered_kv is not None, ( "PCP+DCP prefill step without write-gathered K/V: " "do_kv_cache_update did not run for this layer (kv sharing?)." diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 6debc92788b4..8580aadca1cf 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1381,7 +1381,7 @@ def execute_model( is_padding=input_batch.is_padding, ): if self.pcp_manager is not None: - self.pcp_manager.populate_forward_context() + self.pcp_manager.populate_attn_metadata(attn_metadata) self.kv_connector.pre_forward(scheduler_output) if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE: # Run the PIECEWISE graph (compiled PW cudagraph or breakable diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 74bf28befa0c..9b8dc98b6e39 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -8,8 +8,8 @@ from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group -from vllm.forward_context import get_forward_context from vllm.logger import init_logger +from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu @@ -693,20 +693,27 @@ def restore_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: gathered = get_pcp_group().all_gather(hidden_states, dim=0) return gathered[self._hidden_restore_idx] - def populate_forward_context(self) -> None: - """Stash the PCP metadata the attention forward / kv-cache update need. + def populate_attn_metadata(self, attn_metadata: dict | None) -> None: + """Stamp the PCP per-step state onto each GQA attention metadata object. - Centralized here so the model runner stays a single call. Populates: + Centralized here so the model runner stays a single call. Sets, on every + FlashAttentionMetadata in ``attn_metadata``: - ``pcp_has_prefill``: rank-invariant bool gating the cache-write all-gather (when the global batch has any prefill, every rank gathers the whole batch; otherwise writes stay local). - ``pcp_row_plan``: per-step :class:`PCPRowPlan` for the sharded PCP+DCP prefill path (None for warmup or pure-decode steps, which is also how forward() tells the two paths apart). + MLA uses its own metadata class and has no row plan, so non-GQA objects + are skipped. """ - kwargs = get_forward_context().additional_kwargs - kwargs["pcp_has_prefill"] = self._global_has_prefill - kwargs["pcp_row_plan"] = self.build_cp_row_plan() + if not attn_metadata: + return + plan = self.build_cp_row_plan() + for meta in attn_metadata.values(): + if isinstance(meta, FlashAttentionMetadata): + meta.pcp_row_plan = plan + meta.pcp_has_prefill = self._global_has_prefill def build_cp_row_plan(self) -> "PCPRowPlan | None": """Build the per-step row plan for the sharded PCP+DCP attention path. From 115e11d5450770a8f2343450fe38d2f68c025855 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Fri, 31 Jul 2026 09:29:49 +0000 Subject: [PATCH 15/18] rewrite PCP Plan with global batch Signed-off-by: JaredforReal Co-authored-by: Lucas Wilkinson --- .../layers/attention/mla_attention.py | 88 +-- vllm/model_executor/layers/attention/pcp.py | 138 +---- .../layers/sparse_attn_indexer.py | 6 +- vllm/v1/attention/backends/flash_attn.py | 585 +++++++----------- vllm/v1/worker/gpu/model_runner.py | 4 +- vllm/v1/worker/gpu/pcp_manager.py | 418 +++++-------- 6 files changed, 453 insertions(+), 786 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 415f41274fd1..16fb961e563a 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -211,6 +211,7 @@ from vllm.config.cache import CacheDType from vllm.distributed.parallel_state import ( get_dcp_group, + get_tp_group, is_global_first_rank, ) from vllm.forward_context import ForwardContext, get_forward_context @@ -226,10 +227,8 @@ maybe_transfer_kv_layer, ) from vllm.model_executor.layers.attention.pcp import ( - cp_reconcile_heads, - maybe_all_gather_q_for_dcp, - maybe_gather_cache_inputs, - resolve_dcp_combine_fn, + finalize_mla_pcp_decode, + maybe_gather_mla_latent_cache_inputs, ) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.linear import ( @@ -275,6 +274,8 @@ get_dcp_local_seq_lens, split_decodes_and_prefills, ) +from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs +from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.attention.ops.triton_merge_attn_states import mask_empty_context from vllm.v1.attention.selector import get_attn_backend @@ -342,24 +343,6 @@ def _canonicalize_sparse_mla_kv_cache_dtype( return kv_cache_dtype -def _gather_latent_cache_inputs( - kv_c_normed: torch.Tensor, - k_pe: torch.Tensor, - slot_mapping: torch.Tensor | None, - num_decode_tokens: int | None, - use_pcp: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """PCP latent-cache write gather; flattens k_pe for the all-gather.""" - if not use_pcp or num_decode_tokens is None: - return kv_c_normed, k_pe, slot_mapping - assert slot_mapping is not None - k_pe_flat = k_pe.reshape(kv_c_normed.shape[0], -1) - (cache_kv_c, cache_k_pe_flat), cache_slot_mapping = maybe_gather_cache_inputs( - (kv_c_normed, k_pe_flat), slot_mapping, num_decode_tokens, use_pcp - ) - return cache_kv_c, cache_k_pe_flat.view(-1, *k_pe.shape[1:]), cache_slot_mapping - - class MLAAttention(nn.Module, AttentionLayerBase): """Multi-Head Latent Attention layer. @@ -567,7 +550,12 @@ def __init__( self.use_sparse = use_sparse - self.dcp_combine = resolve_dcp_combine_fn(get_current_vllm_config_or_none()) + _vllm_config = get_current_vllm_config_or_none() + self.dcp_a2a = ( + _vllm_config is not None + and _vllm_config.parallel_config.decode_context_parallel_size > 1 + and _vllm_config.parallel_config.dcp_comm_backend == "a2a" + ) # Initialize q/k/v range constants. self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32) @@ -643,7 +631,7 @@ def forward( ) layer_slot_mapping = slot_mapping.get(self.layer_name) kv_for_cache, kpe_for_cache, layer_slot_mapping = ( - _gather_latent_cache_inputs( + maybe_gather_mla_latent_cache_inputs( kv_c_normed, k_pe, layer_slot_mapping, @@ -887,15 +875,18 @@ def forward_impl( mqa_q = (mqa_ql_nope, mqa_q_pe) # concatenate nope + pe -> (B, N, L + P) (fp8 op above may have fused) if self.impl.dcp_world_size > 1: - if not self.use_pcp and isinstance(mqa_q, tuple): - # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) - mqa_q = torch.cat(mqa_q, dim=-1) - mqa_q = maybe_all_gather_q_for_dcp( - mqa_q, - self.impl.dcp_world_size, - self.impl.pcp_world_size, - already_replicated=qrep_decode and not self.use_pcp, - ) + if self.use_pcp: + if self.impl.dcp_world_size > self.impl.pcp_world_size: + if isinstance(mqa_q, tuple): + mqa_q = torch.cat(mqa_q, dim=-1) + mqa_q = get_tp_group().all_gather(mqa_q, dim=1) + else: + if isinstance(mqa_q, tuple): + # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) + mqa_q = torch.cat(mqa_q, dim=-1) + if not qrep_decode: + # mqa_q do allgather in head dim. + mqa_q = get_dcp_group().all_gather(mqa_q, dim=1) # call decode attn if not self.impl.is_sparse: @@ -905,14 +896,29 @@ def forward_impl( # correct dcp attn_out with lse. if self.impl.dcp_world_size > 1: assert lse is not None - attn_out = self.dcp_combine( - attn_out, - lse, - get_dcp_group(), - is_lse_base_on_e=self.impl.lse_base_on_e, - ) + if self.dcp_a2a: + attn_out = dcp_a2a_lse_reduce( + attn_out, + lse, + get_dcp_group(), + is_lse_base_on_e=self.impl.lse_base_on_e, + ) + elif self.use_pcp: + attn_out = cp_lse_ag_out_ar( + attn_out, + lse, + get_dcp_group(), + is_lse_base_on_e=self.impl.lse_base_on_e, + ) + else: + attn_out = cp_lse_ag_out_rs( + attn_out, + lse, + get_dcp_group(), + is_lse_base_on_e=self.impl.lse_base_on_e, + ) if self.use_pcp: - attn_out = cp_reconcile_heads(attn_out, self.num_heads) + attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads) # v_up projection self._v_up_proj(attn_out, out=mqa_output_slice) @@ -1160,7 +1166,7 @@ def unified_mla_kv_cache_update( layer_name ) if layer_slot_mapping is not None: - kv_c_normed, k_pe, layer_slot_mapping = _gather_latent_cache_inputs( + kv_c_normed, k_pe, layer_slot_mapping = maybe_gather_mla_latent_cache_inputs( kv_c_normed, k_pe, layer_slot_mapping, diff --git a/vllm/model_executor/layers/attention/pcp.py b/vllm/model_executor/layers/attention/pcp.py index 467b20f6097b..75ab1c9e8e13 100644 --- a/vllm/model_executor/layers/attention/pcp.py +++ b/vllm/model_executor/layers/attention/pcp.py @@ -1,24 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING, TypeVar - import torch from vllm.distributed.parallel_state import ( - GroupCoordinator, - get_dcp_group, get_pcp_group, get_tp_group, ) -from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs -from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce - -if TYPE_CHECKING: - from vllm.config import VllmConfig - -# Query as handed to the DCP context attention: a tensor, or MLA's split -# (nope, pe) pair which is only concatenated if a gather happens. -_QueryT = TypeVar("_QueryT", torch.Tensor, tuple[torch.Tensor, ...]) def _gather_prefill_cache_inputs( @@ -58,113 +45,42 @@ def _gather_prefill_cache_inputs( return cache_inputs, cache_slot_mapping -def maybe_gather_cache_inputs( - tensors: tuple[torch.Tensor, ...], +def maybe_gather_mla_latent_cache_inputs( + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, slot_mapping: torch.Tensor | None, num_decode_tokens: int | None, use_pcp: bool, -) -> tuple[tuple[torch.Tensor, ...], torch.Tensor | None]: - """PCP cache-write gather. - - All-gather the prefill portion of ``tensors`` (K/V for GQA, latent KV and - k_pe for MLA, k for the sparse indexer) across PCP ranks so every rank can - write the full prefill contents to its cache, while keeping decode writes - local. Returns the gathered tensors and the gathered cache slot mapping. - No-op when PCP is off. - - Gathering each tensor separately (rather than ``cat``-then-``split``) - keeps each output contiguous, so the cache kernel's head-stride - assumption holds. - """ +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: if not use_pcp or num_decode_tokens is None: - return tensors, slot_mapping + return kv_c_normed, k_pe, slot_mapping assert slot_mapping is not None - return _gather_prefill_cache_inputs(tensors, slot_mapping, num_decode_tokens) - - -def _dcp_q_gather_group( - dcp_world_size: int, - pcp_world_size: int, -) -> GroupCoordinator | None: - """Group holding the query-head shards this rank must gather, or None. - - DCP groups span the PCP axis before TP (see ``init_model_parallel_groups``), - so how far the group reaches decides which ranks hold distinct query heads: - - - ``pcp == 1``: the DCP group is a TP subgroup that splits Q heads, so - gather over the DCP group. - - ``dcp > pcp``: the group spans TP x PCP. Q is head-sharded across TP and - replicated across PCP, so gather over the TP group -- gathering over the - DCP group here would duplicate the PCP-replicated heads. - - otherwise: the DCP group is the PCP group and already holds every head. - """ - if pcp_world_size == 1: - return get_dcp_group() if dcp_world_size > 1 else None - if dcp_world_size > pcp_world_size: - return get_tp_group() - return None - - -def dcp_q_gather_size(dcp_world_size: int, pcp_world_size: int) -> int: - """Head-shard count ``maybe_all_gather_q_for_dcp`` will produce. - - For sizing work that has to be planned before the gather happens, such as - FlashAttention's AOT scheduler metadata. 1 means no gather. - """ - group = _dcp_q_gather_group(dcp_world_size, pcp_world_size) - return 1 if group is None else group.world_size - - -def maybe_all_gather_q_for_dcp( - q: _QueryT, - dcp_world_size: int, - pcp_world_size: int, - already_replicated: bool = False, -) -> _QueryT | torch.Tensor: - """All-gather query heads for attention against the DCP-sharded cache. - - Returns ``q`` untouched when this rank already holds every head the kernel - needs; ``already_replicated`` lets a caller declare that up front. A split - ``q`` -- MLA's ``(nope, pe)`` pair -- is concatenated only when a gather - actually happens, so callers that can consume the split form keep it. - - Pairs with ``resolve_dcp_combine_fn``: changing one without the other breaks - the head layout the combine expects. - """ - if already_replicated: - return q - group = _dcp_q_gather_group(dcp_world_size, pcp_world_size) - if group is None: - return q - gathered = torch.cat(q, dim=-1) if isinstance(q, tuple) else q - return group.all_gather(gathered, dim=1) - - -def resolve_dcp_combine_fn(vllm_config: "VllmConfig | None"): - """LSE-combine to fold per-rank partial attentions over the DCP group. + num_tokens = kv_c_normed.shape[0] + k_pe_flat = k_pe.reshape(num_tokens, -1) + (cache_kv_c, cache_k_pe_flat), cache_slot_mapping = _gather_prefill_cache_inputs( + (kv_c_normed, k_pe_flat), + slot_mapping, + num_decode_tokens, + ) + cache_k_pe = cache_k_pe_flat.view(-1, *k_pe.shape[1:]) + return cache_kv_c, cache_k_pe, cache_slot_mapping - Under PCP the partials cover the full head set on every rank, so they - all-reduce and each rank takes its own heads back out with - ``cp_reconcile_heads``. Without PCP the DCP group is a TP subgroup, so the - reduce-scatter lands each rank's heads directly. - Pairs with ``maybe_all_gather_q_for_dcp``. ``vllm_config`` may be None when - there is no config in scope (unit tests), which implies no CP. - """ - if vllm_config is None: - return cp_lse_ag_out_rs - parallel_config = vllm_config.parallel_config - if ( - parallel_config.decode_context_parallel_size > 1 - and parallel_config.dcp_comm_backend == "a2a" - ): - return dcp_a2a_lse_reduce - if parallel_config.prefill_context_parallel_size > 1: - return cp_lse_ag_out_ar - return cp_lse_ag_out_rs +def maybe_gather_indexer_k( + k: torch.Tensor, + slot_mapping: torch.Tensor, + num_decode_tokens: int, + use_pcp: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + if not use_pcp: + return k, slot_mapping + (cache_k,), cache_slot_mapping = _gather_prefill_cache_inputs( + (k,), slot_mapping, num_decode_tokens + ) + return cache_k, cache_slot_mapping -def cp_reconcile_heads( +def finalize_mla_pcp_decode( output: torch.Tensor, num_heads: int, ) -> torch.Tensor: diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index fa493f3e217a..9a671be55639 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -13,7 +13,7 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp -from vllm.model_executor.layers.attention.pcp import maybe_gather_cache_inputs +from vllm.model_executor.layers.attention.pcp import maybe_gather_indexer_k from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) @@ -388,8 +388,8 @@ def sparse_attn_indexer( if not skip_k_cache_insert: assert k is not None - (k,), slot_mapping_for_cache = maybe_gather_cache_inputs( - (k,), + k, slot_mapping_for_cache = maybe_gather_indexer_k( + k, slot_mapping, num_decode_tokens, use_pcp, diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index f59c0708411d..3038de57b41a 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -4,21 +4,12 @@ import copy from dataclasses import dataclass -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np import torch -if TYPE_CHECKING: - from vllm.v1.worker.gpu.pcp_manager import PCPRowPlan - from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.attention.pcp import ( - dcp_q_gather_size, - maybe_all_gather_q_for_dcp, - maybe_gather_cache_inputs, - resolve_dcp_combine_fn, -) from vllm.platforms import current_platform from vllm.utils.torch_utils import ( canonicalize_singleton_dim_strides, @@ -38,9 +29,15 @@ is_flash_attn_varlen_func_available, ) from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens +from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs +from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce from vllm.v1.attention.ops.merge_attn_states import merge_attn_states +from vllm.v1.attention.ops.triton_merge_attn_states import mask_empty_context from vllm.v1.worker.workspace import current_workspace_manager +if TYPE_CHECKING: + from vllm.v1.worker.gpu.pcp_manager import PCPPlan + if is_flash_attn_varlen_func_available(): from vllm.v1.attention.backends.fa_utils import ( flash_attn_supports_sinks, @@ -56,10 +53,6 @@ ) from vllm.config.cache import CacheDType from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group -from vllm.forward_context import ( - get_forward_context, - is_forward_context_available, -) from vllm.logger import init_logger from vllm.platforms.interface import DeviceCapability from vllm.utils.math_utils import cdiv, round_up @@ -279,16 +272,11 @@ class FlashAttentionMetadata: max_dcp_context_kv_len: int | None = None dcp_context_kv_lens: torch.Tensor | None = None - # PCP per-step state, stamped by PCPManager.populate_attn_metadata (not via - # forward_context). pcp_row_plan selects _forward_pcp_dcp vs _forward_with_dcp; - # pcp_has_prefill gates the rank-invariant cache-write all-gather. - pcp_has_prefill: bool = False - pcp_row_plan: "PCPRowPlan | None" = None - - # Per-segment is_prefilling flag (rank-local view). Used by the builder to - # detect prefill vs decode batches under DualChunkSwap (a prefill chunk can - # be a single token, so max_query_len is unreliable). - is_prefilling: torch.Tensor | None = None + # MRv2 PCP. Stamped per step by PCPManager.populate_attn_metadata (never via + # forward_context). Non-None only on a PCP+DCP step whose global batch has a + # cached context, which is also how forward() picks _forward_pcp_dcp over + # _forward_with_dcp. + pcp_plan: "PCPPlan | None" = None # Split counts for FA2 DCP context attention. num_prefill_* tracks # context-bearing extend rows; pure prefills do not attend to DCP context. @@ -384,11 +372,8 @@ def get_cudagraph_support( vllm_config: "VllmConfig", kv_cache_spec: "AttentionSpec", ) -> AttentionCGSupport: - # Under PCP, do_kv_cache_update runs a PCP all-gather to materialize the - # full replicated KV cache before attention. That collective is not - # captureable into a CUDA graph, so force NEVER and let the runner use - # piecewise graphs (rest of the model captured, attention eager) when - # PCP is on. + # PCP runs an all-gather inside the KV cache update, which cannot be + # captured. Fall back to piecewise graphs (attention stays eager). if vllm_config.parallel_config.prefill_context_parallel_size > 1: return AttentionCGSupport.NEVER return cls._cudagraph_support @@ -428,21 +413,10 @@ def __init__( self.dcp_world_size = 1 self.dcp_rank = 0 - try: - from vllm.distributed.parallel_state import get_pcp_group - - self.pcp_world_size = get_pcp_group().world_size - self.pcp_rank = get_pcp_group().rank_in_group - except AssertionError: - # PCP might not be initialized in testing. - self.pcp_world_size = 1 - self.pcp_rank = 0 - self.cp_kv_cache_interleave_size = ( self.parallel_config.cp_kv_cache_interleave_size ) - - self.use_pcp = self.pcp_world_size > 1 + self.use_pcp = self.parallel_config.prefill_context_parallel_size > 1 self.use_full_cuda_graph = ( self.compilation_config.cudagraph_mode.has_full_cudagraphs() @@ -472,7 +446,11 @@ def __init__( ) if self.dcp_world_size > 1: - max_num_reqs = vllm_config.scheduler_config.max_num_seqs + # DualChunkSwap gives the rank-local batch up to two rows per + # prefilling request, so under PCP it can hold 2x max_num_seqs rows. + max_num_reqs = vllm_config.scheduler_config.max_num_seqs * ( + 2 if self.use_pcp else 1 + ) self._dcp_context_kv_lens = torch.zeros( max_num_reqs, dtype=torch.int32, @@ -496,37 +474,6 @@ def __init__( [self.rswa_window], dtype=torch.int32, device=self.device ) - def _build_dcp_context_lens( - self, - num_reqs: int, - query_start_loc: torch.Tensor, - seq_lens: torch.Tensor, - max_seq_len: int, - ) -> tuple[torch.Tensor, int]: - """This rank's DCP shard of each row's cached-context length. - - The cached context (seq_lens - query_lens) is interleave-sharded - across DCP ranks; each rank attends only its shard. The returned max - is the aligned per-rank upper bound, ceil(L / (N * I)) * I with - L = max_seq_len, N = dcp_world_size, I = interleave size, which avoids - a GPU->CPU sync while minimizing workspace over-allocation. - """ - query_lens = query_start_loc[1:] - query_start_loc[:-1] - context_kv_lens = seq_lens - query_lens - local_context_kv_lens = get_dcp_local_seq_lens( - context_kv_lens, - self.dcp_world_size, - self.dcp_rank, - self.cp_kv_cache_interleave_size, - ) - self._dcp_context_kv_lens[:num_reqs] = local_context_kv_lens - self._dcp_context_kv_lens[num_reqs:] = 0 - num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size - max_dcp_context_kv_len = ( - (max_seq_len + num_partitions - 1) // num_partitions - ) * self.cp_kv_cache_interleave_size - return self._dcp_context_kv_lens[:num_reqs], max_dcp_context_kv_len - def build( self, common_prefix_len: int, @@ -585,28 +532,19 @@ def build( max_num_splits = 1 def schedule( - batch_size, - cu_query_lens, - max_query_len, - seqlens, - max_seq_len, - causal, + batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal ): cache_dtype = self.cache_config.cache_dtype if is_quantized_kv_cache(cache_dtype): qkv_dtype = current_platform.fp8_dtype() else: qkv_dtype = self.kv_cache_dtype - # The kernel sees whatever maybe_all_gather_q_for_dcp() hands it. - num_heads_q = self.num_heads_q * dcp_q_gather_size( - self.dcp_world_size, self.pcp_world_size - ) if aot_schedule: return get_scheduler_metadata( batch_size=batch_size, max_seqlen_q=max_query_len, max_seqlen_k=max_seq_len, - num_heads_q=num_heads_q, + num_heads_q=self.num_heads_q * self.dcp_world_size, num_heads_kv=self.num_heads_kv, headdim=self.headdim, cache_seqlens=seqlens, @@ -635,60 +573,68 @@ def schedule( prefix_scheduler_metadata = None if self.dcp_world_size > 1: - # Sharded (1/dcp) KV cache: every row's cached context - # (seq_lens - query_lens) is attended via this rank's interleaved - # shard. Pure-DCP batches and PCP+DCP decode batches both use - # this metadata; PCP+DCP prefill/mixed steps run on the - # PCPRowPlan from the forward context instead. - dcp_context_kv_lens, max_dcp_context_kv_len = self._build_dcp_context_lens( - num_reqs, query_start_loc, seq_lens, max_seq_len + query_lens = query_start_loc[1:] - query_start_loc[:-1] + context_kv_lens = seq_lens - query_lens + local_context_kv_lens = get_dcp_local_seq_lens( + context_kv_lens, + self.dcp_world_size, + self.dcp_rank, + self.cp_kv_cache_interleave_size, ) - scheduler_metadata = None - if self.use_pcp: - # Detect prefill vs decode for the cache-write gather. Under - # DualChunkSwap a prefill chunk can be as small as 1 token, so - # max_query_len is unreliable -- use the is_prefilling flag. - is_prefilling = common_attn_metadata.is_prefilling - if is_prefilling is not None: - is_prefill_batch = bool(is_prefilling.any().item()) - else: - is_prefill_batch = max_query_len > 1 - if is_prefill_batch: - num_prefill_tokens = num_actual_tokens - else: - num_decode_tokens = num_actual_tokens - else: - skip_dcp_context_attention = False - if common_attn_metadata.seq_lens_cpu_upper_bound is not None: - query_lens_cpu = ( - common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1] - - common_attn_metadata.query_start_loc_cpu[:num_reqs] - ) - context_kv_lens_cpu = ( - common_attn_metadata.seq_lens_cpu_upper_bound[:num_reqs] - - query_lens_cpu - ) - skip_dcp_context_attention = should_skip_dcp_context_attention( - context_kv_lens_cpu - ) + self._dcp_context_kv_lens[:num_reqs] = local_context_kv_lens + self._dcp_context_kv_lens[num_reqs:] = 0 + dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs] + + skip_dcp_context_attention = False + if ( + not self.use_pcp + and common_attn_metadata.seq_lens_cpu_upper_bound is not None + ): + query_lens_cpu = ( + common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1] + - common_attn_metadata.query_start_loc_cpu[:num_reqs] + ) + context_kv_lens_cpu = ( + common_attn_metadata.seq_lens_cpu_upper_bound[:num_reqs] + - query_lens_cpu + ) + skip_dcp_context_attention = should_skip_dcp_context_attention( + context_kv_lens_cpu + ) - if max_query_len > 1: - ( - num_decode_reqs, - num_prefill_reqs, - num_decode_tokens, - num_prefill_tokens, - ) = split_dcp_context_queries( - common_attn_metadata.query_start_loc_cpu, - common_attn_metadata.seq_lens_cpu_upper_bound, - max_query_len, - num_actual_tokens, - ) + if max_query_len > 1 and not self.use_pcp: + ( + num_decode_reqs, + num_prefill_reqs, + num_decode_tokens, + num_prefill_tokens, + ) = split_dcp_context_queries( + common_attn_metadata.query_start_loc_cpu, + common_attn_metadata.seq_lens_cpu_upper_bound, + max_query_len, + num_actual_tokens, + ) - if skip_dcp_context_attention: - max_dcp_context_kv_len = 0 - else: - scheduler_metadata = schedule( + # After DCP distribution, the maximum number of tokens for any rank is + # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size, + # and I is cp_kv_cache_interleave_size. + # This eliminates GPU->CPU sync while minimizing workspace over-allocation. + if skip_dcp_context_attention: + max_dcp_context_kv_len = 0 + scheduler_metadata = None + else: + num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size + max_dcp_context_kv_len = ( + (max_seq_len + num_partitions - 1) // num_partitions + ) * self.cp_kv_cache_interleave_size + + # Under PCP the DCP group holds replicated query heads, so the + # `num_heads_q * dcp_world_size` sizing schedule() assumes does + # not hold. Let FA schedule itself instead. + scheduler_metadata = ( + None + if self.use_pcp + else schedule( batch_size=num_reqs, cu_query_lens=query_start_loc, max_query_len=max_query_len, @@ -696,6 +642,7 @@ def schedule( max_seq_len=max_dcp_context_kv_len, causal=False, ) + ) elif use_cascade: cu_prefix_query_lens = torch.tensor( [0, num_actual_tokens], dtype=torch.int32, device=self.device @@ -761,7 +708,6 @@ def schedule( slot_mapping=slot_mapping, max_dcp_context_kv_len=max_dcp_context_kv_len, dcp_context_kv_lens=dcp_context_kv_lens, - is_prefilling=common_attn_metadata.is_prefilling, num_decode_reqs=num_decode_reqs, num_prefill_reqs=num_prefill_reqs, num_decode_tokens=num_decode_tokens, @@ -904,22 +850,31 @@ def __init__( self.supports_quant_query_input = flash_attn_supports_quant_query_input() vllm_config = get_current_vllm_config_or_none() - # self.pcp_world_size / self.pcp_rank are auto-populated by - # AttentionImplBase.__new__ from get_pcp_group(). + dcp_a2a = ( + vllm_config is not None + and vllm_config.parallel_config.decode_context_parallel_size > 1 + and vllm_config.parallel_config.dcp_comm_backend == "a2a" + ) + # self.pcp_world_size is populated by AttentionImplBase.__new__. self.use_pcp = self.pcp_world_size > 1 - # Per-layer write-gathered K/V from do_kv_cache_update, reused by the - # suffix attention in _forward_pcp_dcp. Keyed by layer_name because this - # impl is shared across the layers of an attention group. - self._pcp_gathered_kv: dict[str, tuple[torch.Tensor, torch.Tensor] | None] = {} - # How Q and the partials move across the DCP group is decided by the - # PCP/DCP topology alone; both backends share the rule (see pcp.py). - self.dcp_combine = resolve_dcp_combine_fn(vllm_config) + if dcp_a2a: + self.dcp_combine = dcp_a2a_lse_reduce + elif self.use_pcp: + # Under PCP the DCP group spans the PCP ranks, so every rank holds + # the full head set: all-reduce instead of reduce-scatter. + self.dcp_combine = cp_lse_ag_out_ar + else: + self.dcp_combine = cp_lse_ag_out_rs if self.use_pcp and self.dcp_world_size > 1: assert self.dcp_world_size == self.pcp_world_size, ( - "FlashAttention MRv2 PCP+DCP requires dcp == pcp (the DCP " - "group spans exactly the PCP ranks), got " - f"dcp={self.dcp_world_size}, pcp={self.pcp_world_size}." + "FlashAttention PCP+DCP requires dcp == pcp (the DCP group spans " + f"exactly the PCP ranks), got dcp={self.dcp_world_size}, " + f"pcp={self.pcp_world_size}." ) + # Write-gathered K/V from do_kv_cache_update, reused by the new-token + # pass of _forward_pcp_dcp. Keyed by layer because this impl is shared + # across the layers of an attention group. + self._pcp_kv: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} self._dcp_dtype: torch.dtype | None = None self._dcp_max_num_tokens: int = 0 @@ -1039,20 +994,14 @@ def forward( v_descale = layer._v_scale.expand(descale_shape) if self.dcp_world_size > 1: - # Sharded (1/dcp) KV cache. One path covers pure DCP and - # PCP+DCP decode batches (rows replicated across ranks, so the - # topology helpers make the Q gather a no-op and the combine an - # all-reduce); only a PCP+DCP prefill/mixed step -- where - # DualChunkSwap partitions the queries and the manager builds - # a row plan -- goes to _forward_pcp_dcp. - plan = attn_metadata.pcp_row_plan if self.use_pcp else None + # Sharded (1/dcp) KV cache. A PCP+DCP step whose global batch + # has a cached context needs the queries replicated across the + # group before the LSE combine; everything else (pure DCP, and + # PCP steps with no cached context) attends its own rows. + plan = attn_metadata.pcp_plan if plan is not None: - # MRv2 PCP+DCP prefill (dcp == pcp): every row is a - # (cached prefix, new-token span) pair -- the suffix - # attends the write-gathered new K/V, the prefix attends - # the local cache shard + LSE-combine. self._forward_pcp_dcp( - query[:num_actual_tokens], + query, key[:num_actual_tokens], value[:num_actual_tokens], key_cache, @@ -1065,23 +1014,21 @@ def forward( k_descale=k_descale, v_descale=v_descale, ) - else: - self._forward_with_dcp( - query[:num_actual_tokens], - key[:num_actual_tokens], - value[:num_actual_tokens], - key_cache, - value_cache, - output[:num_actual_tokens], - attn_metadata, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - ) + return output + self._forward_with_dcp( + query[:num_actual_tokens], + key[:num_actual_tokens], + value[:num_actual_tokens], + key_cache, + value_cache, + output[:num_actual_tokens], + attn_metadata, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) return output else: - # dcp_world_size <= 1 (no DCP): pure PCP or unparallelized -- - # the normal FA path is correct for both decode and prefill. window = ( attn_metadata.sliding_window if attn_metadata.sliding_window is not None @@ -1218,24 +1165,6 @@ def forward( ) return output - def _get_attn_metadata_for_layer( - self, layer: torch.nn.Module - ) -> FlashAttentionMetadata | None: - """Fetch this layer's FlashAttentionMetadata from the forward context. - - ``do_kv_cache_update`` does not receive ``attn_metadata``, so under PCP - we look it up by ``layer_name`` (the same key the runner uses to store - per-layer metadata in ``forward_context.attn_metadata``). - """ - if self.pcp_world_size <= 1 or not is_forward_context_available(): - return None - attn_metadata_map = get_forward_context().attn_metadata - layer_name = getattr(layer, "layer_name", None) - if not isinstance(attn_metadata_map, dict) or layer_name is None: - return None - meta = attn_metadata_map.get(layer_name) - return meta if isinstance(meta, FlashAttentionMetadata) else None - def do_kv_cache_update( self, layer: torch.nn.Module, @@ -1249,53 +1178,26 @@ def do_kv_cache_update( # we use direct Q, K, V tensors without caching return - if self.use_pcp: - # All-gather the prefill K/V across PCP ranks (decode writes stay - # local) so every rank's cache gets the full prefill KV. K/V are - # gathered separately to stay contiguous (cache kernel head-stride). - attn_metadata = self._get_attn_metadata_for_layer(layer) - num_decode_tokens = ( - attn_metadata.num_decode_tokens if attn_metadata is not None else 0 - ) - # Rank-invariant gather decision: per-rank num_decode_tokens can differ - # under DualChunkSwap and desync the all-gather, so if the global batch - # has any prefill every rank gathers the whole batch. - if attn_metadata is not None and attn_metadata.pcp_has_prefill: - num_decode_tokens = 0 - key_cache, value_cache = kv_cache.transpose(1, 2).split( - self.head_size, dim=-1 - ) - (cache_key, cache_value), cache_slot_mapping = maybe_gather_cache_inputs( - (key, value), - slot_mapping, - num_decode_tokens, - self.use_pcp, - ) - # Stash the write-gathered K/V for the suffix attention of the - # sharded PCP+DCP path (a gather only happens on prefill steps). - # Keyed by layer on the impl (shared across an attention group), so a - # layer whose cache update is skipped (kv sharing) cannot pick up - # stale tensors. - self._pcp_gathered_kv[layer.layer_name] = ( - (cache_key, cache_value) if num_decode_tokens == 0 else None - ) - reshape_and_cache_flash( - cache_key, - cache_value, - key_cache, - value_cache, - cache_slot_mapping, - self.kv_cache_dtype, - layer._k_scale, - layer._v_scale, - ) - return - # Scatter write into the KV cache using slot_mapping indices. # No TMA kernel is invoked here, so stride canonicalization is not needed. # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D)) key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + if self.use_pcp and slot_mapping.shape[0] > key.shape[0]: + # PCP partitions the new tokens across ranks, so a slot mapping + # wider than this rank's tokens means the manager laid it out for + # the whole (all-gathered) batch: materialize it. Duplicate decode + # rows and padding are already masked to PAD_SLOT_ID by the manager. + # K and V are gathered separately to keep each output contiguous, + # which the cache kernel's head-stride assumes. + num_local_tokens = slot_mapping.shape[0] // self.pcp_world_size + pcp_group = get_pcp_group() + key = pcp_group.all_gather(key[:num_local_tokens], dim=0) + value = pcp_group.all_gather(value[:num_local_tokens], dim=0) + # The new-token pass of _forward_pcp_dcp reuses this gather instead + # of paying for a second one. + self._pcp_kv[layer.layer_name] = (key, value) + # Reshape the input keys and values and store them in the cache. # Skip this if sharing KV cache with an earlier attention layer. # NOTE(woosuk): Here, key and value are padded while slot_mapping is @@ -1314,24 +1216,6 @@ def do_kv_cache_update( layer._v_scale, ) - def _fa_common_kwargs( - self, - q_descale: torch.Tensor | None, - k_descale: torch.Tensor | None, - v_descale: torch.Tensor | None, - ) -> dict: - sw = list(self.sliding_window) if self.sliding_window is not None else None - return dict( - softmax_scale=self.scale, - alibi_slopes=self.alibi_slopes, - window_size=sw, - softcap=self.logits_soft_cap, - fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - ) - def _forward_pcp_dcp( self, query: torch.Tensor, @@ -1342,110 +1226,110 @@ def _forward_pcp_dcp( output: torch.Tensor, attn_metadata: FlashAttentionMetadata, layer: torch.nn.Module, - plan: "PCPRowPlan", + plan: "PCPPlan", q_descale: torch.Tensor | None = None, k_descale: torch.Tensor | None = None, v_descale: torch.Tensor | None = None, ) -> torch.Tensor: - """MRv2 PCP+DCP GQA prefill/mixed attention over a sharded KV cache. - - Every rank-local row (decode or DualChunkSwap prefill chunk) is a - (cached prefix, new-token span) pair and is evaluated as: - - - suffix: causal attention of the row against the write-gathered new - K/V -- the same all-gather the cache write already performs, so no - extra collective, and every rank computes only its own rows; - - prefix: non-causal attention of the PCP-gathered prefix-row - queries against this rank's DCP cache shard, then an LSE - all-reduce combine across the group (same-queries requirement), - sliced back to this rank's rows; - - merge_attn_states folds the two. - - Pure-decode batches do not reach here: no prefill means no row plan, - and forward() routes those steps to _forward_with_dcp. + """PCP + DCP (dcp == pcp) attention over a sharded KV cache. + + Same two-pass split as :meth:`_forward_with_dcp`, but PCP has scattered + the new tokens across ranks: + + - new tokens: this rank's rows attend the write-gathered new K/V (the + all-gather ``do_kv_cache_update`` already paid for), so no collective + and every rank computes only its own rows; + - cached context: the LSE combine requires identical queries on every + rank, so the queries are all-gathered back into global batch order and + attended against this rank's cache shard with the *global* batch's + rows, then sliced back. Replicating the queries costs nothing overall: + the pcp-fold query replication is cancelled by the dcp-fold KV shard. """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." ) - query = query.contiguous() - fa_kw = self._fa_common_kwargs(q_descale, k_descale, v_descale) - - n = query.shape[0] - suffix_lse = None - if n > 0: - # Suffix: local rows vs the write-gathered new K/V. The gather is - # the one do_kv_cache_update already did for the cache write. - gathered_kv = self._pcp_gathered_kv.get(layer.layer_name) - assert gathered_kv is not None, ( - "PCP+DCP prefill step without write-gathered K/V: " - "do_kv_cache_update did not run for this layer (kv sharing?)." - ) - k_g, v_g = gathered_kv - _, suffix_lse = flash_attn_varlen_func( - q=query, - k=k_g[plan.suffix_kv_idx], - v=v_g[plan.suffix_kv_idx], - out=output, - cu_seqlens_q=attn_metadata.query_start_loc, - max_seqlen_q=attn_metadata.max_query_len, - cu_seqlens_k=plan.suffix_cu_k, - max_seqlen_k=plan.suffix_max_k, - causal=attn_metadata.causal, - return_softmax_lse=True, - num_splits=attn_metadata.max_num_splits, - **fa_kw, - ) - prefix = plan.prefix - if prefix is None: - return output + sliding_window = ( + list(self.sliding_window) if self.sliding_window is not None else None + ) + fa_kwargs: dict[str, Any] = dict( + softmax_scale=self.scale, + alibi_slopes=self.alibi_slopes, + window_size=sliding_window, + softcap=self.logits_soft_cap, + return_softmax_lse=True, + fa_version=self.vllm_flash_attn_version, + num_splits=attn_metadata.max_num_splits, + ) - # Prefix: gather the prefix-row queries (extend chunks + decode rows) - # so every rank evaluates the same rows against its DCP cache shard, - # then LSE all-reduce the partials. - num_prefix_rows = prefix.cu_q.shape[0] - 1 - pfx_descale_shape = (num_prefix_rows, self.num_kv_heads) - pfx_fa_kw = self._fa_common_kwargs( - layer._q_scale.expand(pfx_descale_shape) - if self.supports_quant_query_input - else None, - layer._k_scale.expand(pfx_descale_shape), - layer._v_scale.expand(pfx_descale_shape), + num_tokens = output.shape[0] + # Consumed, not just read: a layer that shares an earlier layer's cache + # never runs do_kv_cache_update, and must not pick up a stale gather. + gathered_kv = self._pcp_kv.pop(layer.layer_name, None) + assert gathered_kv is not None, ( + "PCP+DCP step without write-gathered K/V: do_kv_cache_update did " + "not run for this layer (kv sharing?)." ) - if n > 0: - q_local = query[prefix.q_local_idx] - else: - # This rank holds no rows; still must join the collectives. - q_local = query.new_zeros( - (prefix.padded_num_tokens, self.num_heads, self.head_size) - ) - q_g = get_pcp_group().all_gather(q_local, dim=0)[prefix.q_restore_idx] + k_gathered, v_gathered = gathered_kv + _, new_lse = flash_attn_varlen_func( + q=query[:num_tokens].contiguous(), + k=k_gathered[plan.new_kv_idx], + v=v_gathered[plan.new_kv_idx], + out=output, + cu_seqlens_q=attn_metadata.query_start_loc, + max_seqlen_q=attn_metadata.max_query_len, + cu_seqlens_k=plan.new_cu_kv, + max_seqlen_k=plan.new_max_kv, + causal=attn_metadata.causal, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + **fa_kwargs, + ) + + ctx = plan.ctx + if ctx is None: + # Nothing was cached before this step, so the new-token pass above + # already covered every key. Skip the context pass and both of its + # collectives (rank-invariant: the global batch decides). + return output + + assert query.shape[0] >= ctx.padded_num_tokens + query_global = get_pcp_group().all_gather( + query[: ctx.padded_num_tokens], dim=0 + )[ctx.restore_idx] + descale_shape = (ctx.cu_q.shape[0] - 1, self.num_kv_heads) ctx_out, ctx_lse = flash_attn_varlen_func( - q=q_g, + q=query_global, k=key_cache, v=value_cache, - cu_seqlens_q=prefix.cu_q, - max_seqlen_q=prefix.max_q, - seqused_k=prefix.dcp_ctx_lens, - max_seqlen_k=prefix.max_ctx, + cu_seqlens_q=ctx.cu_q, + max_seqlen_q=ctx.max_q, + seqused_k=ctx.ctx_lens, + max_seqlen_k=ctx.max_ctx, causal=False, - block_table=prefix.block_table, - return_softmax_lse=True, - num_splits=attn_metadata.max_num_splits, - **pfx_fa_kw, + block_table=ctx.block_table, + q_descale=layer._q_scale.expand(descale_shape) + if self.supports_quant_query_input + else None, + k_descale=layer._k_scale.expand(descale_shape), + v_descale=layer._v_scale.expand(descale_shape), + **fa_kwargs, ) - ctx_out_cor, ctx_lse_cor = self.dcp_combine( + # Rows whose cache shard holds no context attended nothing; neutralize + # them before the combine so undefined partials cannot poison it. + mask_empty_context(ctx_lse, ctx_out, ctx.cu_q, ctx.ctx_cu) + ctx_out, ctx_lse = self.dcp_combine( ctx_out, ctx_lse.transpose(0, 1), get_dcp_group(), return_lse=True ) - if n == 0: - return output - assert suffix_lse is not None - ctx_lse_cor = ctx_lse_cor.transpose(0, 1).contiguous() - pfx_out = output.new_zeros((n, self.num_heads, self.head_size)) - pfx_lse = suffix_lse.new_full((suffix_lse.shape[0], n), float("-inf")) - pfx_out[prefix.local_token_idx] = ctx_out_cor[prefix.local_out_idx] - pfx_lse[:, prefix.local_token_idx] = ctx_lse_cor[:, prefix.local_out_idx] - # Rows without a prefix row keep their suffix output (lse -inf -> 0 weight). - merge_attn_states(output, pfx_out, pfx_lse, output, suffix_lse) + ctx_lse = ctx_lse.transpose(0, 1).contiguous() + + merge_attn_states( + output, + ctx_out[ctx.local_idx], + ctx_lse[:, ctx.local_idx], + output, + new_lse, + ) return output def _forward_with_dcp( @@ -1461,13 +1345,6 @@ def _forward_with_dcp( k_descale: torch.Tensor | None = None, v_descale: torch.Tensor | None = None, ) -> torch.Tensor: - """Sharded-cache attention for pure DCP and PCP+DCP decode batches. - - Attend the DCP-sharded cached context with the topology-appropriate - query gather (no-op when DCP spans exactly the PCP ranks, i.e. decode - rows are replicated), LSE-combine across the group, then merge with - causal attention over the new tokens. - """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." ) @@ -1503,13 +1380,15 @@ def _forward_with_dcp( ) return output - query_for_context = maybe_all_gather_q_for_dcp( - query, self.dcp_world_size, self.pcp_world_size + # Under PCP the DCP group *is* the PCP group, whose ranks already hold + # every query head, so the head gather is a no-op. + query_across_dcp = ( + query if self.use_pcp else get_dcp_group().all_gather(query, dim=1) ) sliding_window_size = ( list(self.sliding_window) if self.sliding_window is not None else None ) - n, context_num_heads = query_for_context.shape[:2] + n, context_num_heads = query_across_dcp.shape[:2] num_reqs = cu_seqlens_q.shape[0] - 1 num_decodes = attn_metadata.num_decode_reqs num_context_prefills = attn_metadata.num_prefill_reqs @@ -1540,14 +1419,14 @@ def _forward_with_dcp( # TODO: Remove this DCP + FA2 mixed decode/prefill workaround once # FA4 supports this Qwen3.5 shape. assert not self.use_pcp, ( - "FA2 split-DCP context path does not support PCP+DCP; use FA3/FA4." + "The FA2 split-DCP context path does not support PCP; use FA3/FA4." ) assert attn_metadata.dcp_context_kv_lens is not None assert attn_metadata.max_dcp_context_kv_len is not None assert self.vllm_flash_attn_version is not None context_attn_out, context_lse = run_split_fa2_dcp_context_attention( flash_attn_varlen_func, - query_for_context, + query_across_dcp, key_cache, value_cache, dcp_context_out, @@ -1574,7 +1453,7 @@ def _forward_with_dcp( ) else: context_attn_out, context_lse = flash_attn_varlen_func( - q=query_for_context, + q=query_across_dcp, k=key_cache, v=value_cache, out=dcp_context_out, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 8580aadca1cf..bd3fb5d2df32 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1295,6 +1295,8 @@ def execute_model( # indices from the previous real batch. for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, ) + if self.pcp_manager is not None: + self.pcp_manager.populate_attn_metadata(attn_metadata) input_ids = input_batch.input_ids inputs_embeds = None @@ -1380,8 +1382,6 @@ def execute_model( skip_compiled=skip_compiled, is_padding=input_batch.is_padding, ): - if self.pcp_manager is not None: - self.pcp_manager.populate_attn_metadata(attn_metadata) self.kv_connector.pre_forward(scheduler_output) if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE: # Run the PIECEWISE graph (compiled PW cudagraph or breakable diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 9b8dc98b6e39..d8e4f1c4afb6 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -36,43 +36,40 @@ def num_tokens(self) -> int: return self.global_batch_slice.stop - self.global_batch_slice.start -class PCPPrefixPlan(NamedTuple): - """Prefix (cached-context) attention plan for one PCP+DCP prefill step. - - Every rank gathers the same prefix-row queries (extend chunks + decode - rows), attends its DCP shard of the cached prefix, and LSE-combines - (all-reduce) across the group; the combine requires the same queries on - every rank. Index tensors are views into one packed buffer uploaded with - a single H2D copy per step. +class PCPContextPlan(NamedTuple): + """Cached-context attention plan for one PCP+DCP step. + + The LSE combine requires identical queries on every rank, so the queries + are all-gathered back into global batch order and the *global* batch's rows + are attended against this rank's cache shard. Every index here is either + the global batch's own metadata or one of the two permutations the manager + already maintains for the hidden states. """ padded_num_tokens: int # all-gather slab size (rank-invariant) - q_local_idx: torch.Tensor # [padded_num_tokens] int64, local token coords - q_restore_idx: torch.Tensor # [total_global_prefix_tokens] int64 - cu_q: torch.Tensor # [num_global_prefix_rows + 1] int32 + restore_idx: torch.Tensor # [num_global_tokens] int64, slab -> global order + local_idx: torch.Tensor # [num_local_tokens] int64, global order -> local + cu_q: torch.Tensor # [num_global_reqs + 1] int32 max_q: int - dcp_ctx_lens: torch.Tensor # [num_global_prefix_rows] int32, this DCP rank + ctx_lens: torch.Tensor # [num_global_reqs] int32, this DCP rank's shard + ctx_cu: torch.Tensor # [num_global_reqs + 1] int32 max_ctx: int - block_table: torch.Tensor # [num_global_prefix_rows, max_num_blocks] - local_token_idx: torch.Tensor # [num_local_prefix_tokens] int64 - local_out_idx: torch.Tensor # [num_local_prefix_tokens] int64 + block_table: torch.Tensor # [num_global_reqs, max_num_blocks] -class PCPRowPlan(NamedTuple): - """Per-step row plan for the sharded PCP+DCP prefill path. +class PCPPlan(NamedTuple): + """Per-step plan for the sharded PCP+DCP attention path. - Every rank-local row is a (cached prefix, new-token span) pair. The suffix - (new tokens) attends the write-gathered K/V -- the same all-gather the - cache write already performs -- so it needs no collective. ``prefix`` is - None when no global row has a cached context this step (pure - fresh-prefill batch), letting the prefix attention and its query gather - be skipped on every rank. + The new tokens of the step are fully materialized on every rank by the + cache-write all-gather, so each rank attends its own rows against them with + no collective; ``ctx`` covers the cached prefix and is None when nothing was + cached this step (then the new-token pass alone is exact). """ - suffix_kv_idx: torch.Tensor # [total_suffix_k] int64, into stashed K/V - suffix_cu_k: torch.Tensor # [num_local_rows + 1] int32 - suffix_max_k: int - prefix: PCPPrefixPlan | None + new_kv_idx: torch.Tensor # [total_new_kv] int64, into the gathered slab + new_cu_kv: torch.Tensor # [num_local_rows + 1] int32 + new_max_kv: int + ctx: PCPContextPlan | None class PCPManager: @@ -118,33 +115,29 @@ def __init__( self._global_has_prefill: bool = False self._pad_slot_id = torch.tensor(PAD_SLOT_ID, dtype=torch.int64, device=device) - # Per-step row-plan state for the sharded PCP+DCP attention path. - # ``_segments_by_rank``/``_hidden_restore_idx_np`` are captured at - # partition time; ``build_cp_row_plan`` derives the plan consumed by - # FlashAttentionImpl._forward_pcp_dcp (see PCPRowPlan). - self._segments_by_rank: list[list[RankSegment]] | None = None + # Per-step state for the sharded PCP+DCP attention path, captured at + # partition time and consumed by ``build_plan`` (see PCPPlan). + self._local_segments: list[RankSegment] = [] self._hidden_restore_idx_np: np.ndarray | None = None - # Global prefix rows are decode rows plus every extend chunk row, so a - # request contributes up to 2*pcp rows. Size the gather buffers for the - # worst case. - max_num_prefix_rows = ( - 2 * pcp_world_size * max_num_reqs if max_num_reqs is not None else None - ) - self._prefix_block_tables: tuple[torch.Tensor, ...] | None = ( + self._padded_num_tokens: int = 0 + # The cached-context pass runs on the global batch's rows, so it needs + # its own block tables: the ones prepare_attn() gathers describe the + # rank-local DualChunkSwap rows. + self._global_block_tables: tuple[torch.Tensor, ...] | None = ( tuple( - table.new_zeros((max_num_prefix_rows, table.shape[1])) + table.new_zeros((max_num_reqs, table.shape[1])) for table in block_tables.input_block_tables ) - if block_tables is not None and max_num_prefix_rows is not None + if block_tables is not None and max_num_reqs is not None else None ) - self._prefix_block_table_ptrs: torch.Tensor | None = ( + self._global_block_table_ptrs: torch.Tensor | None = ( torch.tensor( - [t.data_ptr() for t in self._prefix_block_tables], + [t.data_ptr() for t in self._global_block_tables], dtype=torch.uint64, device=device, ) - if self._prefix_block_tables is not None + if self._global_block_tables is not None else None ) @@ -385,6 +378,7 @@ def _build_batch_layout( ) self._hidden_restore_idx_np = hidden_restore_idx + self._padded_num_tokens = padded_num_tokens self._hidden_restore_idx = async_copy_to_gpu( hidden_restore_idx, device=self.device ) @@ -419,7 +413,6 @@ def partition_batch(self, input_batch: InputBatch) -> InputBatch: is_prefilling, global_batch.query_start_loc_np, ) - self._segments_by_rank = segments_by_rank local_segments = segments_by_rank[self.pcp_rank] if not local_segments: @@ -430,6 +423,7 @@ def partition_batch(self, input_batch: InputBatch) -> InputBatch: rank_local_batch_slice=slice(0, 0), ) ] + self._local_segments = local_segments num_local_reqs = len(local_segments) if num_local_reqs > input_buffers.max_num_reqs: @@ -640,6 +634,14 @@ def prepare_attn( return block_tables, slot_mappings def prepare_slot_mappings(self) -> torch.Tensor: + """Slot mappings for the cache write, in the layout the write expects. + + The width doubles as the protocol with the attention backend: a mapping + wider than the rank's own tokens means the new K/V must be all-gathered + across PCP before writing (prefill chunks live on different ranks). A + pure-decode step needs no gather -- its rows are replicated -- so it + gets the plain rank-local mapping. + """ assert self._block_tables is not None assert self._global_batch_slot_mappings is not None assert self._global_batch is not None @@ -651,6 +653,8 @@ def prepare_slot_mappings(self) -> torch.Tensor: global_batch.num_tokens, out=self._global_batch_slot_mappings, ) + if not self._global_has_prefill: + return global_batch_slot_mappings return self._convert_to_gathered_slot_mappings(global_batch_slot_mappings) def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor: @@ -694,258 +698,120 @@ def restore_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: return gathered[self._hidden_restore_idx] def populate_attn_metadata(self, attn_metadata: dict | None) -> None: - """Stamp the PCP per-step state onto each GQA attention metadata object. - - Centralized here so the model runner stays a single call. Sets, on every - FlashAttentionMetadata in ``attn_metadata``: - - ``pcp_has_prefill``: rank-invariant bool gating the cache-write - all-gather (when the global batch has any prefill, every rank - gathers the whole batch; otherwise writes stay local). - - ``pcp_row_plan``: per-step :class:`PCPRowPlan` for the sharded - PCP+DCP prefill path (None for warmup or pure-decode steps, which - is also how forward() tells the two paths apart). - MLA uses its own metadata class and has no row plan, so non-GQA objects - are skipped. + """Stamp this step's :class:`PCPPlan` onto each GQA attention metadata. + + Left unset (None) when the step needs no PCP-specific attention: warmup, + a replicated pure-decode batch, or a non-sharded cache. MLA uses its own + metadata class and its own PCP path, so it is skipped. """ if not attn_metadata: return - plan = self.build_cp_row_plan() + plan = self.build_plan() + if plan is None: + return for meta in attn_metadata.values(): if isinstance(meta, FlashAttentionMetadata): - meta.pcp_row_plan = plan - meta.pcp_has_prefill = self._global_has_prefill + meta.pcp_plan = plan - def build_cp_row_plan(self) -> "PCPRowPlan | None": - """Build the per-step row plan for the sharded PCP+DCP attention path. + def build_plan(self) -> PCPPlan | None: + """Build the per-step plan for the sharded PCP+DCP attention path. - Returns None when the plan is not needed: no global batch (warmup), - the cache is not DCP-sharded, or the global batch has no prefill - (decode rows are replicated across ranks, so the local-metadata path - handles them without any gather). Derived entirely from partition-time - CPU state; rank-invariant except the DCP-sharded context lengths. + Returns None when the rank-local batch is already the global batch as + far as attention is concerned: no global batch (warmup), an unsharded + cache, or a pure-decode step (decode rows are replicated on every rank, + so the ordinary DCP path handles them). """ gb = self._global_batch - segments_by_rank = self._segments_by_rank - hidden_restore_idx = self._hidden_restore_idx_np + restore_idx_np = self._hidden_restore_idx_np if ( gb is None - or segments_by_rank is None - or hidden_restore_idx is None + or restore_idx_np is None or self.dcp_world_size <= 1 or not self._global_has_prefill ): return None assert self._block_tables is not None - assert self._prefix_block_tables is not None - assert self._prefix_block_table_ptrs is not None - - pcp = self.pcp_world_size - num_chunks = 2 * pcp - rank = self.pcp_rank - num_computed = gb.num_computed_tokens_np - is_prefilling = gb.is_prefilling_np - query_start_loc = gb.query_start_loc_np - num_scheduled = gb.num_scheduled_tokens - num_reqs = gb.num_reqs - - def chunk_size(j: int) -> int: - return -(-int(num_scheduled[j]) // num_chunks) - - def seg_chunk_idx(seg: RankSegment) -> int: - j = seg.global_batch_req_idx - if not bool(is_prefilling[j]): - return -1 - chunk_offset = seg.global_batch_slice.start - int(query_start_loc[j]) - return chunk_offset // chunk_size(j) - - def is_prefix_seg(seg: RankSegment) -> bool: - # Prefix rows are exactly the rows with a cached context. - return ( - seg.num_tokens > 0 and int(num_computed[seg.global_batch_req_idx]) > 0 - ) - - local_segments = segments_by_rank[rank] or [ - RankSegment(0, slice(0, 0), slice(0, 0)) - ] + assert self._hidden_restore_idx is not None + assert self._padded_gather_idx is not None - # -- Suffix: per local row, the request's new tokens up to the row end, - # addressed into the stashed write-gathered K/V (padded-slab layout). - suffix_idx_parts = [] - suffix_lens = [] - for seg in local_segments: - j = seg.global_batch_req_idx - req_start = int(query_start_loc[j]) + num_reqs = gb.num_reqs + query_start_loc = gb.query_start_loc_np + num_computed = gb.num_computed_tokens_np[:num_reqs] + + # New tokens: each local row attends its request's new tokens up to the + # row's end, addressed into the write-gathered slab. Rows of the same + # request overlap, so the spans are materialized rather than sliced. + new_kv_parts = [] + new_kv_lens = [] + for seg in self._local_segments: + req_start = int(query_start_loc[seg.global_batch_req_idx]) row_end = seg.global_batch_slice.stop - suffix_lens.append(row_end - req_start) - suffix_idx_parts.append(hidden_restore_idx[req_start:row_end]) - suffix_kv_idx_np = ( - np.concatenate(suffix_idx_parts) - if suffix_idx_parts - else np.empty(0, dtype=np.int64) - ) - suffix_cu_np = np.zeros(len(local_segments) + 1, dtype=np.int32) - np.cumsum(np.asarray(suffix_lens, dtype=np.int32), out=suffix_cu_np[1:]) - suffix_max_k = max(suffix_lens, default=0) - - # -- Prefix rows: slab offsets of every rank's gather contribution. - # Each rank gathers its prefix-row tokens in local-row order; slab r - # starts at r * padded_num_prefix_tokens in the gathered buffer. - slab_offsets: list[dict[tuple[int, int], int]] = [] - per_rank_prefix_tokens = [] - for r in range(pcp): - offsets = {} - offset = 0 - for seg in segments_by_rank[r]: - if not is_prefix_seg(seg): - continue - offsets[(seg.global_batch_req_idx, seg_chunk_idx(seg))] = offset - offset += seg.num_tokens - slab_offsets.append(offsets) - per_rank_prefix_tokens.append(offset) - padded_num_prefix = max(per_rank_prefix_tokens, default=0) - - # Canonical global prefix-row order: decode rows (request order), then - # extend chunk rows (request order, chunk index order). - global_prefix_rows: list[tuple[int, int, int]] = [] # (req, chunk, len) - for j in range(num_reqs): - if not bool(is_prefilling[j]) and int(num_computed[j]) > 0: - global_prefix_rows.append((j, -1, int(num_scheduled[j]))) - for j in range(num_reqs): - if bool(is_prefilling[j]) and int(num_computed[j]) > 0: - cs = chunk_size(j) - for c in range(num_chunks): - chunk_len = min(cs, int(num_scheduled[j]) - c * cs) - if chunk_len > 0: - global_prefix_rows.append((j, c, chunk_len)) - has_prefix = len(global_prefix_rows) > 0 - - restore_parts = [] - prefix_cu_np = np.zeros(len(global_prefix_rows) + 1, dtype=np.int32) - row_req_np = np.zeros(len(global_prefix_rows), dtype=np.int64) - row_ctx_np = np.zeros(len(global_prefix_rows), dtype=np.int32) - for i, (j, c, row_len) in enumerate(global_prefix_rows): - # Chunk c lives on the rank r with r == c or r == num_chunks-1-c. - holding_rank = 0 if c < 0 else min(c, num_chunks - 1 - c) - slab_start = slab_offsets[holding_rank][(j, c)] - flat_start = holding_rank * padded_num_prefix + slab_start - restore_parts.append( - np.arange(flat_start, flat_start + row_len, dtype=np.int64) - ) - prefix_cu_np[i + 1] = prefix_cu_np[i] + row_len - row_req_np[i] = j - row_ctx_np[i] = num_computed[j] - prefix_q_restore_np = ( - np.concatenate(restore_parts) - if restore_parts - else np.empty(0, dtype=np.int64) - ) - prefix_max_q = max((row_len for _, _, row_len in global_prefix_rows), default=0) - dcp_ctx_np = get_dcp_local_seq_lens( - torch.from_numpy(row_ctx_np), - self.dcp_world_size, - self.dcp_rank, - self.cp_interleave, - ).numpy() - prefix_max_ctx = int(dcp_ctx_np.max()) if len(dcp_ctx_np) else 0 - - # Local prefix rows: where their tokens sit in the local batch and in - # the combined (global prefix-token order) prefix attention output. - row_out_start = { - (j, c): int(prefix_cu_np[i]) - for i, (j, c, _) in enumerate(global_prefix_rows) - } - local_token_parts = [] - local_out_parts = [] - for seg in local_segments: - if not is_prefix_seg(seg): - continue - out_start = row_out_start[(seg.global_batch_req_idx, seg_chunk_idx(seg))] - local_token_parts.append( - np.arange( - seg.rank_local_batch_slice.start, - seg.rank_local_batch_slice.stop, - dtype=np.int64, - ) - ) - local_out_parts.append( - np.arange(out_start, out_start + seg.num_tokens, dtype=np.int64) + new_kv_parts.append(restore_idx_np[req_start:row_end]) + new_kv_lens.append(row_end - req_start) + new_kv_idx_np = ( + np.concatenate(new_kv_parts) if new_kv_parts else np.empty(0, np.int64) + ) + new_cu_np = np.zeros(len(new_kv_lens) + 1, dtype=np.int32) + np.cumsum(np.asarray(new_kv_lens, dtype=np.int32), out=new_cu_np[1:]) + + # Cached context: the global batch's rows against this rank's shard. + # Nothing cached anywhere means the new-token pass alone is exact, and + # the whole (collective-bearing) context pass can be skipped -- a + # rank-invariant decision, since the global batch is. + ctx_np = num_computed.astype(np.int32) + has_ctx = bool(ctx_np.any()) + dcp_ctx_np = ( + get_dcp_local_seq_lens( + torch.from_numpy(ctx_np), + self.dcp_world_size, + self.dcp_rank, + self.cp_interleave, ) - prefix_local_token_np = ( - np.concatenate(local_token_parts) - if local_token_parts - else np.empty(0, dtype=np.int64) - ) - prefix_local_out_np = ( - np.concatenate(local_out_parts) - if local_out_parts - else np.empty(0, dtype=np.int64) - ) - num_local_prefix = len(prefix_local_token_np) - prefix_q_local_np = np.zeros(padded_num_prefix, dtype=np.int64) - prefix_q_local_np[:num_local_prefix] = prefix_local_token_np - - # Pack the index arrays into one int64 and one int32 buffer -- one - # H2D copy each; a pin_memory() per small tensor is the dominant cost - # here. The plan hands out views into the packed buffers. - num_prefix_rows = len(global_prefix_rows) - blob_a = async_copy_to_gpu( - np.concatenate( - ( - suffix_kv_idx_np, - prefix_q_local_np, - prefix_q_restore_np, - prefix_local_token_np, - prefix_local_out_np, - row_req_np, - ) - ), - device=self.device, - ) - blob_b = async_copy_to_gpu( - np.concatenate((suffix_cu_np, prefix_cu_np, dcp_ctx_np)), - device=self.device, - ) - n_sfx = len(suffix_kv_idx_np) - n_restore = len(prefix_q_restore_np) - suffix_kv_idx = blob_a[:n_sfx] - prefix_q_local_idx = blob_a[n_sfx : n_sfx + padded_num_prefix] - q_restore_off = n_sfx + padded_num_prefix - prefix_q_restore_idx = blob_a[q_restore_off : q_restore_off + n_restore] - loc_tok_off = q_restore_off + n_restore - prefix_local_token_idx = blob_a[loc_tok_off : loc_tok_off + num_local_prefix] - loc_out_off = loc_tok_off + num_local_prefix - prefix_local_out_idx = blob_a[loc_out_off : loc_out_off + num_local_prefix] - row_req_gpu = blob_a[loc_out_off + num_local_prefix :] - n_cu_k = len(suffix_cu_np) - suffix_cu_k = blob_b[:n_cu_k] - prefix_cu_q = blob_b[n_cu_k : n_cu_k + num_prefix_rows + 1] - prefix_dcp_ctx_lens = blob_b[n_cu_k + num_prefix_rows + 1 :] - - prefix_plan = None - if has_prefix: - prefix_block_table = self._block_tables.gather_block_tables( - torch.index_select(gb.idx_mapping, 0, row_req_gpu), - num_prefix_rows, - out=self._prefix_block_tables, - out_ptrs=self._prefix_block_table_ptrs, + .numpy() + .astype(np.int32) + if has_ctx + else np.zeros(0, dtype=np.int32) + ) + ctx_cu_np = np.zeros(len(dcp_ctx_np) + 1, dtype=np.int32) + np.cumsum(dcp_ctx_np, out=ctx_cu_np[1:]) + + new_kv_idx = async_copy_to_gpu(new_kv_idx_np, device=self.device) + int32_blob = async_copy_to_gpu( + np.concatenate((new_cu_np, dcp_ctx_np, ctx_cu_np)), device=self.device + ) + num_cu = len(new_cu_np) + new_cu_kv = int32_blob[:num_cu] + + ctx = None + if has_ctx: + ctx_lens = int32_blob[num_cu : num_cu + num_reqs] + ctx_cu = int32_blob[num_cu + num_reqs :] + block_table = self._block_tables.gather_block_tables( + gb.idx_mapping, + num_reqs, + out=self._global_block_tables, + out_ptrs=self._global_block_table_ptrs, )[0] - prefix_plan = PCPPrefixPlan( - padded_num_tokens=padded_num_prefix, - q_local_idx=prefix_q_local_idx, - q_restore_idx=prefix_q_restore_idx, - cu_q=prefix_cu_q, - max_q=prefix_max_q, - dcp_ctx_lens=prefix_dcp_ctx_lens, - max_ctx=prefix_max_ctx, - block_table=prefix_block_table, - local_token_idx=prefix_local_token_idx, - local_out_idx=prefix_local_out_idx, + rank_start = self.pcp_rank * self._padded_num_tokens + num_local_tokens = sum(seg.num_tokens for seg in self._local_segments) + ctx = PCPContextPlan( + padded_num_tokens=self._padded_num_tokens, + restore_idx=self._hidden_restore_idx, + local_idx=self._padded_gather_idx[ + rank_start : rank_start + num_local_tokens + ], + cu_q=gb.query_start_loc[: num_reqs + 1], + max_q=int(gb.num_scheduled_tokens.max()), + ctx_lens=ctx_lens, + ctx_cu=ctx_cu, + max_ctx=int(dcp_ctx_np.max()), + block_table=block_table, ) - return PCPRowPlan( - suffix_kv_idx=suffix_kv_idx, - suffix_cu_k=suffix_cu_k, - suffix_max_k=suffix_max_k, - prefix=prefix_plan, + return PCPPlan( + new_kv_idx=new_kv_idx, + new_cu_kv=new_cu_kv, + new_max_kv=max(new_kv_lens, default=0), + ctx=ctx, ) def restore_for_sampling( From 872e91b1c7990d76759d5d2b5d962160b1e4b308 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Tue, 4 Aug 2026 02:35:12 +0000 Subject: [PATCH 16/18] Result Signed-off-by: JaredforReal --- hy3_pcp_dcp_benchmark/RESULTS.md | 153 ++++++++++++++++++++++++++ qwen_pcp_dcp_benchmark/RESULTS.md | 173 ++++++++++++++++++++++++++++++ 2 files changed, 326 insertions(+) create mode 100644 hy3_pcp_dcp_benchmark/RESULTS.md create mode 100644 qwen_pcp_dcp_benchmark/RESULTS.md diff --git a/hy3_pcp_dcp_benchmark/RESULTS.md b/hy3_pcp_dcp_benchmark/RESULTS.md new file mode 100644 index 000000000000..a1f7f3884094 --- /dev/null +++ b/hy3_pcp_dcp_benchmark/RESULTS.md @@ -0,0 +1,153 @@ +# Hy3-290B-FP8 PCP/DCP Benchmark + +Model: `tencent/Hy3-FP8` (HYV3ForCausalLM, GQA 64Q/8KV, 192-expert MoE + 1 shared, +80 layers, hidden 4096, 279 GB weights). Hardware: 8×H100 80GB. + +All configs use `--enforce-eager`, `--enable-expert-parallel`, +`--no-enable-flashinfer-autotune` (skips the FlashInfer warmup that OOMs under +TP1 where all 64 attention heads land on one GPU), `--gpu-memory-utilization 0.82` +(headroom for warmup activations; 0.88 OOMs on TP2+PCP4), `--max-model-len 32800` +(TP1+PCP8 has only 12.4 GiB KV cache; 40960 needs 12.5 GiB), `--safetensors-load-strategy +prefetch` (boot 9 min → 2.5 min), `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`. + +Boot date: 2026-08-03. Branch: `pcp-gqa` @ `115e11d54` (rewritten PCP plan — 448 +lines / 6 files, no forward_context / InputBatch / CommonAttentionMetadata changes). + +## Config matrix — all 7 use exactly 8 GPUs + +`TP × PCP = 8` throughout; `DCP == PCP` so DCP ranks share PCP ranks (no extra +world size). EP spans TP × PCP = 8 in every config. + +| # | Config | TP | PCP | DCP | KV cache | +|---|---|---|---|---|---| +| 1 | tp8 | 8 | 1 | 1 | replicated (baseline) | +| 2 | tp4_pcp2 | 4 | 2 | 1 | replicated | +| 3 | tp2_pcp4 | 2 | 4 | 1 | replicated | +| 4 | tp1_pcp8 | 1 | 8 | 1 | replicated (pure PCP) | +| 5 | tp4_pcp2_dcp2_sharded | 4 | 2 | 2 | sharded (1/dcp) | +| 6 | tp2_pcp4_dcp4_sharded | 2 | 4 | 4 | sharded (1/dcp) | +| 7 | tp1_pcp8_dcp8_sharded | 1 | 8 | 8 | sharded (1/dcp) | + +## TPOT — decode latency (mean / P99 ms) + +ShareGPT subset, prefix caching **on** (decode is prefix-insensitive; matches +how the Qwen-30B baselines were taken). Concurrency 1 / 16 / 32. + +| Config | c=1 | c=16 | c=32 | c1→32 | +|---|---|---|---|---| +| TP8 (baseline) | 158 / 161 | 158 / 165 | 162 / 165 | +2.5% | +| TP4+PCP2 | 177 / 181 | 180 / 183 | 180 / 183 | +1.7% | +| TP2+PCP4 | 174 / 180 | 177 / 180 | 178 / 182 | +2.3% | +| TP1+PCP8 | 167 / 171 | 170 / 175 | 174 / 179 | +4.2% | +| TP4+PCP2+DCP2 | 205 / 215 | 213 / 225 | 213 / 219 | +3.9% | +| TP2+PCP4+DCP4 | 205 / 208 | 215 / 220 | 215 / 220 | +4.9% | +| TP1+PCP8+DCP8 | 195 / 198 | 202 / 208 | 206 / 215 | +5.6% | + +## TTFT — prefill latency (median ms) + +Random-input probe, prefix caching **off** + 6 warmup requests discarded, +20 measured. + +| Config | 4K | 16K | 32K | +|---|---|---|---| +| TP8 (baseline) | 256 | 926 | 1973 | +| TP4+PCP2 | 243 | 875 | 1858 | +| TP2+PCP4 | 238 | 832 | 1741 | +| TP1+PCP8 | 226 | 783 | 1682 | +| TP4+PCP2+DCP2 | 244 | 911 | 1979 | +| TP2+PCP4+DCP4 | 233 | 891 | 1980 | +| TP1+PCP8+DCP8 | 228 | 927 | 2123 | + +## Accuracy & KV capacity + +GSM8K 200 questions, 5-shot, greedy, conc 256. + +| Config | GSM8K | KV tokens | DCP capacity win | +|---|---|---|---| +| TP8 | 0.885 | 735,056 | 1× | +| TP4+PCP2 | 0.905 | 338,288 | — | +| TP2+PCP4 | 0.915 | 139,840 | — | +| TP1+PCP8 | 0.940 | 40,736 | — | +| TP4+PCP2+DCP2 | 0.915 | 676,576 | **2.0×** vs TP4+PCP2 | +| TP2+PCP4+DCP4 | 0.895 | 558,814 | **4.0×** vs TP2+PCP4 | +| TP1+PCP8+DCP8 | 0.910 | 324,936 | **8.0×** vs TP1+PCP8 | + +GSM8K spread (0.885–0.940) is 200-question sample noise (each question = 0.5%, +so ±2–3 questions = ±1–1.5%), not a regression. + +## Key findings + +### 1. TPOT is nearly flat across 32× concurrency (+2–6%) + +Same as Qwen-30B. DCP's LSE-combine is a fixed per-step collective that +amortizes across the batch rather than scaling with it. + +### 2. DCP decode cost is a fixed ~28 ms + +| Comparison | Δ TPOT (c=1) | +|---|---| +| TP4+PCP2 (177) → +DCP2 (205) | +28 ms (+16%) | +| TP1+PCP8 (167) → +DCP8 (195) | +28 ms (+17%) | + +DCP2 ≈ DCP4 ≈ DCP8 in TPOT (205 / 215 / 195) — the collective is +latency-bound, not bandwidth-bound, so doubling `dcp` buys 2× more KV +concurrency for ~zero extra decode TPOT. + +### 3. DCP KV capacity win is exactly dcp× + +| Config | KV tokens | × over pure-PCP counterpart | +|---|---|---| +| TP1+PCP8 | 40,736 | 1× | +| TP1+PCP8+DCP8 | 324,936 | **8.0×** | + +Each DCP rank stores 1/dcp of the KV (interleaved by position), so the cluster +holds dcp× the sequences. This is the only path for `dcp_shares_pcp_ranks`. + +### 4. Prefill: PCP accelerates, DCP penalises + +**PCP parallelises prefill** — more PCP ranks = faster prefill: + +| Config | 32K TTFT | +|---|---| +| TP8 | 1973 ms | +| TP1+PCP8 | 1682 ms (−15%) | + +**DCP adds a prefill penalty** — the chunked-prefill extend path: + +| Comparison | Δ TTFT @ 32K | +|---|---| +| TP1+PCP8 (1682) → +DCP8 (2123) | +441 ms (+26%) | + +Same mechanism as Qwen-30B: with the default `max_num_batched_tokens=8192`, a +32K prefill splits into 4 chunks; chunks 2–4 are extends that run the replicated +prefix pass (all-gather prefix-Q across PCP + LSE-combine). Disabling chunking +removes this penalty (verified on Qwen-30B: 764 → 567 ms). + +### 5. Scaling to 500B-class GQA + +This 290B GQA model is the near-500B-class case. The results confirm: + +- **DCP8 delivers 8× KV capacity for +17% decode / +26% prefill overhead.** +- **Scaling dcp 2→4→8 costs ~zero extra decode TPOT** (205 / 215 / 195 ms), + while KV capacity scales linearly. +- The prefill penalty is chunked-prefill, not the attention mechanism — + fixable with larger `max_num_batched_tokens`. + +For KV-memory-bound, decode-heavy serving of large GQA models, the trade is +favourable: the decode-time price of DCP is small and roughly constant, while +the KV-capacity gain scales with `dcp` for free. + +## Methodology notes + +- **TPOT** (prefix-on) and **TTFT** (prefix-off) come from separate boots with + different prefix-caching settings. This is deliberate: decode is + prefix-insensitive, prefill is not. The per-section flag + (`serve_tpot_prefix_caching` / `ttft_prefix_caching`) is recorded in each + result JSON. +- TPOT prompt counts: `{1: 16, 16: 48, 32: 64}` (reduced from the Qwen-30B + sizing of `{1: 32, 16: 128, 32: 256}` because this model's ~155 ms TPOT makes + longer runs impractical). +- GSM8K reduced to 200 questions (from 1319) for the same reason. +- The `kill_server()` function uses PID-based cleanup via `nvidia-smi + --query-compute-apps` instead of `pkill -f` (which self-matches its own shell + and can leave orphans holding GPU memory). diff --git a/qwen_pcp_dcp_benchmark/RESULTS.md b/qwen_pcp_dcp_benchmark/RESULTS.md new file mode 100644 index 000000000000..94f637c1d937 --- /dev/null +++ b/qwen_pcp_dcp_benchmark/RESULTS.md @@ -0,0 +1,173 @@ +# Qwen3-30B-A3B-FP8 PCP/DCP Benchmark + +Model: `Qwen3-30B-A3B-FP8` (GQA, MoE, 48 layers). Hardware: 4×H100 80GB. + +All configs use `--enforce-eager`, `--enable-expert-parallel`, +`--gpu-memory-utilization 0.88`, `--max-model-len 40960`. Weight load via +default strategy (no prefetch needed — 30 GB loads in ~30 s). + +Boot date: 2026-07-31. Branch: `pcp-gqa` @ `115e11d54` (rewritten PCP plan — 448 +lines / 6 files, down from 927 / 9; no forward_context / InputBatch / +CommonAttentionMetadata changes). + +## Config matrix — all 5 use 4 GPUs + +When `dcp == pcp` the DCP ranks share the PCP ranks, so the CP world size is +just `pcp`. Config 4 (PCP2+DCP2) therefore runs on **2 GPUs**, not 4. + +| # | Config | TP | PCP | DCP | KV cache | GPUs | +|---|---|---|---|---|---|---| +| 1 | tp4 | 4 | 1 | 1 | replicated (baseline) | 4 | +| 2 | tp2_pcp2 | 2 | 2 | 1 | replicated | 4 | +| 3 | tp1_pcp4 | 1 | 4 | 1 | replicated (pure PCP) | 4 | +| 4 | tp1_pcp2_dcp2_sharded | 1 | 2 | 2 | sharded (1/dcp) | **2** | +| 5 | tp1_pcp4_dcp4_sharded | 1 | 4 | 4 | sharded (1/dcp) | 4 | + +## TPOT — decode latency (mean / P99 ms) + +ShareGPT subset, prefix caching **on** (decode is prefix-insensitive). +Concurrency 1 / 16 / 32. + +| Config | c=1 | c=16 | c=32 | c1→32 | +|---|---|---|---|---| +| TP4 (baseline) | 60.2 / 64.5 | 60.9 / 64.1 | 65.7 / 76.5 | +9.1% | +| TP2+PCP2 | 73.6 / 78.9 | 74.7 / 79.1 | 79.5 / 84.9 | +8.0% | +| TP1+PCP4 | 68.6 / 70.2 | 70.2 / 76.1 | 75.9 / 80.6 | +10.6% | +| TP1+PCP2+DCP2 | 83.9 / 87.7 | 89.6 / 95.7 | 94.6 / 100.8 | +12.8% | +| TP1+PCP4+DCP4 | 85.6 / 91.7 | 90.3 / 95.1 | 96.4 / 105.3 | +12.6% | + +### vs. pre-rewrite (same host, same session, A/B against `c65076b0e`) + +| Config | old (c=32) | new (c=32) | Δ | +|---|---|---|---| +| TP4 | 66.4 | 65.7 | −1.1% (noise floor — path untouched) | +| TP2+PCP2 | 87.7 | 79.5 | −9.4% | +| TP1+PCP4 | 85.8 | 75.9 | **−11.5%** (directly A/B'd, reproduced twice within 0.8%) | +| TP1+PCP2+DCP2 | 98.2 | 94.6 | −3.6% | +| TP1+PCP4+DCP4 | 98.4 | 96.4 | −2.1% | + +The PCP gain is from decode steps no longer doing a per-layer `forward_context` +lookup + store, nor building the pcp-wide gathered slot mapping. DCP configs +gain less because their per-step LSE-combine collective dominates what was +removed. + +## TTFT — prefill latency (median ms) + +Random-input probe, prefix caching **off** + 6 warmup requests discarded, +20 measured. + +| Config | 4K | 16K | 32K | +|---|---|---|---| +| TP4 (baseline) | 96 | 289 | 645 | +| TP2+PCP2 | 116 | 301 | 648 | +| TP1+PCP4 | 110 | 287 | 622 | +| TP1+PCP2+DCP2 | 109 | 423 | 1038 | +| TP1+PCP4+DCP4 | 114 | 333 | 771 | + +The three no-DCP configs converge at 32K (622–648 ms — same parallelised +prefill). DCP4 is +126 ms (+20%) over that floor; DCP2 is much worse but runs +on **2 GPUs**. + +### Chunked-prefill penalty breakdown (DCP4 @ 32K, prefix-off, warmup) + +The DCP4 residual at 32K is entirely chunked prefill, not the attention +mechanism — verified by a chunk-size sweep: + +| max_num_batched_tokens | chunks | extend chunks | 32K TTFT | Δ vs floor | +|---|---|---|---|---| +| 8192 (default) | 4 | 3 | 764 ms | +206 ms (+37%) | +| 16384 | 2 | 1 | 668 ms | +110 ms (+20%) | +| 32768 | 1 | 0 | 558 ms | floor | +| 40960 | 1 | 0 | 567 ms | floor | + +Penalty ≈ α × (prefix FLOPs) + β × (num extends), with β ≈ 27 ms/extend (fixed +collective: all-gather prefix-Q + LSE combine) + compute-proportional term. + +## Accuracy & KV capacity + +GSM8K 1319 questions, 5-shot, greedy, conc 256. + +| Config | GSM8K | KV tokens | Max concurrency | DCP capacity win | +|---|---|---|---|---| +| TP4 | 0.901 | 2,657,504 | — | — | +| TP2+PCP2 | 0.905 | 1,302,384 | — | — | +| TP1+PCP4 | 0.886 | 626,016 | — | — | +| TP1+PCP2+DCP2 | 0.896 | 1,132,096 | — | **1.81×** vs TP1+PCP4 | +| TP1+PCP4+DCP4 | 0.891 | 2,504,064 | 61.13× | **4.00×** vs TP1+PCP4 | + +All within the pre-rewrite range (0.884–0.901); the rewrite is accuracy-neutral. + +### Pre-rewrite accuracy comparison + +| Config | before | after | +|---|---|---| +| TP4 | 0.896 | 0.901 | +| TP2+PCP2 | 0.901 | 0.905 | +| TP1+PCP4 | 0.890 | 0.886 | +| TP1+PCP2+DCP2 | 0.884 | 0.896 | +| TP1+PCP4+DCP4 | 0.887 | 0.891 | + +## Key findings + +### 1. TPOT is nearly flat across 32× concurrency (+8–13%) + +DCP's LSE-combine is a fixed per-step collective that amortises across the +batch rather than scaling with it. + +### 2. DCP decode cost is small and roughly constant + +- Pure PCP4 (68.6 ms) → PCP4+DCP4 (85.6 ms) = **+17.0 ms (+25%)** at c=1. +- DCP4 costs ~the same TPOT as DCP2 (96.4 vs 94.6 ms @ c=32), yet holds **4× + the KV capacity**. Doubling `dcp` from 2→4 buys 2× more KV concurrency for + ~zero extra decode TPOT. + +### 3. Sharded DCP gives dcp× more concurrency + +Config 3 (TP1+PCP4, no DCP) → config 5 (TP1+PCP4+DCP4, sharded) — same TP, +same PCP, same 4 GPUs, only the DCP dimension differs: + +| | KV tokens | Max concurrency | ratio | +|---|---|---|---| +| TP1+PCP4 (no DCP) | 626,016 | 15.28× | 1.00× | +| TP1+PCP4+DCP4 sharded | 2,504,064 | 61.13× | **4.00×** | + +### 4. Prefill: single-chunk has no DCP penalty + +Single-chunk prefills (short prompts, or chunking disabled) have **no DCP +penalty** — the refactored prefill is fully parallelised. The residual penalty +seen at 32K under the default chunked-prefill config comes from extend chunks +running the replicated prefix path; disabling chunking removes it (DCP4 32K: +764 → 567 ms ≈ PCP4 floor). + +## Rewrite impact + +The PCP implementation was rewritten from 927 lines / 9 files to **448 lines / +6 files**, with `forward_context`, `InputBatch`, and `CommonAttentionMetadata` +left completely untouched. The key simplification: the prefix (cached-context) +attention pass runs on the **global batch's own rows** and reuses the two +permutations the PCPManager already maintains for hidden-state restore +(`_hidden_restore_idx`, `_padded_gather_idx`), eliminating the custom +`PCPPrefixPlan` / `PCPRowPlan` / slab-offset / holding-rank machinery. + +Two bugs were found and fixed during the rewrite (both from reviewing PR #2 on +the same branch): + +1. **Mixed prefill+decode crash**: the builder's `_dcp_context_kv_lens` buffer + was sized `max_num_seqs` but DualChunkSwap gives up to 2 local rows per + prefilling request → overflow. Fixed by sizing `2 × max_num_seqs`. +2. **Write-gathered K/V staleness**: changed from `get` to `pop` so a + kv-sharing layer can't pick up a stale gather. + +## Methodology notes + +- **TPOT** (prefix-on) and **TTFT** (prefix-off) come from separate boots. + Decode is prefix-insensitive; prefill is not. The per-section flag + (`serve_tpot_prefix_caching` / `ttft_prefix_caching`) is recorded in each + result JSON. +- Prefix caching is ON by default in vLLM; the off-form flag is + `--no-enable-prefix-caching` (BooleanOptionalAction, not `--no-prefix-caching`). +- TPOT prompt counts: `{1: 32, 16: 128, 32: 256}`. +- GSM8K: full 1319 questions. +- The pre-rewrite A/B comparison was measured by checking out `c65076b0e`, + running `tp1_pcp4` in the same session/host, and comparing (reproduced twice + within 0.8%). From 560469c2dea31acd1b5f60296211d3581f002b58 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Tue, 4 Aug 2026 02:43:03 +0000 Subject: [PATCH 17/18] revert Signed-off-by: JaredforReal --- hy3_pcp_dcp_benchmark/RESULTS.md | 153 -------------------------- qwen_pcp_dcp_benchmark/RESULTS.md | 173 ------------------------------ 2 files changed, 326 deletions(-) delete mode 100644 hy3_pcp_dcp_benchmark/RESULTS.md delete mode 100644 qwen_pcp_dcp_benchmark/RESULTS.md diff --git a/hy3_pcp_dcp_benchmark/RESULTS.md b/hy3_pcp_dcp_benchmark/RESULTS.md deleted file mode 100644 index a1f7f3884094..000000000000 --- a/hy3_pcp_dcp_benchmark/RESULTS.md +++ /dev/null @@ -1,153 +0,0 @@ -# Hy3-290B-FP8 PCP/DCP Benchmark - -Model: `tencent/Hy3-FP8` (HYV3ForCausalLM, GQA 64Q/8KV, 192-expert MoE + 1 shared, -80 layers, hidden 4096, 279 GB weights). Hardware: 8×H100 80GB. - -All configs use `--enforce-eager`, `--enable-expert-parallel`, -`--no-enable-flashinfer-autotune` (skips the FlashInfer warmup that OOMs under -TP1 where all 64 attention heads land on one GPU), `--gpu-memory-utilization 0.82` -(headroom for warmup activations; 0.88 OOMs on TP2+PCP4), `--max-model-len 32800` -(TP1+PCP8 has only 12.4 GiB KV cache; 40960 needs 12.5 GiB), `--safetensors-load-strategy -prefetch` (boot 9 min → 2.5 min), `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`. - -Boot date: 2026-08-03. Branch: `pcp-gqa` @ `115e11d54` (rewritten PCP plan — 448 -lines / 6 files, no forward_context / InputBatch / CommonAttentionMetadata changes). - -## Config matrix — all 7 use exactly 8 GPUs - -`TP × PCP = 8` throughout; `DCP == PCP` so DCP ranks share PCP ranks (no extra -world size). EP spans TP × PCP = 8 in every config. - -| # | Config | TP | PCP | DCP | KV cache | -|---|---|---|---|---|---| -| 1 | tp8 | 8 | 1 | 1 | replicated (baseline) | -| 2 | tp4_pcp2 | 4 | 2 | 1 | replicated | -| 3 | tp2_pcp4 | 2 | 4 | 1 | replicated | -| 4 | tp1_pcp8 | 1 | 8 | 1 | replicated (pure PCP) | -| 5 | tp4_pcp2_dcp2_sharded | 4 | 2 | 2 | sharded (1/dcp) | -| 6 | tp2_pcp4_dcp4_sharded | 2 | 4 | 4 | sharded (1/dcp) | -| 7 | tp1_pcp8_dcp8_sharded | 1 | 8 | 8 | sharded (1/dcp) | - -## TPOT — decode latency (mean / P99 ms) - -ShareGPT subset, prefix caching **on** (decode is prefix-insensitive; matches -how the Qwen-30B baselines were taken). Concurrency 1 / 16 / 32. - -| Config | c=1 | c=16 | c=32 | c1→32 | -|---|---|---|---|---| -| TP8 (baseline) | 158 / 161 | 158 / 165 | 162 / 165 | +2.5% | -| TP4+PCP2 | 177 / 181 | 180 / 183 | 180 / 183 | +1.7% | -| TP2+PCP4 | 174 / 180 | 177 / 180 | 178 / 182 | +2.3% | -| TP1+PCP8 | 167 / 171 | 170 / 175 | 174 / 179 | +4.2% | -| TP4+PCP2+DCP2 | 205 / 215 | 213 / 225 | 213 / 219 | +3.9% | -| TP2+PCP4+DCP4 | 205 / 208 | 215 / 220 | 215 / 220 | +4.9% | -| TP1+PCP8+DCP8 | 195 / 198 | 202 / 208 | 206 / 215 | +5.6% | - -## TTFT — prefill latency (median ms) - -Random-input probe, prefix caching **off** + 6 warmup requests discarded, -20 measured. - -| Config | 4K | 16K | 32K | -|---|---|---|---| -| TP8 (baseline) | 256 | 926 | 1973 | -| TP4+PCP2 | 243 | 875 | 1858 | -| TP2+PCP4 | 238 | 832 | 1741 | -| TP1+PCP8 | 226 | 783 | 1682 | -| TP4+PCP2+DCP2 | 244 | 911 | 1979 | -| TP2+PCP4+DCP4 | 233 | 891 | 1980 | -| TP1+PCP8+DCP8 | 228 | 927 | 2123 | - -## Accuracy & KV capacity - -GSM8K 200 questions, 5-shot, greedy, conc 256. - -| Config | GSM8K | KV tokens | DCP capacity win | -|---|---|---|---| -| TP8 | 0.885 | 735,056 | 1× | -| TP4+PCP2 | 0.905 | 338,288 | — | -| TP2+PCP4 | 0.915 | 139,840 | — | -| TP1+PCP8 | 0.940 | 40,736 | — | -| TP4+PCP2+DCP2 | 0.915 | 676,576 | **2.0×** vs TP4+PCP2 | -| TP2+PCP4+DCP4 | 0.895 | 558,814 | **4.0×** vs TP2+PCP4 | -| TP1+PCP8+DCP8 | 0.910 | 324,936 | **8.0×** vs TP1+PCP8 | - -GSM8K spread (0.885–0.940) is 200-question sample noise (each question = 0.5%, -so ±2–3 questions = ±1–1.5%), not a regression. - -## Key findings - -### 1. TPOT is nearly flat across 32× concurrency (+2–6%) - -Same as Qwen-30B. DCP's LSE-combine is a fixed per-step collective that -amortizes across the batch rather than scaling with it. - -### 2. DCP decode cost is a fixed ~28 ms - -| Comparison | Δ TPOT (c=1) | -|---|---| -| TP4+PCP2 (177) → +DCP2 (205) | +28 ms (+16%) | -| TP1+PCP8 (167) → +DCP8 (195) | +28 ms (+17%) | - -DCP2 ≈ DCP4 ≈ DCP8 in TPOT (205 / 215 / 195) — the collective is -latency-bound, not bandwidth-bound, so doubling `dcp` buys 2× more KV -concurrency for ~zero extra decode TPOT. - -### 3. DCP KV capacity win is exactly dcp× - -| Config | KV tokens | × over pure-PCP counterpart | -|---|---|---| -| TP1+PCP8 | 40,736 | 1× | -| TP1+PCP8+DCP8 | 324,936 | **8.0×** | - -Each DCP rank stores 1/dcp of the KV (interleaved by position), so the cluster -holds dcp× the sequences. This is the only path for `dcp_shares_pcp_ranks`. - -### 4. Prefill: PCP accelerates, DCP penalises - -**PCP parallelises prefill** — more PCP ranks = faster prefill: - -| Config | 32K TTFT | -|---|---| -| TP8 | 1973 ms | -| TP1+PCP8 | 1682 ms (−15%) | - -**DCP adds a prefill penalty** — the chunked-prefill extend path: - -| Comparison | Δ TTFT @ 32K | -|---|---| -| TP1+PCP8 (1682) → +DCP8 (2123) | +441 ms (+26%) | - -Same mechanism as Qwen-30B: with the default `max_num_batched_tokens=8192`, a -32K prefill splits into 4 chunks; chunks 2–4 are extends that run the replicated -prefix pass (all-gather prefix-Q across PCP + LSE-combine). Disabling chunking -removes this penalty (verified on Qwen-30B: 764 → 567 ms). - -### 5. Scaling to 500B-class GQA - -This 290B GQA model is the near-500B-class case. The results confirm: - -- **DCP8 delivers 8× KV capacity for +17% decode / +26% prefill overhead.** -- **Scaling dcp 2→4→8 costs ~zero extra decode TPOT** (205 / 215 / 195 ms), - while KV capacity scales linearly. -- The prefill penalty is chunked-prefill, not the attention mechanism — - fixable with larger `max_num_batched_tokens`. - -For KV-memory-bound, decode-heavy serving of large GQA models, the trade is -favourable: the decode-time price of DCP is small and roughly constant, while -the KV-capacity gain scales with `dcp` for free. - -## Methodology notes - -- **TPOT** (prefix-on) and **TTFT** (prefix-off) come from separate boots with - different prefix-caching settings. This is deliberate: decode is - prefix-insensitive, prefill is not. The per-section flag - (`serve_tpot_prefix_caching` / `ttft_prefix_caching`) is recorded in each - result JSON. -- TPOT prompt counts: `{1: 16, 16: 48, 32: 64}` (reduced from the Qwen-30B - sizing of `{1: 32, 16: 128, 32: 256}` because this model's ~155 ms TPOT makes - longer runs impractical). -- GSM8K reduced to 200 questions (from 1319) for the same reason. -- The `kill_server()` function uses PID-based cleanup via `nvidia-smi - --query-compute-apps` instead of `pkill -f` (which self-matches its own shell - and can leave orphans holding GPU memory). diff --git a/qwen_pcp_dcp_benchmark/RESULTS.md b/qwen_pcp_dcp_benchmark/RESULTS.md deleted file mode 100644 index 94f637c1d937..000000000000 --- a/qwen_pcp_dcp_benchmark/RESULTS.md +++ /dev/null @@ -1,173 +0,0 @@ -# Qwen3-30B-A3B-FP8 PCP/DCP Benchmark - -Model: `Qwen3-30B-A3B-FP8` (GQA, MoE, 48 layers). Hardware: 4×H100 80GB. - -All configs use `--enforce-eager`, `--enable-expert-parallel`, -`--gpu-memory-utilization 0.88`, `--max-model-len 40960`. Weight load via -default strategy (no prefetch needed — 30 GB loads in ~30 s). - -Boot date: 2026-07-31. Branch: `pcp-gqa` @ `115e11d54` (rewritten PCP plan — 448 -lines / 6 files, down from 927 / 9; no forward_context / InputBatch / -CommonAttentionMetadata changes). - -## Config matrix — all 5 use 4 GPUs - -When `dcp == pcp` the DCP ranks share the PCP ranks, so the CP world size is -just `pcp`. Config 4 (PCP2+DCP2) therefore runs on **2 GPUs**, not 4. - -| # | Config | TP | PCP | DCP | KV cache | GPUs | -|---|---|---|---|---|---|---| -| 1 | tp4 | 4 | 1 | 1 | replicated (baseline) | 4 | -| 2 | tp2_pcp2 | 2 | 2 | 1 | replicated | 4 | -| 3 | tp1_pcp4 | 1 | 4 | 1 | replicated (pure PCP) | 4 | -| 4 | tp1_pcp2_dcp2_sharded | 1 | 2 | 2 | sharded (1/dcp) | **2** | -| 5 | tp1_pcp4_dcp4_sharded | 1 | 4 | 4 | sharded (1/dcp) | 4 | - -## TPOT — decode latency (mean / P99 ms) - -ShareGPT subset, prefix caching **on** (decode is prefix-insensitive). -Concurrency 1 / 16 / 32. - -| Config | c=1 | c=16 | c=32 | c1→32 | -|---|---|---|---|---| -| TP4 (baseline) | 60.2 / 64.5 | 60.9 / 64.1 | 65.7 / 76.5 | +9.1% | -| TP2+PCP2 | 73.6 / 78.9 | 74.7 / 79.1 | 79.5 / 84.9 | +8.0% | -| TP1+PCP4 | 68.6 / 70.2 | 70.2 / 76.1 | 75.9 / 80.6 | +10.6% | -| TP1+PCP2+DCP2 | 83.9 / 87.7 | 89.6 / 95.7 | 94.6 / 100.8 | +12.8% | -| TP1+PCP4+DCP4 | 85.6 / 91.7 | 90.3 / 95.1 | 96.4 / 105.3 | +12.6% | - -### vs. pre-rewrite (same host, same session, A/B against `c65076b0e`) - -| Config | old (c=32) | new (c=32) | Δ | -|---|---|---|---| -| TP4 | 66.4 | 65.7 | −1.1% (noise floor — path untouched) | -| TP2+PCP2 | 87.7 | 79.5 | −9.4% | -| TP1+PCP4 | 85.8 | 75.9 | **−11.5%** (directly A/B'd, reproduced twice within 0.8%) | -| TP1+PCP2+DCP2 | 98.2 | 94.6 | −3.6% | -| TP1+PCP4+DCP4 | 98.4 | 96.4 | −2.1% | - -The PCP gain is from decode steps no longer doing a per-layer `forward_context` -lookup + store, nor building the pcp-wide gathered slot mapping. DCP configs -gain less because their per-step LSE-combine collective dominates what was -removed. - -## TTFT — prefill latency (median ms) - -Random-input probe, prefix caching **off** + 6 warmup requests discarded, -20 measured. - -| Config | 4K | 16K | 32K | -|---|---|---|---| -| TP4 (baseline) | 96 | 289 | 645 | -| TP2+PCP2 | 116 | 301 | 648 | -| TP1+PCP4 | 110 | 287 | 622 | -| TP1+PCP2+DCP2 | 109 | 423 | 1038 | -| TP1+PCP4+DCP4 | 114 | 333 | 771 | - -The three no-DCP configs converge at 32K (622–648 ms — same parallelised -prefill). DCP4 is +126 ms (+20%) over that floor; DCP2 is much worse but runs -on **2 GPUs**. - -### Chunked-prefill penalty breakdown (DCP4 @ 32K, prefix-off, warmup) - -The DCP4 residual at 32K is entirely chunked prefill, not the attention -mechanism — verified by a chunk-size sweep: - -| max_num_batched_tokens | chunks | extend chunks | 32K TTFT | Δ vs floor | -|---|---|---|---|---| -| 8192 (default) | 4 | 3 | 764 ms | +206 ms (+37%) | -| 16384 | 2 | 1 | 668 ms | +110 ms (+20%) | -| 32768 | 1 | 0 | 558 ms | floor | -| 40960 | 1 | 0 | 567 ms | floor | - -Penalty ≈ α × (prefix FLOPs) + β × (num extends), with β ≈ 27 ms/extend (fixed -collective: all-gather prefix-Q + LSE combine) + compute-proportional term. - -## Accuracy & KV capacity - -GSM8K 1319 questions, 5-shot, greedy, conc 256. - -| Config | GSM8K | KV tokens | Max concurrency | DCP capacity win | -|---|---|---|---|---| -| TP4 | 0.901 | 2,657,504 | — | — | -| TP2+PCP2 | 0.905 | 1,302,384 | — | — | -| TP1+PCP4 | 0.886 | 626,016 | — | — | -| TP1+PCP2+DCP2 | 0.896 | 1,132,096 | — | **1.81×** vs TP1+PCP4 | -| TP1+PCP4+DCP4 | 0.891 | 2,504,064 | 61.13× | **4.00×** vs TP1+PCP4 | - -All within the pre-rewrite range (0.884–0.901); the rewrite is accuracy-neutral. - -### Pre-rewrite accuracy comparison - -| Config | before | after | -|---|---|---| -| TP4 | 0.896 | 0.901 | -| TP2+PCP2 | 0.901 | 0.905 | -| TP1+PCP4 | 0.890 | 0.886 | -| TP1+PCP2+DCP2 | 0.884 | 0.896 | -| TP1+PCP4+DCP4 | 0.887 | 0.891 | - -## Key findings - -### 1. TPOT is nearly flat across 32× concurrency (+8–13%) - -DCP's LSE-combine is a fixed per-step collective that amortises across the -batch rather than scaling with it. - -### 2. DCP decode cost is small and roughly constant - -- Pure PCP4 (68.6 ms) → PCP4+DCP4 (85.6 ms) = **+17.0 ms (+25%)** at c=1. -- DCP4 costs ~the same TPOT as DCP2 (96.4 vs 94.6 ms @ c=32), yet holds **4× - the KV capacity**. Doubling `dcp` from 2→4 buys 2× more KV concurrency for - ~zero extra decode TPOT. - -### 3. Sharded DCP gives dcp× more concurrency - -Config 3 (TP1+PCP4, no DCP) → config 5 (TP1+PCP4+DCP4, sharded) — same TP, -same PCP, same 4 GPUs, only the DCP dimension differs: - -| | KV tokens | Max concurrency | ratio | -|---|---|---|---| -| TP1+PCP4 (no DCP) | 626,016 | 15.28× | 1.00× | -| TP1+PCP4+DCP4 sharded | 2,504,064 | 61.13× | **4.00×** | - -### 4. Prefill: single-chunk has no DCP penalty - -Single-chunk prefills (short prompts, or chunking disabled) have **no DCP -penalty** — the refactored prefill is fully parallelised. The residual penalty -seen at 32K under the default chunked-prefill config comes from extend chunks -running the replicated prefix path; disabling chunking removes it (DCP4 32K: -764 → 567 ms ≈ PCP4 floor). - -## Rewrite impact - -The PCP implementation was rewritten from 927 lines / 9 files to **448 lines / -6 files**, with `forward_context`, `InputBatch`, and `CommonAttentionMetadata` -left completely untouched. The key simplification: the prefix (cached-context) -attention pass runs on the **global batch's own rows** and reuses the two -permutations the PCPManager already maintains for hidden-state restore -(`_hidden_restore_idx`, `_padded_gather_idx`), eliminating the custom -`PCPPrefixPlan` / `PCPRowPlan` / slab-offset / holding-rank machinery. - -Two bugs were found and fixed during the rewrite (both from reviewing PR #2 on -the same branch): - -1. **Mixed prefill+decode crash**: the builder's `_dcp_context_kv_lens` buffer - was sized `max_num_seqs` but DualChunkSwap gives up to 2 local rows per - prefilling request → overflow. Fixed by sizing `2 × max_num_seqs`. -2. **Write-gathered K/V staleness**: changed from `get` to `pop` so a - kv-sharing layer can't pick up a stale gather. - -## Methodology notes - -- **TPOT** (prefix-on) and **TTFT** (prefix-off) come from separate boots. - Decode is prefix-insensitive; prefill is not. The per-section flag - (`serve_tpot_prefix_caching` / `ttft_prefix_caching`) is recorded in each - result JSON. -- Prefix caching is ON by default in vLLM; the off-form flag is - `--no-enable-prefix-caching` (BooleanOptionalAction, not `--no-prefix-caching`). -- TPOT prompt counts: `{1: 32, 16: 128, 32: 256}`. -- GSM8K: full 1319 questions. -- The pre-rewrite A/B comparison was measured by checking out `c65076b0e`, - running `tp1_pcp4` in the same session/host, and comparing (reproduced twice - within 0.8%). From cb7e192d510e9b95ba2b05123803ea09187146d7 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Tue, 4 Aug 2026 05:59:30 +0000 Subject: [PATCH 18/18] merging forward with dcp and forward pcp dcp Signed-off-by: JaredforReal --- vllm/v1/attention/backends/flash_attn.py | 421 ++++++----------------- 1 file changed, 100 insertions(+), 321 deletions(-) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 3038de57b41a..66a6a986cf8f 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -33,7 +33,6 @@ from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.attention.ops.triton_merge_attn_states import mask_empty_context -from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: from vllm.v1.worker.gpu.pcp_manager import PCPPlan @@ -64,10 +63,7 @@ from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.kv_cache_interface import AttentionSpec from vllm.v1.worker.cp_utils import ( - run_split_fa2_dcp_context_attention, should_skip_dcp_context_attention, - should_split_fa2_dcp_context_attention, - split_dcp_context_queries, ) logger = init_logger(__name__) @@ -278,13 +274,6 @@ class FlashAttentionMetadata: # _forward_with_dcp. pcp_plan: "PCPPlan | None" = None - # Split counts for FA2 DCP context attention. num_prefill_* tracks - # context-bearing extend rows; pure prefills do not attend to DCP context. - num_decode_reqs: int = 0 - num_prefill_reqs: int = 0 - num_decode_tokens: int = 0 - num_prefill_tokens: int = 0 - # Optional aot scheduling scheduler_metadata: torch.Tensor | None = None prefix_scheduler_metadata: torch.Tensor | None = None @@ -562,11 +551,6 @@ def schedule( use_cascade = common_prefix_len > 0 max_dcp_context_kv_len = 0 dcp_context_kv_lens = None - num_decode_reqs = 0 - num_prefill_reqs = 0 - num_decode_tokens = 0 - num_prefill_tokens = 0 - cu_prefix_query_lens = None prefix_kv_lens = None suffix_kv_lens = None @@ -585,11 +569,10 @@ def schedule( self._dcp_context_kv_lens[num_reqs:] = 0 dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs] - skip_dcp_context_attention = False - if ( - not self.use_pcp - and common_attn_metadata.seq_lens_cpu_upper_bound is not None - ): + # Skip the context pass entirely when no row has any cached + # context (rank-invariant: derived from the global seq_lens). + skip = False + if common_attn_metadata.seq_lens_cpu_upper_bound is not None: query_lens_cpu = ( common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1] - common_attn_metadata.query_start_loc_cpu[:num_reqs] @@ -598,28 +581,9 @@ def schedule( common_attn_metadata.seq_lens_cpu_upper_bound[:num_reqs] - query_lens_cpu ) - skip_dcp_context_attention = should_skip_dcp_context_attention( - context_kv_lens_cpu - ) + skip = should_skip_dcp_context_attention(context_kv_lens_cpu) - if max_query_len > 1 and not self.use_pcp: - ( - num_decode_reqs, - num_prefill_reqs, - num_decode_tokens, - num_prefill_tokens, - ) = split_dcp_context_queries( - common_attn_metadata.query_start_loc_cpu, - common_attn_metadata.seq_lens_cpu_upper_bound, - max_query_len, - num_actual_tokens, - ) - - # After DCP distribution, the maximum number of tokens for any rank is - # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size, - # and I is cp_kv_cache_interleave_size. - # This eliminates GPU->CPU sync while minimizing workspace over-allocation. - if skip_dcp_context_attention: + if skip: max_dcp_context_kv_len = 0 scheduler_metadata = None else: @@ -627,10 +591,6 @@ def schedule( max_dcp_context_kv_len = ( (max_seq_len + num_partitions - 1) // num_partitions ) * self.cp_kv_cache_interleave_size - - # Under PCP the DCP group holds replicated query heads, so the - # `num_heads_q * dcp_world_size` sizing schedule() assumes does - # not hold. Let FA schedule itself instead. scheduler_metadata = ( None if self.use_pcp @@ -708,10 +668,6 @@ def schedule( slot_mapping=slot_mapping, max_dcp_context_kv_len=max_dcp_context_kv_len, dcp_context_kv_lens=dcp_context_kv_lens, - num_decode_reqs=num_decode_reqs, - num_prefill_reqs=num_prefill_reqs, - num_decode_tokens=num_decode_tokens, - num_prefill_tokens=num_prefill_tokens, use_cascade=use_cascade, common_prefix_len=common_prefix_len, scheduler_metadata=scheduler_metadata, @@ -876,14 +832,6 @@ def __init__( # across the layers of an attention group. self._pcp_kv: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} - self._dcp_dtype: torch.dtype | None = None - self._dcp_max_num_tokens: int = 0 - if vllm_config is not None and self.dcp_world_size > 1: - self._dcp_dtype = vllm_config.model_config.dtype - self._dcp_max_num_tokens = ( - vllm_config.scheduler_config.max_num_batched_tokens - ) - def forward( self, layer: torch.nn.Module, @@ -994,35 +942,15 @@ def forward( v_descale = layer._v_scale.expand(descale_shape) if self.dcp_world_size > 1: - # Sharded (1/dcp) KV cache. A PCP+DCP step whose global batch - # has a cached context needs the queries replicated across the - # group before the LSE combine; everything else (pure DCP, and - # PCP steps with no cached context) attends its own rows. - plan = attn_metadata.pcp_plan - if plan is not None: - self._forward_pcp_dcp( - query, - key[:num_actual_tokens], - value[:num_actual_tokens], - key_cache, - value_cache, - output[:num_actual_tokens], - attn_metadata, - layer, - plan, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - ) - return output - self._forward_with_dcp( - query[:num_actual_tokens], + self._forward_dcp( + query, key[:num_actual_tokens], value[:num_actual_tokens], key_cache, value_cache, output[:num_actual_tokens], attn_metadata, + layer, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, @@ -1216,7 +1144,7 @@ def do_kv_cache_update( layer._v_scale, ) - def _forward_pcp_dcp( + def _forward_dcp( self, query: torch.Tensor, key: torch.Tensor, @@ -1226,28 +1154,30 @@ def _forward_pcp_dcp( output: torch.Tensor, attn_metadata: FlashAttentionMetadata, layer: torch.nn.Module, - plan: "PCPPlan", q_descale: torch.Tensor | None = None, k_descale: torch.Tensor | None = None, v_descale: torch.Tensor | None = None, ) -> torch.Tensor: - """PCP + DCP (dcp == pcp) attention over a sharded KV cache. - - Same two-pass split as :meth:`_forward_with_dcp`, but PCP has scattered - the new tokens across ranks: - - - new tokens: this rank's rows attend the write-gathered new K/V (the - all-gather ``do_kv_cache_update`` already paid for), so no collective - and every rank computes only its own rows; - - cached context: the LSE combine requires identical queries on every - rank, so the queries are all-gathered back into global batch order and - attended against this rank's cache shard with the *global* batch's - rows, then sliced back. Replicating the queries costs nothing overall: - the pcp-fold query replication is cancelled by the dcp-fold KV shard. + """DCP attention over a sharded KV cache: context + new-token split. + + Every step is two non-overlapping attentions merged by LSE: + + - new tokens: causal attention of this rank's rows against the new K/V + (local for pure DCP / decode; write-gathered for PCP+DCP prefill); + - cached context: non-causal attention against this rank's 1/dcp shard + of the KV cache, then LSE-combine across the group. + + ``attn_metadata.pcp_plan`` (non-None only on a PCP+DCP prefill step with + a cached context) selects the non-trivial Q-gather / K-V-source / scatter + variant; when it is None the path is the identity (local rows, head-axis + Q gather or no-op, no scatter). """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." ) + num_tokens = output.shape[0] + query = query.contiguous() + plan = attn_metadata.pcp_plan sliding_window = ( list(self.sliding_window) if self.sliding_window is not None else None ) @@ -1261,24 +1191,31 @@ def _forward_pcp_dcp( num_splits=attn_metadata.max_num_splits, ) - num_tokens = output.shape[0] - # Consumed, not just read: a layer that shares an earlier layer's cache - # never runs do_kv_cache_update, and must not pick up a stale gather. - gathered_kv = self._pcp_kv.pop(layer.layer_name, None) - assert gathered_kv is not None, ( - "PCP+DCP step without write-gathered K/V: do_kv_cache_update did " - "not run for this layer (kv sharing?)." - ) - k_gathered, v_gathered = gathered_kv + # --- New-token pass (causal) --- + if plan is not None: + gathered_kv = self._pcp_kv.pop(layer.layer_name, None) + assert gathered_kv is not None, ( + "PCP+DCP step without write-gathered K/V: do_kv_cache_update " + "did not run for this layer (kv sharing?)." + ) + new_k = gathered_kv[0][plan.new_kv_idx] + new_v = gathered_kv[1][plan.new_kv_idx] + new_cu_k = plan.new_cu_kv + new_max_k = plan.new_max_kv + else: + new_k = key + new_v = value + new_cu_k = attn_metadata.query_start_loc + new_max_k = attn_metadata.max_query_len _, new_lse = flash_attn_varlen_func( - q=query[:num_tokens].contiguous(), - k=k_gathered[plan.new_kv_idx], - v=v_gathered[plan.new_kv_idx], + q=query[:num_tokens], + k=new_k, + v=new_v, out=output, cu_seqlens_q=attn_metadata.query_start_loc, max_seqlen_q=attn_metadata.max_query_len, - cu_seqlens_k=plan.new_cu_kv, - max_seqlen_k=plan.new_max_kv, + cu_seqlens_k=new_cu_k, + max_seqlen_k=new_max_k, causal=attn_metadata.causal, q_descale=q_descale, k_descale=k_descale, @@ -1286,28 +1223,51 @@ def _forward_pcp_dcp( **fa_kwargs, ) - ctx = plan.ctx - if ctx is None: - # Nothing was cached before this step, so the new-token pass above - # already covered every key. Skip the context pass and both of its - # collectives (rank-invariant: the global batch decides). - return output - - assert query.shape[0] >= ctx.padded_num_tokens - query_global = get_pcp_group().all_gather( - query[: ctx.padded_num_tokens], dim=0 - )[ctx.restore_idx] - descale_shape = (ctx.cu_q.shape[0] - 1, self.num_kv_heads) + # --- Context pass (non-causal, sharded cache) --- + # Skip entirely if no row has any cached context (rank-invariant). + if plan is not None: + ctx = plan.ctx + if ctx is None: + return output + else: + ctx = None + if attn_metadata.max_dcp_context_kv_len == 0: + return output + + if plan is not None: + assert ctx is not None + q_ctx = get_pcp_group().all_gather(query[: ctx.padded_num_tokens], dim=0)[ + ctx.restore_idx + ] + cu_q = ctx.cu_q + max_q = ctx.max_q + ctx_lens = ctx.ctx_lens + max_ctx = ctx.max_ctx + ctx_cu = ctx.ctx_cu + block_table = ctx.block_table + else: + assert attn_metadata.dcp_context_kv_lens is not None + assert attn_metadata.max_dcp_context_kv_len is not None + q_ctx = query if self.use_pcp else get_dcp_group().all_gather(query, dim=1) + cu_q = attn_metadata.query_start_loc + max_q = attn_metadata.max_query_len + ctx_lens = attn_metadata.dcp_context_kv_lens + max_ctx = attn_metadata.max_dcp_context_kv_len + ctx_cu = None # not needed — pure DCP uses seqused_k=0 + block_table = attn_metadata.block_table + fa_kwargs["scheduler_metadata"] = attn_metadata.scheduler_metadata + + descale_shape = (cu_q.shape[0] - 1, self.num_kv_heads) ctx_out, ctx_lse = flash_attn_varlen_func( - q=query_global, + q=q_ctx, k=key_cache, v=value_cache, - cu_seqlens_q=ctx.cu_q, - max_seqlen_q=ctx.max_q, - seqused_k=ctx.ctx_lens, - max_seqlen_k=ctx.max_ctx, + cu_seqlens_q=cu_q, + max_seqlen_q=max_q, + seqused_k=ctx_lens, + max_seqlen_k=max_ctx, causal=False, - block_table=ctx.block_table, + block_table=block_table, q_descale=layer._q_scale.expand(descale_shape) if self.supports_quant_query_input else None, @@ -1315,205 +1275,24 @@ def _forward_pcp_dcp( v_descale=layer._v_scale.expand(descale_shape), **fa_kwargs, ) - # Rows whose cache shard holds no context attended nothing; neutralize - # them before the combine so undefined partials cannot poison it. - mask_empty_context(ctx_lse, ctx_out, ctx.cu_q, ctx.ctx_cu) + # Sanitize rows whose cache shard holds no context. Only needed for + # PCP (global-batch context pass where some rows may have 0 context + # and FA leaves undefined output); pure DCP's seqused_k=0 handles + # this internally. + if plan is not None: + mask_empty_context(ctx_lse, ctx_out, cu_q, ctx_cu) ctx_out, ctx_lse = self.dcp_combine( ctx_out, ctx_lse.transpose(0, 1), get_dcp_group(), return_lse=True ) ctx_lse = ctx_lse.transpose(0, 1).contiguous() - merge_attn_states( - output, - ctx_out[ctx.local_idx], - ctx_lse[:, ctx.local_idx], - output, - new_lse, - ) - return output - - def _forward_with_dcp( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - output: torch.Tensor, - attn_metadata: FlashAttentionMetadata, - q_descale: torch.Tensor | None = None, - k_descale: torch.Tensor | None = None, - v_descale: torch.Tensor | None = None, - ) -> torch.Tensor: - assert self.vllm_flash_attn_version is not None, ( - "FlashAttention version not detected." - ) - - cu_seqlens_q = attn_metadata.query_start_loc - max_seqlen_q = attn_metadata.max_query_len - block_table = attn_metadata.block_table - - query = query.contiguous() - if attn_metadata.max_dcp_context_kv_len == 0: - flash_attn_varlen_func( - q=query, - k=key, - v=value, - out=output, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - cu_seqlens_k=cu_seqlens_q, - max_seqlen_k=max_seqlen_q, - softmax_scale=self.scale, - causal=attn_metadata.causal, - alibi_slopes=self.alibi_slopes, - window_size=list(self.sliding_window) - if self.sliding_window is not None - else None, - softcap=self.logits_soft_cap, - return_softmax_lse=True, - fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - num_splits=attn_metadata.max_num_splits, - ) - return output + if plan is not None: + assert ctx is not None + ctx_out = ctx_out[ctx.local_idx] + ctx_lse = ctx_lse[:, ctx.local_idx] - # Under PCP the DCP group *is* the PCP group, whose ranks already hold - # every query head, so the head gather is a no-op. - query_across_dcp = ( - query if self.use_pcp else get_dcp_group().all_gather(query, dim=1) - ) - sliding_window_size = ( - list(self.sliding_window) if self.sliding_window is not None else None - ) - n, context_num_heads = query_across_dcp.shape[:2] - num_reqs = cu_seqlens_q.shape[0] - 1 - num_decodes = attn_metadata.num_decode_reqs - num_context_prefills = attn_metadata.num_prefill_reqs - num_decode_tokens = attn_metadata.num_decode_tokens - num_context_prefill_tokens = attn_metadata.num_prefill_tokens - split_dcp_context = should_split_fa2_dcp_context_attention( - self.vllm_flash_attn_version, - max_seqlen_q, - num_reqs, - num_decodes, - num_context_prefills, - ) - dcp_context_out_tokens = max(n, self._dcp_max_num_tokens) - dcp_context_out_spec = ( - ( - dcp_context_out_tokens, - context_num_heads, - self.head_size, - ), - self._dcp_dtype, - ) - (dcp_context_out_workspace,) = current_workspace_manager().get_simultaneous( - dcp_context_out_spec, - ) - dcp_context_out = dcp_context_out_workspace[:n] - - if split_dcp_context: - # TODO: Remove this DCP + FA2 mixed decode/prefill workaround once - # FA4 supports this Qwen3.5 shape. - assert not self.use_pcp, ( - "The FA2 split-DCP context path does not support PCP; use FA3/FA4." - ) - assert attn_metadata.dcp_context_kv_lens is not None - assert attn_metadata.max_dcp_context_kv_len is not None - assert self.vllm_flash_attn_version is not None - context_attn_out, context_lse = run_split_fa2_dcp_context_attention( - flash_attn_varlen_func, - query_across_dcp, - key_cache, - value_cache, - dcp_context_out, - cu_seqlens_q, - max_seqlen_q, - attn_metadata.dcp_context_kv_lens, - attn_metadata.max_dcp_context_kv_len, - self.scale, - self.alibi_slopes, - sliding_window_size, - block_table, - self.logits_soft_cap, - self.vllm_flash_attn_version, - q_descale, - k_descale, - v_descale, - attn_metadata.max_num_splits, - self.num_heads, - self.dcp_world_size, - num_decodes, - num_context_prefills, - num_decode_tokens, - num_context_prefill_tokens, - ) - else: - context_attn_out, context_lse = flash_attn_varlen_func( - q=query_across_dcp, - k=key_cache, - v=value_cache, - out=dcp_context_out, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - seqused_k=attn_metadata.dcp_context_kv_lens, - max_seqlen_k=attn_metadata.max_dcp_context_kv_len, - softmax_scale=self.scale, - causal=False, - alibi_slopes=self.alibi_slopes, - window_size=sliding_window_size, - block_table=block_table, - softcap=self.logits_soft_cap, - return_softmax_lse=True, - scheduler_metadata=attn_metadata.scheduler_metadata, - fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - num_splits=attn_metadata.max_num_splits, - ) - # FA returns LSE in shape [ H, B ] but DCP combine wants [ B, H ] - context_attn_out_cor, context_lse_cor = self.dcp_combine( - context_attn_out, - context_lse.transpose(0, 1), - get_dcp_group(), - return_lse=True, - ) - context_lse_cor = context_lse_cor.transpose(0, 1).contiguous() - - query_attn_out, query_lse = flash_attn_varlen_func( - q=query, - k=key, - v=value, - out=output, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - cu_seqlens_k=cu_seqlens_q, - max_seqlen_k=max_seqlen_q, - softmax_scale=self.scale, - causal=attn_metadata.causal, - alibi_slopes=self.alibi_slopes, - window_size=sliding_window_size, - softcap=self.logits_soft_cap, - return_softmax_lse=True, - fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - num_splits=attn_metadata.max_num_splits, - ) - assert context_attn_out_cor.shape == query_attn_out.shape - assert context_lse_cor.shape == query_lse.shape - merge_attn_states( - output, - context_attn_out_cor, - context_lse_cor, - query_attn_out, - query_lse, - ) + merge_attn_states(output, ctx_out, ctx_lse, output, new_lse) + return output def _forward_encoder_attention( self,