[PCP][Spec Decode] Prefer sharded MTP with replicated draft fallback by pisceskkk · Pull Request #53427 · vllm-project/vllm

pisceskkk · 2026-08-23T04:17:03Z

Summary

Extend the PCP speculative-decoding support already merged in #56107 with backend-aware draft execution:

Prefer the upstream PCP-sharded draft prefill for single-module MTP when the selected draft attention backend supports PCP.
Fall back to a global replicated drafter when the selected backend does not support PCP, while keeping the target model PCP-sharded.
Keep DSpark on its existing replicated path.

This allows PCP-capable and non-PCP-capable draft backends to coexist without disabling the upstream sharded implementation.

Implementation

Draft mode is resolved before model construction from the effective draft attention backend, including attention_backend from the speculative config and backend_per_kind precedence:

Automatic backend selection remains shard-first because the selector receives use_pcp=True and considers PCP-capable candidates.
An explicitly selected backend with supports_pcp() == false builds the drafter with prefill_context_parallel_size=1.

The two runtime contracts remain separate:

Draft mode	Hidden states	Prefill metadata	Input/capture buffers
PCP sharded	Rank-local target states	Reuse rank-local target metadata	PCP local buffers
Replicated fallback	PCP-restored global states	Rebuild from global positions, sequence lengths, and block tables	Drafter-owned global buffers

For PIECEWISE graphs, replicated drafting captures with drafter-owned attention groups and buffers. Runtime graph-padding rows receive zero hidden states and the global is_padding mask.

Warmup follows the same selection: sharded MTP stays rank-local, while replicated drafters restore the global batch before propose().

Validation

Current commit:

Ruff check passed on all changed files.
Ruff format check passed on all changed files.
Python bytecode compilation passed on all changed files.
Unit coverage added for automatic shard preference, explicit backend capability selection, per-KV-kind precedence, config isolation, and global replicated metadata construction.

The available 5090 development container did not have a main-compatible Python/test environment, so model execution for this rebased commit remains to be rerun on the H20 environment.

Replicated-mode model evidence

The following full GSM8K results were collected on the earlier global-replicated implementation using 4x NVIDIA H20, /home/weight/GLM-4.7-Flash, FLASH_ATTN_MLA, EP4, MTP3, all 1319 questions, 5-shot, temperature 0, seed 42, max output tokens 4096, and concurrency 32. They are retained as reference evidence for the replicated execution contract; they are not reported as a rerun of this rebased commit.

Configuration	Correct	Accuracy	Invalid	Acceptance rate	Acceptance length
TP4 + greedy fixed K	960/1319	72.7824%	0	0.409285	2.227856
PCP4 + greedy fixed K	933/1319	70.7354%	0	0.410057	2.230170
PCP4 + greedy fixed K + PIECEWISE	949/1319	71.9484%	1	0.409514	2.228542
PCP4 + probabilistic fixed K	932/1319	70.6596%	0	0.407958	2.223873
PCP4 + probabilistic dynamic K	967/1319	73.3131%	1	0.920441	1.929728

Dynamic K used K=3 for batch sizes 1-8, K=2 for 9-16, and K=1 for 17-32. Its raw acceptance rate is not directly comparable with fixed K.

The fixed-K PCP and TP acceptance metrics were aligned (acceptance-rate delta 0.000771; acceptance-length delta 0.002314). PIECEWISE was also aligned with the eager PCP reference (acceptance-rate delta 0.000543; acceptance-length delta 0.001628).

Scope

Single-module MTP uses capability-selected sharded or replicated draft prefill.
DSpark remains replicated as implemented by [PCP][Spec Decode] Adds PCP support for single-module MTP and replicated DSpark. #56107.
PCP continues to support PIECEWISE CUDA graphs; full graph modes remain outside the current PCP contract.
Existing PCP+DCP restrictions are unchanged.

AI assistance was used for implementation and validation. The human submitter reviewed the changes.

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.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f31a8af8fd

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-08-23T04:22:09Z

+    attn_metadata = speculator._build_draft_attn_metadata(
+        num_reqs=num_reqs,
+        num_reqs_padded=num_reqs_padded,
+        num_tokens_padded=num_tokens_padded,
+        seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,


Refresh the draft max sequence length before metadata

For the standard AutoRegressiveSpeculator MTP path, this helper builds metadata before propose() updates draft_max_seq_len from the current batch, so _build_draft_attn_metadata() uses the preceding batch's value. If the longest request accepts all K drafts plus its bonus token, the next batch's sequence upper bound grows by K+1 while the saved bound only included K lookahead tokens, making max_seqlen_k one token smaller than seq_lens; FlashAttention MLA may then reject the launch or process invalid draft attention. Compute the current bound before this call or pass it directly rather than relying on mutable state from the previous proposal.

Useful? React with 👍 / 👎.

Enable MTP speculative decoding in the MRV2 GPU runner when prefill context parallelism is enabled. The target model stays PCP-sharded while the MTP drafter runs replicated over the global batch on every rank: - Drop PCP from the drafter's config copy so its attention groups, forward context, and cudagraphs are built without PCP. Since the drafter builds its own draft prefill attention metadata from the (global) input batch, no metadata plumbing is needed. - partition_batch supports multi-token decode rows: reuse the GPU request-state positions verbatim (the CPU num_computed_tokens is only an async upper bound after rejection), derive local seq_lens from the segment end positions, and take last-token logits indices directly instead of recombining sampled/draft tokens. - restore_hidden_states appends explicit zero rows when the global batch is graph-padded, and restore_hidden_state_buffer restores persistent max-token buffers (e.g. DeepSeek V4's pre-hc_head residual). Unsupported combinations (non-MTP methods, sparse MLA, DCP, adaptive verification) are rejected at config validation. Adapted from #53427. Co-authored-by: QiuChunshuo <qiuchunshuo@huawei.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

LucasWilkinson · 2026-08-26T10:58:47Z

Im thinking we should maybe do something more like pisceskkk#6

… glm-53-blog/prefiller

mergify · 2026-08-26T12:03:04Z

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

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

…m-project#53427 merge 732f952 enabled PIECEWISE graphs for sparse MLA under PCP (with vllm-project#53515's persistent input buffers); the vllm#53427 merge brought the old guard back.

pisceskkk · 2026-08-26T14:25:54Z

Im thinking we should maybe do something more like pisceskkk#6

have modified the pr with pisceskkk#6 and resolved the merge conflicts. ptal when you have time, thanks!

Enable MTP speculative decoding in the MRV2 GPU runner when prefill context parallelism is enabled. The target model stays PCP-sharded while the MTP drafter runs replicated over the global batch on every rank: - Drop PCP from the drafter's config copy so its attention groups, forward context, and cudagraphs are built without PCP. Since the drafter builds its own draft prefill attention metadata from the (global) input batch, no metadata plumbing is needed. - partition_batch supports multi-token decode rows: reuse the GPU request-state positions verbatim (the CPU num_computed_tokens is only an async upper bound after rejection), derive local seq_lens from the segment end positions, and take last-token logits indices directly instead of recombining sampled/draft tokens. - restore_hidden_states appends explicit zero rows when the global batch is graph-padded, and restore_hidden_state_buffer restores persistent max-token buffers (e.g. DeepSeek V4's pre-hc_head residual). Unsupported combinations (non-MTP methods, sparse MLA, DCP, adaptive verification) are rejected at config validation. Adapted from vllm-project#53427. Co-authored-by: QiuChunshuo <qiuchunshuo@huawei.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

…coding with PCP Enable MTP speculative decoding in the MRV2 GPU runner when prefill context parallelism is enabled. The target model stays PCP-sharded while the MTP drafter runs replicated over the global batch on every rank: - Drop PCP from the drafter's config copy so its attention groups, forward context, and cudagraphs are built without PCP. Since the drafter builds its own draft prefill attention metadata from the (global) input batch, no metadata plumbing is needed. - partition_batch supports multi-token decode rows: reuse the GPU request-state positions verbatim (the CPU num_computed_tokens is only an async upper bound after rejection), derive local seq_lens from the segment end positions, and take last-token logits indices directly instead of recombining sampled/draft tokens. - restore_hidden_states appends explicit zero rows when the global batch is graph-padded, and restore_hidden_state_buffer restores persistent max-token buffers (e.g. DeepSeek V4's pre-hc_head residual). Unsupported combinations (non-MTP methods, sparse MLA, DCP, adaptive verification) are rejected at config validation. Adapted from vllm-project#53427. Co-authored-by: QiuChunshuo <qiuchunshuo@huawei.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

mergify · 2026-09-10T20:34:02Z

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

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

mergify · 2026-09-11T09:07:49Z

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

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

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

mergify · 2026-09-21T20:02:34Z

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

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

pisceskkk requested review from WoosukKwon, njhill and yewentao256 as code owners August 23, 2026 04:17

claude Bot reviewed Aug 23, 2026

View reviewed changes

mergify Bot added speculative-decoding mrv2 Model Runner V2 specific labels Aug 23, 2026

chatgpt-codex-connector Bot reviewed Aug 23, 2026

View reviewed changes

pisceskkk changed the title ~~[Spec Decode] Support MTP with PCP in the V2 GPU runner~~ feat(PCP)(Spec Decode): Support MTP with PCP in the V2 GPU runner Aug 24, 2026

pisceskkk force-pushed the pcp/mtp branch from f31a8af to 22801af Compare August 24, 2026 07:46

This was referenced Aug 24, 2026

[Spec Decode] Support MTP with PCP GirasoleY/vllm#4

Draft

[Spec Decode] Support PCP-sharded MTP prefill #53653

Closed

[RFC]: Support Prefill Context Parallel (PCP) #25749

Open

This was referenced Aug 25, 2026

[Spec Decode] Let the speculator build its own draft prefill attention metadata #53660

Draft

[PCP][Spec Decode] Consolidate MTP support and add DSpark groundwork #53661

Closed

mergify Bot added the dflash label Aug 25, 2026

github-project-automation Bot added this to Sprint - DFlash Aug 25, 2026

github-project-automation Bot moved this to Backlog in Sprint - DFlash Aug 25, 2026

LucasWilkinson mentioned this pull request Aug 25, 2026

[RFC]: Decouple draft-model parallelism from the target model (PCP/DCP + speculative decoding) #53673

Open

1 task done

LucasWilkinson mentioned this pull request Aug 26, 2026

[PCP][Spec Decode] Simplify PCP+MTP draft metadata handling; add DSpark support pisceskkk/vllm#6

Merged

LucasWilkinson added a commit to neuralmagic/vllm that referenced this pull request Aug 26, 2026

Merge pisceskkk#6 (PCP+MTP simplification on vllm-project#53427) into…

b42af97

… glm-53-blog/prefiller

mergify Bot added the needs-rebase label Aug 26, 2026

pisceskkk force-pushed the pcp/mtp branch from 09ec600 to 5a509a4 Compare August 26, 2026 13:53

pisceskkk changed the title ~~feat(PCP)(Spec Decode): Support MTP with PCP in the V2 GPU runner~~ feat(PCP)(Spec Decode): Support MTP and DSpark with PCP Aug 26, 2026

mergify Bot removed the needs-rebase label Aug 26, 2026

LopezCastroRoberto mentioned this pull request Aug 27, 2026

[Perf][PCP][DSpark] Shard context-KV precompute across PCP ranks #54036

Draft

LucasWilkinson mentioned this pull request Sep 9, 2026

[PCP][Spec Decode] Adds PCP support for single-module MTP and replicated DSpark. #56107

Merged

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

pisceskkk force-pushed the pcp/mtp branch from 5a509a4 to 73ec79b Compare September 11, 2026 01:39

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

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

This was referenced Sep 14, 2026

[PCP][DCP] Support DCP target model with non-DCP Dspark #56723

Merged

[Bugfix][Spec Decode][PCP] Scope CP impl checks to the target model's layers #56737

Closed

pisceskkk force-pushed the pcp/mtp branch from 73ec79b to e47a448 Compare September 15, 2026 03:56

pisceskkk changed the title ~~feat(PCP)(Spec Decode): Support MTP and DSpark with PCP~~ [PCP][Spec Decode] Prefer sharded MTP with replicated draft fallback Sep 15, 2026

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

[PCP][Spec Decode] Add replicated draft fallback

9961a86

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

pisceskkk force-pushed the pcp/mtp branch from e47a448 to 9961a86 Compare September 15, 2026 04:05

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

pisceskkk closed this Sep 22, 2026

github-project-automation Bot moved this from Backlog to Done in Sprint - DFlash Sep 22, 2026

Navigation Menu

Repository navigation

[PCP][Spec Decode] Prefer sharded MTP with replicated draft fallback - #53427

[PCP][Spec Decode] Prefer sharded MTP with replicated draft fallback#53427
pisceskkk wants to merge 1 commit into
vllm-project:mainfrom
pisceskkk:pcp/mtp

pisceskkk commented Aug 23, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Aug 23, 2026

Uh oh!

LucasWilkinson commented Aug 26, 2026

Uh oh!

mergify Bot commented Aug 26, 2026

Uh oh!

pisceskkk commented Aug 26, 2026

Uh oh!

mergify Bot commented Sep 10, 2026

Uh oh!

mergify Bot commented Sep 11, 2026

Uh oh!

mergify Bot commented Sep 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

pisceskkk commented Aug 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Implementation

Validation

Replicated-mode model evidence

Scope

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

LucasWilkinson commented Aug 26, 2026

Uh oh!

mergify Bot commented Aug 26, 2026

Uh oh!

pisceskkk commented Aug 26, 2026

Uh oh!

mergify Bot commented Sep 10, 2026

Uh oh!

mergify Bot commented Sep 11, 2026

Uh oh!

mergify Bot commented Sep 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pisceskkk commented Aug 23, 2026 •

edited

Loading