[PCP][DCP] Enable PCP+DCP on sparse-MLA models by PatrykSaffer · Pull Request #56157 · vllm-project/vllm

PatrykSaffer · 2026-09-09T21:28:21Z

Purpose

Enable DCP to run alongside PCP on sparse-MLA models, so a deployment can shard the KV cache for decode and still split prefill query work across ranks.

This change:

gives the sparse indexer a DCP-gathered prefill path: this rank's KV shard is all-gathered across the DCP group
gives the FlashMLA-sparse backend the matching prefill path, with a rank-major gathered KV workspace and a PCP-invariant chunk plan every rank derives identically without communicating;
replicates prefills too short to fill all 2*PCP chunks when DCP is on, which keeps the per-region row count identical on every rank — the invariant the gathered layout rests on;
adds a MoE all-reduce fast path for steps whose rows are all PCP-replicated, replacing an unconditional all-gather + reduce-scatter that duplicated every decode row PCP times;
relaxes the PCP CUDA graph restriction for PCP+DCP to everything up to FULL_AND_PIECEWISE, keeping only FULL over mixed batches rejected;

fixes #53573
Additionally

fixes a pre-existing kernel/reference divergence: the top-k index kernel applied the block-table bound to prefill-workspace rows, while _triton_convert_reference_impl has always specified that only -1 is invalid there. The existing test missed it because token_indices[6, 150:160] is an empty slice on a 128-column tensor;

Test Plan

NVIDIA GB200 x4
GLM-5.3, NVFP4 weights
TP=1, PCP=4, DCP=4

Test Result

GSM8K, 1319 questions, TP1+PCP4+DCP4, FULL_AND_PIECEWISE. Seven runs across the

Run	Accuracy	Invalid
TP4	0.830	0.000
PCP4DCP4	0.832	0.000

Benchmarks

Prefill, 32768 in / 1 out:

Layout	TTFT ms	P99 TTFT ms	vs tp4
tp4	5512.1	8141.1	1.00x
pcp4	3278.1	4814.0	1.68x
pcp4dcp4	3326.8	4888.7	1.66x

Decode, 1 in / 1024 out:

Layout	TPOT ms	out tok/s	vs tp4
tp4	18.6	3412.7	1.00x
dcp4	28.4	2236.8	0.66x
pcp4dcp4, PIECEWISE (this PR)	86.7	735.9	0.22x
pcp4dcp4, with FULL decode graphs and moe deduplication PatrykSaffer/vllm@`patryk/pcp-dcp-sparse-mla...patryk/pcp-dcp-full-cg-no-duplication`	32.6	1944.1	0.57x

Why pcp4dcp4 decode is slower than dcp4 decode:
CudaCommunicator.init enables custom all-reduce, torch symm-mem and FlashInfer fused all-reduce for tp groups only, so that per-layer all-reduce runs on NCCL ring.
followup is to enable it, sth like 590e9b7

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-1.blc1.mistralcompute.ai>

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-0.blc1.mistralcompute.ai>

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

mergify · 2026-09-09T21:28:59Z

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

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

Signed-off-by: PatrykSaffer <patryk.saffer@mistral.ai>

Add comment to clarify group name matching for ETP. Signed-off-by: PatrykSaffer <patryk.saffer@mistral.ai>

…se-mla

…patryk/pcp-dcp-sparse-mla

Drop the FlashMLA-sparse PCP+DCP unit test file and the three indexer chunk-plan tests; the kernel-level deinterleave test stays. Add a GLM-5.2-NVFP4 TP1+PCP4+DCP4 config to the existing 4xB200 LM Eval PCP job so PCP+DCP is covered end to end. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

LucasWilkinson · 2026-09-12T18:36:45Z

/ci run

github-actions · 2026-09-12T18:37:03Z

✅ Triggered Buildkite CI #88553 for commit 81050300cb21.

…t param - Give the fake parallel_config in the index-conversion warmup test the prefill_context_parallel_size the warmup keys now read. - Give the fake FlashMLA impl in the decode-subset test pcp_dcp_kv_gather. - Drop the [5, 0, 6] case from the PCP deinterleave test: a request with an empty context is never scheduled, and the plan builder asserts on it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

LucasWilkinson · 2026-09-12T19:54:16Z

/ci run

github-actions · 2026-09-12T19:54:34Z

✅ Triggered Buildkite CI #88558 for commit 6830dd62a90a.

…review-fixes

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

LucasWilkinson · 2026-09-12T21:08:13Z

/ci run

github-actions · 2026-09-12T21:09:16Z

✅ Triggered Buildkite CI #88566 for commit 6e93ff1f51b8.

Rebase vllm-project#54496 onto vllm-project#56157. Publish each TP1 PCP+DCP producer shard through the existing DCP transfer topology and retain upstream completion aggregation. Share transfer-size calculation between the pull and push schedulers and validate the PCP/DCP group in the connector. Cover shard publication, completion, push scheduling, and advertised transfer topology for replicated and sharded producers. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Required to run MTP in PCP+DCP behind NIXL P/D at all. NIXL with DCP > 1 raises cp_kv_cache_interleave_size to the block size (64), and cp_utils.validate_cp_support then asserts supports_mtp_with_cp_non_trivial_interleave_size for every layer impl when a speculative config is present: AssertionError: MTP with cp_kv_cache_interleave_size > 1 is not supported in FlashMLASparseImpl. The flag is declared False in vllm/v1/attention/backend.py and set True by no backend upstream, so MTP + DCP > 1 + NIXL is unreachable on plain main. Neither the flag nor the assertion comes from vllm-project#56157; both predate it. No open upstream PR touches either - searched by flag name and by "MTP cp_kv_cache_interleave". The frankenstein branch sets it True on the same impl in ebe6cc2, on the grounds that the mixed fp8 path processes all scheduled query tokens together and its slot mapping and sparse-index conversion both honour the interleave. That reasoning is plausible here too - the conversion calls thread cp_kv_cache_interleave_size through - but vllm-project#56157 rewrote this prefill path, so the claim is NOT verified for it. Treat this as a hypothesis under test. Startup succeeding only proves the assertion is gone; the arm must be judged on MTP acceptance rate and output sanity before this is proposed upstream. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit 95474a0)

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-1.blc1.mistralcompute.ai> Signed-off-by: patryk.saffer <patryk.saffer@blc-login-0.blc1.mistralcompute.ai> Signed-off-by: PatrykSaffer <patryk.saffer@mistral.ai> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: patryk.saffer <patryk.saffer@blc-login-1.blc1.mistralcompute.ai> Co-authored-by: patryk.saffer <patryk.saffer@blc-login-0.blc1.mistralcompute.ai> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Patryk Saffer <PatrykSaffer@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

patryk.saffer and others added 7 commits September 8, 2026 22:33

PCP+DCP

89ae7f6

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-1.blc1.mistralcompute.ai>

Nits cleanups

1b5bd4c

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-1.blc1.mistralcompute.ai>

better sparse checks

f87826c

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-0.blc1.mistralcompute.ai>

nits

9ec2eea

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-0.blc1.mistralcompute.ai>

nits

06bdf53

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-0.blc1.mistralcompute.ai>

nits

d354ced

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-0.blc1.mistralcompute.ai>

nits

7172d65

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-0.blc1.mistralcompute.ai>

PatrykSaffer requested review from LucasWilkinson, MatthewBonanni, WoosukKwon, mgoin, njhill, pavanimajety, yewentao256 and zyongye as code owners September 9, 2026 21:28

claude Bot reviewed Sep 9, 2026

View reviewed changes

PatrykSaffer changed the title ~~Patryk/pcp dcp sparse mla~~ [PCP][DCP] Enable PCP+DCP on sparse-MLA models Sep 9, 2026

mergify Bot added deepseek Related to DeepSeek models nvidia mrv2 Model Runner V2 specific labels Sep 9, 2026

github-project-automation Bot added this to NVIDIA Sep 9, 2026

mergify Bot added the needs-rebase label Sep 9, 2026

Merge branch 'main' into patryk/pcp-dcp-sparse-mla

8a659c5

Signed-off-by: PatrykSaffer <patryk.saffer@mistral.ai>

mergify Bot removed the needs-rebase label Sep 9, 2026

PatrykSaffer and others added 5 commits September 9, 2026 23:43

Clarify group name matching for ETP

81283e7

Add comment to clarify group name matching for ETP. Signed-off-by: PatrykSaffer <patryk.saffer@mistral.ai>

Merge remote-tracking branch 'upstream/main' into patryk/pcp-dcp-spar…

ce2ae61

…se-mla

Retarget indexer chunk test onto the static method moved upstream

02cbba5

Merge remote-tracking branch 'origin/patryk/pcp-dcp-sparse-mla' into …

b6bf591

…patryk/pcp-dcp-sparse-mla

Move fused all-reduce for CP groups to its own PR

590e9b7

LucasWilkinson requested review from AndreasKaratzas, Harry-Chen, khluu and vadiklyutiy as code owners September 12, 2026 18:35

mergify Bot added the ci/build label Sep 12, 2026

mergify Bot removed the needs-rebase label Sep 12, 2026

LucasWilkinson enabled auto-merge (squash) September 12, 2026 20:12

LucasWilkinson and others added 2 commits September 12, 2026 21:05

Merge remote-tracking branch 'upstream/main' into lwilkinson/pcp-dcp-…

fd403b7

…review-fixes

[PCP][DCP] Fix CI: add pcp_dcp_kv_gather to sparse-MLA test fakes

6e93ff1

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

LucasWilkinson merged commit ebe1dec into vllm-project:main Sep 12, 2026
175 checks passed

github-project-automation Bot moved this from Ready to Done in NVIDIA Sep 12, 2026

LucasWilkinson mentioned this pull request Sep 12, 2026

[NIXL][PCP][DCP] Expose PCP producer KV shards as transfer ranks #56645

Merged

he-yufeng mentioned this pull request Sep 13, 2026

[Bugfix][Attention] Declare layer-compact KV layouts for V3.2 sparse MLA #55438

Open

khluu mentioned this pull request Sep 13, 2026

[CI] Select the supported DCP backend for PCP eval #56677

Merged

Navigation Menu

Uh oh!

[PCP][DCP] Enable PCP+DCP on sparse-MLA models - #56157

[PCP][DCP] Enable PCP+DCP on sparse-MLA models#56157
LucasWilkinson merged 31 commits into
vllm-project:mainfrom
PatrykSaffer:patryk/pcp-dcp-sparse-mla

PatrykSaffer commented Sep 9, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

mergify Bot commented Sep 9, 2026

Uh oh!

LucasWilkinson commented Sep 12, 2026

Uh oh!

github-actions Bot commented Sep 12, 2026

Uh oh!

LucasWilkinson commented Sep 12, 2026

Uh oh!

github-actions Bot commented Sep 12, 2026

Uh oh!

LucasWilkinson commented Sep 12, 2026

Uh oh!

github-actions Bot commented Sep 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

PatrykSaffer commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Benchmarks

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Sep 9, 2026

Uh oh!

LucasWilkinson commented Sep 12, 2026

Uh oh!

github-actions Bot commented Sep 12, 2026

Uh oh!

LucasWilkinson commented Sep 12, 2026

Uh oh!

github-actions Bot commented Sep 12, 2026

Uh oh!

LucasWilkinson commented Sep 12, 2026

Uh oh!

github-actions Bot commented Sep 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

PatrykSaffer commented Sep 9, 2026 •

edited

Loading