Skip to content

[WIP] pcp alternative impl#33403

Draft
LucasWilkinson wants to merge 1 commit intomainfrom
pcp-alt
Draft

[WIP] pcp alternative impl#33403
LucasWilkinson wants to merge 1 commit intomainfrom
pcp-alt

Conversation

@LucasWilkinson
Copy link
Copy Markdown
Collaborator

@LucasWilkinson LucasWilkinson commented Jan 30, 2026 •
edited by github-actions bot
Loading

Alternative implementation to: #28988

Makes DCP/PCP interface conform to to: #25749 (comment)

Rank layout:

================================================================================
LEGEND
================================================================================
  q0-3   = query heads 0 through 3
  k0     = KV head 0
  s0     = KV shard 0 (decode sequence shard)
  p0     = prefill sequence shard 0
  [0, 2] = ranks in group

================================================================================
TP=2, PCP=2, DCP=2  (8 query heads, 2 KV heads)
================================================================================

Rank Layout (query heads / KV head / KV shard / prefill shard):
      tp0                    tp1
     +-------------------+-------------------+
pcp0 | 0: q0-3/k0/s0/p0  | 1: q4-7/k1/s0/p0  |
     +-------------------+-------------------+
pcp1 | 2: q0-3/k0/s1/p1  | 3: q4-7/k1/s1/p1  |
     +-------------------+-------------------+

PCP Groups (shard sequence during prefill):
  PCP0: [0, 2]  0: q0-3/k0, s0, p0 | 2: q0-3/k0, s1, p1
  PCP1: [1, 3]  1: q4-7/k1, s0, p0 | 3: q4-7/k1, s1, p1

TP Groups (shard heads):
  TP0:  [0, 1]  0: q0-3/k0, s0, p0 | 1: q4-7/k1, s0, p0
  TP1:  [2, 3]  2: q0-3/k0, s1, p1 | 3: q4-7/k1, s1, p1

DCP Groups (shard KV during decode):
  DCP0: [0, 2]  0: q0-3/k0, s0, p0 | 2: q0-3/k0, s1, p1
  DCP1: [1, 3]  1: q4-7/k1, s0, p0 | 3: q4-7/k1, s1, p1


================================================================================
TP=2, PCP=2, DCP=2  (8 query heads, 1 KV head - MQA)
================================================================================

Rank Layout (query heads / KV head / KV shard / prefill shard):
      tp0                    tp1
     +-------------------+-------------------+
pcp0 | 0: q0-3/k0/s0/p0  | 1: q4-7/k0/s0/p0  |
     +-------------------+-------------------+
pcp1 | 2: q0-3/k0/s1/p1  | 3: q4-7/k0/s1/p1  |
     +-------------------+-------------------+

NOTE: (k0, s0) and (k0, s1) are each replicated once (since num_gpus = TP * PCP and num_gpus / DCP > num_kv_heads)

PCP Groups (shard sequence during prefill):
  PCP0: [0, 2]  0: q0-3/k0, s0, p0 | 2: q0-3/k0, s1, p1
  PCP1: [1, 3]  1: q4-7/k0, s0, p0 | 3: q4-7/k0, s1, p1

TP Groups (shard heads):
  TP0:  [0, 1]  0: q0-3/k0, s0, p0 | 1: q4-7/k0, s0, p0
  TP1:  [2, 3]  2: q0-3/k0, s1, p1 | 3: q4-7/k0, s1, p1

DCP Groups (shard KV during decode):
  DCP0: [0, 2]  0: q0-3/k0, s0, p0 | 2: q0-3/k0, s1, p1
  DCP1: [1, 3]  1: q4-7/k0, s0, p0 | 3: q4-7/k0, s1, p1


================================================================================
TP=2, PCP=2, DCP=4  (8 query heads, 2 KV heads)
================================================================================

Rank Layout (query heads / KV head / KV shard / prefill shard):
      tp0                    tp1
     +-------------------+-------------------+
pcp0 | 0: q0-3/k0/s0/p0  | 1: q4-7/k1/s2/p0  |
     +-------------------+-------------------+
pcp1 | 2: q0-3/k0/s1/p1  | 3: q4-7/k1/s3/p1  |
     +-------------------+-------------------+

PCP Groups (shard sequence during prefill):
  PCP0: [0, 2]  0: q0-3/k0, s0, p0 | 2: q0-3/k0, s1, p1
  PCP1: [1, 3]  1: q4-7/k1, s2, p0 | 3: q4-7/k1, s3, p1

TP Groups (shard heads):
  TP0:  [0, 1]  0: q0-3/k0, s0, p0 | 1: q4-7/k1, s2, p0
  TP1:  [2, 3]  2: q0-3/k0, s1, p1 | 3: q4-7/k1, s3, p1

DCP Groups (shard KV during decode):
  DCP0: [0, 2, 1, 3]  0: q0-3/k0, s0, p0 | 2: q0-3/k0, s1, p1 | 1: q4-7/k1, s2, p0 | 3: q4-7/k1, s3, p1


================================================================================
TP=2, PCP=2, DCP=4  (8 query heads, 1 KV head - MQA)
================================================================================

Rank Layout (query heads / KV head / KV shard / prefill shard):
      tp0                    tp1
     +-------------------+-------------------+
pcp0 | 0: q0-3/k0/s0/p0  | 1: q4-7/k0/s2/p0  |
     +-------------------+-------------------+
pcp1 | 2: q0-3/k0/s1/p1  | 3: q4-7/k0/s3/p1  |
     +-------------------+-------------------+

PCP Groups (shard sequence during prefill):
  PCP0: [0, 2]  0: q0-3/k0, s0, p0 | 2: q0-3/k0, s1, p1
  PCP1: [1, 3]  1: q4-7/k0, s2, p0 | 3: q4-7/k0, s3, p1

TP Groups (shard heads):
  TP0:  [0, 1]  0: q0-3/k0, s0, p0 | 1: q4-7/k0, s2, p0
  TP1:  [2, 3]  2: q0-3/k0, s1, p1 | 3: q4-7/k0, s3, p1

DCP Groups (shard KV during decode):
  DCP0: [0, 2, 1, 3]  0: q0-3/k0, s0, p0 | 2: q0-3/k0, s1, p1 | 1: q4-7/k0, s2, p0 | 3: q4-7/k0, s3, p1

@github-project-automation github-project-automation bot moved this to Todo in AMD Jan 30, 2026
gemini-code-assist[bot]
gemini-code-assist bot reviewed Jan 30, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces Prefill Context Parallelism (PCP), a significant feature for distributing the computational load of prefill requests across multiple GPUs. The implementation is extensive, touching on attention mechanisms, worker logic, and configuration. Key changes include the introduction of a PCPManager to handle input partitioning and output restoration, a fused Triton kernel for optimized QKV selection, and a generalization of context parallelism code from dcp to a more generic cp. The changes are well-structured, with necessary updates to tests and compatibility checks, such as disabling full CUDA graphs when PCP is active. The PR also includes improvements for asynchronous pipeline parallelism and multimodal input handling. Overall, this is a solid and well-thought-out implementation of a complex feature.

@mergify
Copy link
Copy Markdown

mergify bot commented Jan 30, 2026

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LucasWilkinson.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify
Copy link
Copy Markdown

mergify bot commented Jan 31, 2026

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LucasWilkinson.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify
Copy link
Copy Markdown

mergify bot commented Feb 5, 2026

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LucasWilkinson.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Feb 5, 2026
@mergify
Copy link
Copy Markdown

mergify bot commented Mar 2, 2026

Documentation preview: https://vllm--33403.org.readthedocs.build/en/33403/

@mergify mergify bot added documentation Improvements or additions to documentation and removed needs-rebase labels Mar 2, 2026
@LucasWilkinson LucasWilkinson force-pushed the pcp-alt branch 3 times, most recently from dc1e212 to b3c3f97 Compare March 2, 2026 05:31
@LucasWilkinson
[Core] Add Prefill Context Parallelism (PCP) and All-to-All DCP commu…
a97bd60 
…nication

This PR adds Prefill Context Parallelism (PCP) support for splitting prefill
tokens across ranks using a DualChunkSwap pattern, and integrates an All-to-All
communication backend for Decode Context Parallelism (DCP).

Key changes:
- Add PCP with DualChunkSwap token partitioning for balanced prefill computation
- Add All-to-All DCP communication backend reducing NCCL calls from 3 to 2
- Restrict DCP+PCP to two clean configurations:
  - Case 1: DCP = PCP (same TP position, all-reduce only)
  - Case 2: DCP = TP × PCP (full TP all-gather, all-reduce + slice)
- Add PCPManager for buffer management and input partitioning
- Update attention backends (FlashAttention, FlashInfer, MLA) for PCP support
- Add comprehensive tests for DCP operations

Co-Authored-By: QiuChunshuo <qiuchunshuo@huawei.com>
Co-Authored-By: zhenwenqi2024 <zhenwenqi_2022@qq.com>
Co-Authored-By: FENP <yuanyongjie.yyj@antgroup.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
@mergify
Copy link
Copy Markdown

mergify bot commented Mar 2, 2026

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LucasWilkinson.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Mar 2, 2026
@LucasWilkinson LucasWilkinson mentioned this pull request Mar 7, 2026
Draft
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mgoin mgoin Awaiting requested review from mgoin mgoin will be requested when the pull request is marked ready for review mgoin is a code owner

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat will be requested when the pull request is marked ready for review robertgshaw2-redhat is a code owner

@yewentao256 yewentao256 Awaiting requested review from yewentao256 yewentao256 will be requested when the pull request is marked ready for review yewentao256 is a code owner

@pavanimajety pavanimajety Awaiting requested review from pavanimajety pavanimajety will be requested when the pull request is marked ready for review pavanimajety is a code owner

@njhill njhill Awaiting requested review from njhill njhill will be requested when the pull request is marked ready for review njhill is a code owner

@heheda12345 heheda12345 Awaiting requested review from heheda12345 heheda12345 will be requested when the pull request is marked ready for review heheda12345 is a code owner

@benchislett benchislett Awaiting requested review from benchislett benchislett will be requested when the pull request is marked ready for review benchislett is a code owner

@luccafong luccafong Awaiting requested review from luccafong luccafong will be requested when the pull request is marked ready for review luccafong is a code owner

@MatthewBonanni MatthewBonanni Awaiting requested review from MatthewBonanni MatthewBonanni will be requested when the pull request is marked ready for review MatthewBonanni is a code owner

@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon will be requested when the pull request is marked ready for review WoosukKwon is a code owner

@ywang96 ywang96 Awaiting requested review from ywang96 ywang96 will be requested when the pull request is marked ready for review ywang96 is a code owner

@alexm-redhat alexm-redhat Awaiting requested review from alexm-redhat alexm-redhat will be requested when the pull request is marked ready for review alexm-redhat is a code owner

@ApostaC ApostaC Awaiting requested review from ApostaC ApostaC will be requested when the pull request is marked ready for review ApostaC is a code owner

@orozery orozery Awaiting requested review from orozery orozery will be requested when the pull request is marked ready for review orozery is a code owner

@vadiklyutiy vadiklyutiy Awaiting requested review from vadiklyutiy vadiklyutiy will be requested when the pull request is marked ready for review vadiklyutiy is a code owner

@zhuohan123 zhuohan123 Awaiting requested review from zhuohan123 zhuohan123 will be requested when the pull request is marked ready for review zhuohan123 is a code owner

@youkaichao youkaichao Awaiting requested review from youkaichao youkaichao will be requested when the pull request is marked ready for review youkaichao is a code owner

@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth will be requested when the pull request is marked ready for review tlrmchlsmth is a code owner

@houseroad houseroad Awaiting requested review from houseroad houseroad will be requested when the pull request is marked ready for review houseroad is a code owner

@hmellor hmellor Awaiting requested review from hmellor hmellor will be requested when the pull request is marked ready for review hmellor is a code owner

@ProExpertProg ProExpertProg Awaiting requested review from ProExpertProg ProExpertProg will be requested when the pull request is marked ready for review ProExpertProg is a code owner

1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

documentation Improvements or additions to documentation needs-rebase nvidia rocm Related to AMD ROCm speculative-decoding v1

Projects

AMD
Status: Todo
NVIDIA
Status: No status

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@LucasWilkinson