[MRv2] FlashAttention PCP support for GQA on MRv2

vllm/config/model.py

@@ -1251,7 +1251,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 "

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")
 

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 (
@@ -66,6 +68,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,
@@ -80,6 +112,102 @@ 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 finalize_mla_pcp_decode(
     output: torch.Tensor,
     num_heads: int,