[MRv2] FlashAttention PCP support for GQA on MRv2 - #49564
Open
JaredforReal wants to merge 21 commits into
Open
JaredforReal wants to merge 21 commits into
JaredforReal wants to merge 21 commits into
Conversation
Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: JaredforReal <w13431838023@gmail.com>
4 tasks
Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: JaredforReal <w13431838023@gmail.com>
JaredforReal
marked this pull request as ready for review
July 24, 2026 03:48
JaredforReal
requested review from
ApostaC,
LucasWilkinson,
MatthewBonanni,
ProExpertProg,
WoosukKwon,
alexm-redhat,
heheda12345,
houseroad,
ivanium,
mgoin,
njhill,
orozery,
robertgshaw2-redhat,
tlrmchlsmth,
yewentao256,
youkaichao and
ywang96
as code owners
July 24, 2026 03:48
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
Enables MRv2 Prefill Context Parallelism (PCP) for non-MLA (GQA/MHA) models on the FlashAttention backend, including a sharded PCP+DCP path that supports extend/prefix-caching and mixed prefill+decode batches without desynchronizing collectives.
Changes:
- Lift MRv2 PCP restriction to MLA-only by moving capability gating to attention-backend opt-in (
supports_pcp) + existing compatibility checks. - Add PCPManager→ForwardContext plumbing to provide rank-invariant global batch flags and prefill gather/context metadata needed by FlashAttention’s MRv2 PCP+DCP path.
- Implement FlashAttention MRv2 PCP+DCP execution: prefill Q/K/V all-gather + optional sharded cached-prefix attention + mixed-batch decode overwrite; plus KV-cache update gathering for PCP.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| vllm/v1/worker/gpu/pcp_manager.py | Adds global batch flags and forward-context metadata (prefill gather indices + global context metadata) for sharded PCP+DCP attention. |
| vllm/v1/worker/gpu/model_runner.py | Populates PCP forward-context metadata at the start of each forward pass. |
| vllm/v1/core/sched/scheduler.py | Propagates actual pcp_world_size into scheduler/kv-connector setup (was hardcoded to 1). |
| vllm/v1/attention/backends/flash_attn.py | Implements PCP support for GQA/MHA and an MRv2 PCP+DCP sharded-cache path, including mixed-batch handling and PCP-aware KV-cache updates. |
| vllm/model_executor/layers/attention/pcp.py | Adds helpers for PCP prefill QKV gather/slice and for extracting decode subsets from mixed batches. |
| vllm/config/vllm.py | Removes MRv2 PCP “unsupported feature” block for non-MLA; relies on backend opt-in + compatibility checks. |
| vllm/config/model.py | Adjusts DCP validation logic when DCP spans the PCP axis (dcp == pcp) so TP-head sharding constraints aren’t misapplied. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+1469
to
+1472
| 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) |
Comment on lines
+701
to
+708
| 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 |
Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
…cp-gqa Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: JaredforReal <w13431838023@gmail.com>
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
4 tasks
MatthewBonanni
left a comment
MatthewBonanni
left a comment
Member
There was a problem hiding this comment.
Thanks! A few comments
| layer._v_scale, | ||
| ) | ||
|
|
||
| def _forward_with_dcp( |
Member
There was a problem hiding this comment.
Make sure to update FlashAttentionDiffKV too please
Comment on lines
-1213
to
-1215
| if split_dcp_context: | ||
| # TODO: Remove this DCP + FA2 mixed decode/prefill workaround once | ||
| # FA4 supports this Qwen3.5 shape. |
Member
There was a problem hiding this comment.
Are you sure we can remove this?
| self._forward_with_dcp( | ||
| query[:num_actual_tokens], | ||
| self._forward_dcp( | ||
| query, |
Member
There was a problem hiding this comment.
Why get rid of this slice? Are we computing over padding rows now?
| 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( |
Member
There was a problem hiding this comment.
Please use the workspace and pass out so we aren't allocating each step
9 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Hy3-290B-FP8 PCP/DCP Benchmark in 8*H100
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 underTP1 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.Config matrix — all 7 use exactly 8 GPUs
TP × PCP = 8throughout;DCP == PCPso DCP ranks share PCP ranks (no extraworld size). EP spans TP × PCP = 8 in every config.
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.
TTFT — prefill latency (median ms)
Random-input probe, prefix caching off + 6 warmup requests discarded,
20 measured.
Accuracy & KV capacity
GSM8K 200 questions, 5-shot, greedy, conc 256.
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.
Qwen3-30B-A3B-FP8 PCP/DCP Benchmark in 4xH100
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.Config matrix — all 5 use 4 GPUs
When
dcp == pcpthe DCP ranks share the PCP ranks, so the CP world size isjust
pcp. Config 4 (PCP2+DCP2) therefore runs on 2 GPUs, not 4.TPOT — decode latency (mean / P99 ms)
ShareGPT subset: concurrency 1 / 16 / 32.
The PCP gain is from decode steps no longer doing a per-layer
forward_contextlookup + 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.
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.
Accuracy & KV capacity
GSM8K 1319 questions, 5-shot, greedy, conc 256.