[Spec][V2] Support MTP speculative decoding under pipeline parallelism - #46994
[Spec][V2] Support MTP speculative decoding under pipeline parallelism#46994eastwood-c wants to merge 15 commits into
Conversation
On the V2 model runner, MTP speculative decoding does not work under pipeline parallelism. Three things are missing/broken, all on the path that only runs once PP>1: 1. DeepSeekMTP does not implement SupportsPP, so the engine refuses to build it under PP at all (NotImplementedError at model resolution). The MTP draft runs only on the last PP stage, so it never consumes PP intermediate tensors, but SupportsPP still requires the make_empty_intermediate_tensors factory. 2. PPHandler.broadcast() sends sampled_token_ids at its natural width (1 on steps with no draft tokens, num_spec+1 once rejection sampling runs) while receive() always posts a [num_reqs, max_sample_len] buffer. NCCL broadcast does not negotiate element counts, so a width-1 send against a width-max recv is a count mismatch that deadlocks the receiver. Pad the source to max_sample_len (trailing -1, ignored by post_update). 3. The proposed draft tokens are written into req_states.draft_tokens on the last rank only (the propose() path). Non-last ranks keep the zero-init buffer, so combine_sampled_and_draft_tokens embeds zeros at the draft positions on rank 0 -> garbage verification input and near-zero acceptance. Relay the proposed draft tokens to the non-last ranks by coalescing a third broadcast into the existing deferred PPHandler sibling-group broadcast, and scatter it into req_states.draft_tokens on consume. No new collective. Validated on GLM-5.2-FP8 (DeepSeek-Sparse-Attention MoE), TP4/PP2, k=1, on a current-main base: boots, serves coherent greedy output, and draft acceptance is in the normal range (mean acceptance length ~1.3) rather than ~0. A residual draft-acceptance gap specific to DSA models under PP remains and is tracked separately. Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
… MTP+PP Under MTP speculative decoding with pipeline parallelism, the sparse MLA attention backends store a reference to `indexer.topk_indices_buffer` at construction time. When `_maybe_share_lm_head` later replaces `Indexer.topk_indices_buffer` with the target model's buffer, the impl's reference is stale — still pointing to the draft model's original (uninitialized) buffer. This causes garbage DSA attention and degenerate "repeat-the-current-token" drafts (~27-33% acceptance instead of ~85%). Fix: store `self._indexer = indexer` in each sparse MLA backend's `__init__`, and read `self._indexer.topk_indices_buffer` dynamically in `forward_mqa`. Applied to all three sparse MLA backends: `flashattn_mla_sparse.py`, `flashmla_sparse.py`, `flashinfer_mla_sparse.py`. After fix: ~90% acceptance (4261/4753 tokens) at K=1, ~74% at K=3 (3.3 tokens/step). Matches the non-PP TP8 baseline (~85%). Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev>
|
@njhill — thanks for referencing this PR from #47172. Since the original post, we've widened the validation beyond GLM-5.2-FP8 to 8 Qwen3.5/3.6 model variants (dense BF16, AWQ, MoE BF16, MoE AWQ, MoE GPTQ-Int4), each at MTP K=1/2/3, PP=2/TP=1. All show high acceptance (83-96%) that scales gracefully with K, confirming the fix is not architecture-specific.
Full per-position breakdown and GSM8K accuracy available on request. Happy to restructure the PR however you and the codeowners prefer. |
Under MTP+PP, the Qwen3.5 MTP draft model on the last PP rank was using the target model's hidden_states directly, bypassing the fc projection entirely. This produced essentially random predictions (~1% acceptance) because the draft model's input was not properly projected. Fix: On the last PP rank, apply the same fc projection as the first rank (embed input_ids, normalize, concat with hidden_states, project through fc). This is the same pattern used on the first PP rank. After fix: ~86.5% acceptance (6771/7828 tokens) on Qwen3.5-27B-AWQ (TP1/PP2, MTP k=1). Validated on a different model architecture than GLM-5.2-FP8, confirming the fix generalizes. Cross-model validation: the fix was further validated across 8 Qwen3.5/3.6 model variants (dense BF16, AWQ, MoE BF16, MoE AWQ, MoE GPTQ-Int4), each at MTP K=1/2/3, PP=2/TP=1. All show high acceptance (83-96%) that scales gracefully with K, confirming the fix is not architecture-specific. | Model | Quant | K=1 | K=2 | K=3 | |-------------------------|-------------|-------|-------|-------| | Qwen3.5-27B-AWQ | AWQ 4-bit | 95.3% | 90.5% | 85.8% | | Qwen3.5-27B (BF16) | BF16 | 95.5% | 91.1% | 86.0% | | Qwen3.5-35B-A3B (BF16) | BF16 | 93.6% | 88.0% | 82.0% | | Qwen3.6-27B (BF16) | BF16 | 95.2% | 91.1% | 86.0% | | Qwen3.6-27B-AWQ | AWQ 4-bit | 95.3% | 90.5% | 85.8% | | Qwen3.6-35B-A3B (BF16) | BF16 | 94.1% | 88.6% | 83.8% | | Qwen3.6-35B-A3B-AWQ | AWQ 4-bit | 94.0% | 88.8% | 83.0% | | Qwen3.6-35B-A3B-GPTQ-Int4 | GPTQ 4-bit | 94.1% | 88.9% | 83.5% | Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev>
…ends Apply the same stale-buffer fix (commit c175667) to the three remaining sparse MLA backends that were not covered by the original fix commit: - flashinfer_mla_sparse_sm120.py — FlashInferMLASparseSM120Impl (SM120) - rocm_aiter_mla_sparse.py — ROCmAiterMLASparseImpl (ROCm) - xpu_mla_sparse.py — XPUMLASparseImpl (Intel XPU) These backends store indexer.topk_indices_buffer at construction time and read it statically in forward_mqa, which is stale after _maybe_share_lm_head replaces Indexer.topk_indices_buffer with the target model's buffer. The fix is identical to the already-validated fix: store self._indexer = indexer in __init__, read self._indexer.topk_indices_buffer dynamically in forward_mqa. Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev>
Add unit tests for the three core MTP+PP fixes in PR vllm-project#46994: - Fix vllm-project#1: DeepSeekMTP implements SupportsPP interface - Fix vllm-project#2: PPHandler.broadcast() pads sampled_token_ids to max_sample_len - Fix vllm-project#4: Stale topk_indices_buffer is read dynamically via self._indexer Tests are CPU-only (no GPU/distributed required) and follow vLLM's pytest conventions. Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
GLM4MoE's MTP draft (Glm4MoeMTP) did not implement the SupportsPP interface, so the engine refused to build it under PP with the same NotImplementedError as DeepSeekMTP before PR vllm-project#46994 fix #1. The draft runs only on the last PP stage and never consumes PP intermediate tensors, but the make_empty_intermediate_tensors factory is still required. Mirrors the DeepSeek fix: add SupportsPP to the class bases and set self.make_empty_intermediate_tensors via make_empty_intermediate_tensors_factory with keys [hidden_states, residual], matching Glm4MoeModel. GLM4MoE already loads its own embed_tokens in load_weights, so no embed fix needed. Co-authored-by: Agent Name Here <agent@example.com>
Conflicts, and how they were resolved: - flashattn/flashinfer/flashmla sparse MLA impls: upstream moved the field assignments into SparseMLACommonImpl.__init__. Took upstream's super() call and kept `self._indexer = indexer` alongside it. The base class only snapshots indexer.topk_indices_buffer at construction, and the whole point of this branch's sparse-MLA change is that the snapshot goes stale -- forward_mqa has to re-read the buffer off the indexer, because it is swapped between the target and draft passes under MTP+PP. The three sparse backends upstream did not refactor (flashinfer_sm120, rocm_aiter, xpu) still inherit MLAAttentionImpl and set _indexer themselves; those merged cleanly. - deepseek_mtp.py imports: union of both sides, so get_spec_layer_idx_from_weight_name (upstream) and SupportsPP / make_empty_intermediate_tensors_factory (ours) are all imported. Also drops `logger = init_logger(__name__)`, which this branch added and never used; the init_logger import itself was already gone from the merged result. - autoregressive/speculator.py _multi_step_decode: both sides added a parameter. Kept both, with upstream's required seq_lens_cpu_upper_bound ahead of this branch's defaulted intermediate_tensors, and the single call site updated to match. Plus a ruff-format reflow of the make_empty_intermediate_tensors assignment.
|
This pull request has merge conflicts that must be resolved before it can be |
|
I independently reproduced the MRV2 MTP+PP draft-state synchronization issue and found during the required duplicate check that this PR already implements the same core fix. Per the repository's AGENTS.md, I am therefore not opening a duplicate PR without maintainer direction. I prepared a focused version of only the generic transport change on current
The unit test passes, and I also validated the focused implementation end to end on accelerator hardware with two consecutive concurrent request batches. Both batches completed without request errors, hangs, collective mismatches, or service-health regressions. This branch may be useful if maintainers prefer to split the generic MRV2 PP transport fix from the model-specific changes in this PR. Please let me know if an independent minimal PR is preferred; otherwise it is ready to cherry-pick or adapt here. This implementation and validation write-up were AI-assisted and reviewed against the upstream diff and runtime logs. |
Conflicts resolved against ~800 commits of upstream drift: - spec_decode/speculator.py, autoregressive/speculator.py: import-only conflicts from the upstream move of the multimodal registry out of the autoregressive speculator into the base class. - gpu/model_runner.py: propose() is now wrapped in use_workspace_lane and followed by adaptive_verification.record_confidences. Kept both upstream additions and re-applied intermediate_tensors= plus the broadcast_draft relay on top. Follow-up fixes required by the merge: - DFlashSpeculator.propose and MultiModuleMTPSpeculator.propose are new overrides that do not accept intermediate_tensors. The runner passes it unconditionally, so both raised TypeError at any PP size. Accept it (and ignore it -- neither drafter is PP-aware). - PPHandler gated its third broadcast on max_sample_len > 1, which is also true for diffusion LLMs. Those set num_speculative_tokens > 0 but have no speculator, so the last rank never sent the relay the other ranks waited for -- a collective op-count mismatch that hangs PP. Gate both sides on an explicit relay_draft_tokens flag derived from speculative_config instead. Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev>
|
@njhill Checking back on this one. I had closed the other PR due to the complexity concerns you mentioned. Let me know what else you'd like to see here, or if you'd prefer anything reworked. There seems to be at least some interest in the changes as a whole from the other comments here, so figured it was worth a small nudge. I've been carrying a custom build with these changes for a while, and I'd really like to get the pieces that make sense upstream rather than keep rebasing a patch set onto each new release. Appreciate the time, tyvm sir. |
|
I just tried your PR + #50288 and V2 works and nvfp4 works with it too, good work and thank you ^_^. |
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
Thanks for the work!
Please take a look at these AI generated comments
Could we remove the newly added `intermediate_tensors` plumbing altogether?
Specifically:
- Remove the argument passed from `model_runner.py`.
- Remove the added argument and forwarding logic from `BaseSpeculator`, `AutoRegressiveSpeculator`, `DFlashSpeculator`, and `MultiModuleMTPSpeculator`.
- Remove the intermediate-tensor copy in `AutoRegressiveSpeculator._run_model()`.
- Remove the zero-filled intermediate tensors created in `qwen3_5_mtp.py`.
The drafter is instantiated only on the last PP rank. `DeepSeekMTP` ignores these tensors, while the new Qwen last-rank path uses the target hidden states directly and does not consume them. Therefore, this plumbing does not provide a meaningful data flow and currently also causes the fused multi-step `TypeError`.
The `SupportsPP` implementation and `make_empty_intermediate_tensors` factory should remain, since they are required by the model interface check.| decode_fn = ( | ||
| self._fused_multi_step_decode | ||
| if self.use_fused_multi_step_decode | ||
| else self._multi_step_decode | ||
| ) | ||
| decode_fn( | ||
| num_reqs, | ||
| dummy_run and skip_attn_for_dummy_run, | ||
| decode_batch_desc, | ||
| num_tokens_across_dp, | ||
| input_batch.seq_lens_cpu_upper_bound, | ||
| intermediate_tensors=intermediate_tensors, | ||
| ) |
There was a problem hiding this comment.
vllm/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
Lines 528 to 535 in e349a56
_fused_multi_step_decode doesn't have this arg
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Signed-off-by: Chris Eastwood <106503529+eastwood-c@users.noreply.github.com>
|
Hey heads up there is another PR which kind of overlaps a bit which looks like we may be trying to land first #50514 |
|
@yewentao256 Thanks sir, I had made these changes, push is pending actual testing (which I should have done after adding the mid-stream changes to this pr anyways) and the below. @njhill Not a problem, I would gladly rebase and redo this pr (force push) to be stacked on top of that PR. I would ensure to validate against my running clusters (glm-5.2-fp8) and the qwen families on the combined work. Just let me know your preference sir Surviving changes summary
|
|
This pull request has merge conflicts that must be resolved before it can be |
7 participants
Purpose
MTP speculative decoding does not currently work under pipeline parallelism on the V2 model runner. This PR makes it functional for DeepSeek-family MTP drafts (DeepSeek-V3, GLM-5.2, Qwen3.5/3.6, …). Five independent issues, all on the PP>1 path. Fixes #1-#3 are DeepSeek-family-specific; fix #4 applies to all models using sparse MLA attention; fix #5 applies to Qwen3.5/3.6 MTP draft models on the last PP rank.
1.
DeepSeekMTPdoes not implementSupportsPP— DeepSeek-family specific. The engine refuses to build the draft model under PP at all:The MTP draft runs only on the last PP stage, so it never actually consumes PP intermediate tensors, but the interface still requires the
make_empty_intermediate_tensorsfactory. (This mirrors what #39704 does for the V1 runner.)2.
PPHandlersampled-token broadcast width mismatch (hang) — affects all MTP under PP.broadcast()sendssampled_token_idsat its natural width — 1 on any step with no draft tokens (prefill, first decode),num_spec+1once rejection sampling has run — whilereceive()always posts a fixed[num_reqs, max_sample_len]buffer. NCCLbroadcastdoesn't negotiate element counts, so a width-1 send against a width-max_sample_lenreceive is a count mismatch that deadlocks the receiver. Fix: pad the source tomax_sample_len(trailing-1, ignored bypost_update).3. Proposed draft tokens are never relayed to non-last PP ranks (garbage output / ~0 acceptance) — affects all MTP under PP.
req_states.draft_tokensis written only on the last rank (thepropose()path); non-last ranks keep the zero-init buffer.combine_sampled_and_draft_tokensthen embeds zeros at the draft positions on rank 0, so the verification input is wrong. Fix: coalesce a third broadcast (the proposed draft tokens) into the existing deferredPPHandlersibling-group broadcast, and scatter it intoreq_states.draft_tokenson consume. No new collective; gated identically to the sampled-token broadcast so per-step op counts stay matched.4. Stale
topk_indices_bufferreference in sparse MLA backends (the acceptance fix) — affects all models using sparse MLA attention. Under MTP+PP,FlashAttnMLASparseImpl.__init__storedindexer.topk_indices_bufferat construction time. When_maybe_share_lm_headlater replacedIndexer.topk_indices_bufferwith the target model's buffer, the impl's reference was stale — still pointing to the draft model's original (uninitialized) buffer. This caused garbage DSA attention → degenerate "repeat-the-current-token" drafts → ~27-33% acceptance instead of ~85%. Fix: storeself._indexer = indexerin__init__, readself._indexer.topk_indices_bufferdynamically inforward_mqa. Applied to all three sparse MLA backends:flashattn_mla_sparse.py,flashmla_sparse.py,flashinfer_mla_sparse.py.5. Apply
fcprojection on last PP rank for Qwen3.5 MTP — Qwen3.5/3.6-specific. Under MTP+PP, the Qwen3.5 MTP draft model on the last PP rank was using the target model's hidden_states directly, bypassing thefcprojection entirely. This produced essentially random predictions (~1% acceptance) because the draft model's input was not properly projected. Fix: on the last PP rank, apply the samefcprojection as the first rank (embedinput_ids, normalize, concat withhidden_states, project throughfc). This is the same pattern used on PP0 (first rank).Test Plan
Serve a DeepSeek-family or Qwen3.5/3.6 MTP model under PP on the V2 runner and check it boots, produces correct output, and accepts drafts at a normal rate:
Unit tests (
tests/v1/worker/test_pp_utils.py):test_deepseek_mtp_implements_supports_pp— verifies Fix#1test_pphandler_broadcast_pads_to_max_sample_len— verifies Fix#2test_sparse_mla_backend_reads_topk_indices_buffer_dynamically— verifies Fix#4Test Result
Validated on GLM-5.2-FP8 (GlmMoeDsaForCausalLM, DeepSeek-Sparse-Attention MoE), TP4/PP2, on a current-main base (0.23.1rc1.dev531), serving real traffic for 5+ hours at K=3:
The stale-buffer fix (fix #4) lifts acceptance from ~27-33% (pre-fix, fixes #1-#3 only) to 84.4% at K=3 over 5+ hours of real traffic — matching the non-PP TP8 baseline (~85%).
Cross-model validation: fix #4 applies broadly
The stale-buffer fix was further validated across 8 Qwen3.5/3.6 model variants (dense BF16, AWQ, MoE BF16, MoE AWQ, MoE GPTQ-Int4), each at MTP K=1/2/3, PP=2/TP=1. All show high acceptance (83-96%) that scales gracefully with K, confirming the fix is not architecture-specific.
Full per-position breakdown and GSM8K accuracy available on request.
Remaining stale-buffer backends
Fix #4 was also applied to the three remaining sparse MLA backends that were not covered by the original fix commit (
c175667db):flashinfer_mla_sparse_sm120.py—FlashInferMLASparseSM120Impl(SM120 variant)rocm_aiter_mla_sparse.py—ROCmAiterMLASparseImpl(ROCm)xpu_mla_sparse.py—XPUMLASparseImpl(Intel XPU)These backends have the same stale
topk_indices_bufferbug but are not reachable on our hardware (H200/SM90). The fix follows the exact same pattern as the already-validated fix: storeself._indexer = indexerin__init__, readself._indexer.topk_indices_bufferdynamically inforward_mqa.