[PP&Spec] enable speculative decoding (eagle_worker_v2) under PP by liuqi-1 · Pull Request #31139 · sgl-project/sglang

liuqi-1 · 2026-07-14T07:34:49Z

Roadmap: #11857

PP+Dspark is supported base on this PR, the link is #32281

Motivation

server_args.py previously asserted that pp_size > 1 and a speculative algorithm could not be enabled together. This PR removes that restriction and enables EAGLE speculative decoding to run under pipeline parallelism, tested on GLM52 and Deepseek-V4-Flash.

Modifications

The draft worker is created only on the last PP rank; every rank runs its own target forward and verify, the draft tree is carried across stages and rebuilt locally per rank, and the accepted-token KV bookkeeping is moved out of the verify epilogue into a per-rank post-processing step.

Draft worker created only on the last PP rank. Non-last ranks have no draft worker, draft KV pool, or draft lm_head; they forward target layers and relay hidden states / proxy tensors downstream. All draft-accessing paths are null-guarded.
Target hidden capture (capture_hidden = FULL) enabled only on the last rank; upstream ranks capture nothing and pass the hidden/proxy along.
The draft tree produced on the last rank (draft/bonus/parent indices, accept info) is serialized as plain lists, added to the existing PP tensor dict, and reconstructed on every other pp rank into a verify input matching that rank's own KV layout and attention buffers, so each rank builds its tree mask / positions locally.
Sampling / acceptance / grammar run only on the last rank; non-last ranks skip that epilogue and execute only the target-layer forward for the verified tokens.
Accept-token write-back and sequence-length update are moved into the scheduler's batch_result_processor, which runs on every PP rank; the old single-point finalize path is skipped under PP.
The scheduler takes a dedicated non-overlap path that threads pp_proxy_tensors through generation and triggers the per-rank cache update afterwards.

Other changes:

MTP/nextn model layer: the context-parallel split is applied to the spec hidden states under PP + MTP. On GLM-5.2 (chunked-prefill-size 8192) this lowers per-step draft_extend_for_prefill time on the last PP rank from 27 ms to 20 ms (The data is profiled on 2 8*H20 nodes).
Weight loading: the nextn draft's embed_tokens is loaded on the last rank directly rather than recycled from the target; non-last ranks keep skipping it. lm_head / embed access is split into separate hooks so the PP draft shares only lm_head.
CUDA graph compatibility: the static pp_proxy_tensors buffers are sized by max token count (not batch size); the decode graph runner copies the proxy tensors into the graph buffer before replay.

Limitations asserted: PP + spec disables mixed-chunk prefill and speculative_adaptive; only EAGLE algorithm are permitted under PP for now.

Accuracy Tests

Setup: DeepSeek-V4-Flash-FP8, single 8*H20 node;

GSM8K accuracy (200 examples, 5-shot, temperature=0)

Configuration	GSM8K Accuracy (%)
PP2-TP4 (no spec)	98.50
PP2-TP4 + MTP314 (this PR)	98.50
TP8 + MTP314	98.50

Speed Tests

Setup: DeepSeek-V4-Flash-FP8, single 8*H20 node; EAGLE draft with speculative-num-steps=3, speculative-eagle-topk=1, speculative-num-draft-tokens=4. Three configs sweep the input length (16k / 32k / 64k):

TTFT (ms)

Configuration	16k	32k	64k
PP2-TP4 (no spec)	653.59	1142.21	2197.28
PP2-TP4 + MTP314	658.30	1135.22	2124.52
TP8 + MTP314	534.70	1080.04	2212.80

TPOT (ms)

Configuration	16k	32k	64k
PP2-TP4	11.67	11.71	11.68
PP2-TP4 + MTP314 (this PR)	5.42	5.53	5.65
TP8 + MTP314	4.19	4.26	4.45

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to the Write documentations.
Provide accuracy and speed benchmark results according to the Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ❌ Run #32382552977
Latest PR Test (Extra): ❌ Run #32382552472
Latest PR Test (AMD ROCm 7.2): ❌ Run #32382552886

gemini-code-assist

Code Review

This pull request implements support for Eagle speculative decoding under Pipeline Parallelism (PP). It restricts the draft worker to the last PP rank, introduces EaglePPVerifyInputRaw for transmitting draft information across ranks, and updates the scheduler, weight loaders, and KV cache configurators to handle PP-specific speculative execution. The review comments identify critical shape mismatch bugs in draft token reshaping and dummy draft token creation, a potential AttributeError in the CUDA graph runner, and opportunities to optimize tensor allocation on the hot path in the batch result processor.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Copilot

Pull request overview

This PR enables EAGLE-v2 speculative decoding to run with pipeline parallelism (PP>1) by hosting the draft worker on the last PP rank only and relaying a compact “raw” draft tree to other PP stages so all ranks can perform target verification in lockstep.

Changes:

Removed the hard guard that forbade pp_size>1 with speculative decoding; added related feature gates (e.g., no adaptive spec under PP, no mixed-chunk + PP + spec-v2).
Implemented a PP-specific spec-v2 flow: last-rank drafting + cross-PP relay of draft-tree metadata; all ranks run verify.
Optimized NextN/CP by moving CP split ahead of token-wise draft projections; adjusted KV pool sizing so only the last PP stage reserves draft KV.

Reviewed changes

Copilot reviewed 20 out of 21 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
python/sglang/srt/speculative/spec_info.py	Adds a new SpecInputType for PP-relayed raw verify input.
python/sglang/srt/speculative/eagle_worker_v2.py	Core PP+spec-v2 orchestration: last-rank draft, cross-rank raw relay, verify changes.
python/sglang/srt/speculative/eagle_info.py	Introduces `EaglePPVerifyInputRaw` for PP tensor-dict relay and dummy-first-decode handling.
python/sglang/srt/speculative/adaptive_spec_params.py	Disables adaptive speculative decoding when PP is enabled.
python/sglang/srt/server_args.py	Lifts the PP vs speculative-decoding mutual exclusion (keeps overlap-schedule restriction).
python/sglang/srt/models/deepseek_v4.py	Adjusts NextN draft weight loading + adds head-only share/set helpers for PP.
python/sglang/srt/models/deepseek_v4_nextn.py	Moves CP split earlier for spec hidden states to reduce projection cost.
python/sglang/srt/models/deepseek_v2.py	Adds head-only share/set helpers for PP draft behavior.
python/sglang/srt/models/deepseek_nextn.py	Moves CP split ahead of token-wise projection for compute/memory savings.
python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py	Updates draft weight-loading rules (embed/head sharing) for PP last-rank draft design.
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py	Updates PPProxyTensors replay/return behavior under cuda-graph, incl. PP+spec copying.
python/sglang/srt/model_executor/runner_utils/buffers.py	Resizes PP proxy buffers to max token count (not max batch size).
python/sglang/srt/model_executor/pool_configurator.py	Gates draft-KV inflation to last PP rank only; avoids shrinking pools on other stages.
python/sglang/srt/model_executor/model_runner.py	Removes PP/MTP incompat assertion; adjusts memory profiling sync group for draft worker under PP.
python/sglang/srt/mem_cache/kv_cache_builder.py	Avoids creating draft KV pools on PP non-last ranks.
python/sglang/srt/managers/utils.py	Extends GenerationBatchResult with PP+spec relay fields (`pp_verify_input_raw`, accept_index).
python/sglang/srt/managers/tp_worker.py	Allows PP-local draft-worker seed reuse (no world broadcast) via `pp_global_random_seed`.
python/sglang/srt/managers/scheduler.py	Routes PP+spec-v2 batches through a PP-aware forward path; disallows mixed-chunk + PP + spec-v2.
python/sglang/srt/managers/scheduler_pp_mixin.py	Relays PP+spec raw tree via tensor dict; rebuilds dummy-first-decode draft info.
python/sglang/srt/managers/scheduler_components/batch_result_processor.py	Adds PP+spec KV compaction-before-free and seq_len advancement-after-free logic.
python/sglang/srt/managers/schedule_batch.py	Adds (or relocates) ScheduleBatch implementation used by the PP+spec plumbing.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

            draft_probs=draft_probs,
        )
+        return verify_input, parent_list, top_scores_index



dongyibo · 2026-07-14T08:45:01Z

nice work!

enternal111 · 2026-07-18T14:28:13Z

Also need PP with speculative decoding. Willing to test if it is ready to merge.

liuqi-1 · 2026-07-19T08:05:36Z

I updated this PR to add support for DSpark under Pipeline Parallelism (PP). It has been tested and verified working on Deepseekv4-pro-dspark. The branch has also been rebased onto main to minimize merge conflicts. Additionally, I noticed that another PR (#30775) also implements MTP under pipeline parallelism. Below is a comparison of the implementation approaches:

When loading the Draft weights, this PR reuses the existing load_weights method. The advantage is that it reuses the existing weight-loading mechanism; the downside is that every model needs to be adapted to use it. PR Pipeline parallelism x speculative decoding (EAGLE/MTP) compatibility #30775 instead implements a separate _load_checkpoint_tensor method in eagle_worker_v2.py to load the embedding weights. The advantage is that it adapts to different models directly; the downside is that it maintains an additional weight-loading path independent of load_weights.
Both PRs choose to generate the draft information (e.g., draft_tokens) for the next decode round on the last rank after verification. However, this PR stores the information in PPVerifyInputRaw (held within batch.specinfo), which has the advantage of being extensible to different speculative-decoding algorithms in the future. The downside is that it is less concise than PR Pipeline parallelism x speculative decoding (EAGLE/MTP) compatibility #30775, and the code changes are less self-contained. PR Pipeline parallelism x speculative decoding (EAGLE/MTP) compatibility #30775 stores the information in a new tensor dict in the scheduler, which is more concise.
This PR additionally implements support for PP + DSpark.

Perhaps we can find an optimal approach to make PP and speculative decoding compatible. In the future, this could be extended to support new speculative-decoding algorithms, or the PP + Spec feature could be adapted to PD (Prefill-Decode) disaggregation scenarios to further improve performance. Thank you.

liuqi-1 · 2026-07-29T09:56:10Z

I compared the original dummy bootstrap with generating a real draft tree after prefill.

Dummy bootstrap Real draft tree
After prefill Builds a dummy tree by repeating the bonus token Runs draft() and builds a real tree
First decode step Mainly serves as a bootstrap step Can verify real draft tokens immediately
Extra work No extra draft call One extra draft call after prefill
Implementation Simpler More state handling between prefill and decode
Environment: GLM5.2-FP8, TP8CP8PP2 + MTP 3,1,4, 16xH800 Benchmark: ShareGPT, 10 warmups, 100 requests, concurrency 16, PP 40/38.

Metric Dummy Real tree Change
Total throughput 1148.32 tok/s 1143.19 tok/s -0.45%
Mean TPOT 33.34 ms 31.70 ms -4.92%
Median TPOT 32.45 ms 30.30 ms -6.63%
P99 TPOT 59.21 ms 53.13 ms -10.27%
Mean TTFT 352.40 ms 359.83 ms +2.11%
Accept length 2.931 3.003 +2.46%
The real tree gives better decode latency, especially P99, while throughput is basically unchanged and TTFT is slightly worse.

@liuqi-1, could you help evaluate which approach fits this PR better? If you prefer the real draft-tree approach and would like to test it further, I can submit a PR to your branch.

If possible, could you open a PR against my branch (antgroup:sglang-communiy-main-pp-mtp)? I’ll test it first.

lmyybh · 2026-07-29T10:06:55Z

@liuqi-1 Okay, I'll submit a PR to you later. Also, I've already submitted a PR (merge the main branch ) to you earlier. Please take a look first.

Backs out 5140960 (cherry-pick of an earlier sgl-project#31139 revision, 95cc93a) and the five follow-ups built on it -- cd35832, 9a14913, 07f9d3a, 88e69f9, 770f501 -- plus e9da59e, so the current sgl-project#31139 head can be cherry-picked cleanly instead of merged into a diverged copy of itself. The PP/spec files are byte-identical to bf8af2f again; the DSA/SM80 and dp-attn work that landed later is untouched. 88e69f9 and 770f501 are the two commits that were upstreamed into sgl-project#31139 (117bffb, 4fdd3d8), so they come back with the new cherry-pick. e9da59e is not in sgl-project#31139 and has to be re-applied on top.

Cherry-picked from sgl-project/sglang PR sgl-project#31139 at its current head, commit 8c4c3ae, replacing the earlier revision 95cc93a that 5140960 had taken. The new revision is EAGLE-only: it drops the DSpark/dflash PP plumbing the old one carried (dspark_*, deepseek_v4_dspark, dflash_info_v2, multi_layer_eagle_worker_v2 are no longer touched), which is fine here because this deployment has no DSpark draft model. The PR is based on an upstream main 426 commits ahead of our fork point, so four files needed hand resolution: eagle_worker_v2.py signature takes pp_proxy_tensors but not grammar_barrier -- that plumbing does not exist in this tree, and the PR only ever passes it through to itself. The prefill-path condition keeps the un-gated form (see below). The PR's per-PP-rank capture_hidden_mode if/else is taken as-is. eagle_worker_common.py same: pp_proxy_tensors only. eagle_info.py keep this tree's BaseGrammarObject/envs imports and add the PR's ScheduleBatch import. tp_worker.py keep this tree's is_ep_scale_joiner name (upstream has since renamed it to is_ep_joiner) and put the PP draft-worker seed branch ahead of it. PR commit 117bffb ("run idle batch through prefill path") is not cherry-picked separately: its run_as_prefill change is subsumed by the un-gated form carried here, which also covers pp_size=1. PR commit 4fdd3d8 follows. Also folds in the two lockstep fixes from e9da59e, which sgl-project#31139 does not have. Without them tp16 + dp-attn8 + ep16 + EAGLE deadlocks on the first request: an idle DP rank must take the prefill path when a peer is extending (is_extend_in_batch=1) and must run the draft loop when the peers are decoding (is_extend_in_batch=0). Verified by A/B on 2-node 16xA100: reverting just this file hangs at a single 64-token request, restoring it answers in 0.7 s.

Cherry-picked from sgl-project/sglang PR sgl-project#31139, commit 4fdd3d8 (the same change this fork previously carried as 770f501, which the revert backed out).

ScheduleBatch.filter_batch calls spec_info.filter_batch(new_indices=..., has_been_filtered=False, new_indices_cpu=...) by keyword in this tree, but sgl-project#31139 was written against an upstream main whose SpecInput.filter_batch signature no longer has that parameter, so pp2 + EAGLE died with TypeError: EaglePPVerifyInputRaw.filter_batch() got an unexpected keyword argument 'has_been_filtered' as soon as a batch was filtered with chunked_req_to_exclude -- reached by gsm8k with several concurrent requests, not by a single-request probe. Accept and ignore the flag: these are plain per-request Python lists relayed from the last PP rank, so unlike EagleDraftInput they are never pre-filtered by verify() and always need indexing.

The sgl-project#31139 cherry-pick carried a 100644 -> 100755 mode change on this file; it is a plain module with no shebang and is never executed directly.

lmyybh · 2026-08-14T08:52:09Z

@liuqi-1 Hello, Are you still maintaining this PR? Should we follow the changes from the main branch?

liuqi-1 · 2026-08-14T09:16:25Z

@liuqi-1 Hello, Are you still maintaining this PR? Should we follow the changes from the main branch?

Sorry, I've been busy with other things recently. I've made some optimizations to the code in this PR, and I'll update it soon and rebase it onto the latest main branch.

…parallelism

Under PP + DP attention, an idle rank (no local requests) on a global prefill step must follow the prefill path, not the verify path. Routing it through verify emits speculative_num_steps extra draft-decode collectives that desync the MoE cross-DP all-gather and hang. Gate the prefill branch on is_extend OR (pp_enabled and idle and is_extend_in_batch), mirroring the scheduler's is_extend_in_batch lockstep. _draft_extend_for_prefill already short-circuits idle batches, so the target+single-draft_extend collective sequence stays identical across DP ranks.

need_topk gated only on spec_algo would still take the topk stash path when the payload carries no topk_p (topk_p is None), crashing at payload.topk_p[0]. Add the payload.topk_p is not None guard, matching the existing payload.hidden_states check for need_hidden_states.

liuqi-1 · 2026-08-18T06:09:15Z

@ShangmingCai Hi, this PR is ready to merge.

ziang663 · 2026-08-18T06:56:44Z

nice work! btw,does this support PD disaggregation

liuqi-1 · 2026-08-18T07:20:05Z

nice work! btw,does this support PD disaggregation

Not yet.

1e4ves · 2026-08-18T08:29:11Z

@ziang663 #33584 supported Prefil + PP + EAGLE

The unconditional even-rank-sends-first ordering was introduced for PP+MTP, where the EAGLE worker's spec verify state relays multiple GPU tensors in the output dict and every rank posting send first forms a ring wait of device P2P ops that deadlocks on CUDA. Keep the parity ordering only for EAGLE, the only algorithm the PP+MTP support in eagle_worker_v2.py was adapted and verified for. Everything else stays on the upstream default: on CUDA every rank sends first, on XPU adjacent stages are paired by parity.

…iants Three gaps found while comparing this branch against the other in-flight PP x spec efforts (sgl-project#31139, sgl-project#33863); all are confined to the gated path. - PP ring send/recv ordering. The relay carries several extra GPU tensors under spec (accept_lens, new_seq_lens, bonus tokens, next chain). Device P2P work stays ordered on the stream, so enqueueing that many sends before any recv on every rank can form a ring wait even on CUDA, where the existing comment assumes send-first is always safe. Pair adjacent stages by pp_rank parity when the spec relay is active, which makes rank 1 post its recv first and breaks the cycle for any pp_size > 1. Non-spec PP keeps the existing send-first behavior. Credit to sgl-project#31139 for finding this. - Chunked prefill middle chunks. Only the chunk that finishes the prompt samples a real token; a middle chunk's relayed next_token_ids is a placeholder that no consumer reads. Storing it seeded the next verify round with a token the model never emitted. - Batch composition invariants. Every stage rebuilds the same verify input from relayed per-request state, so all stages must see the same batch. Assert against --enable-mixed-chunk (splices decode requests into an extend batch) and --enable-dp-attention (partitions the batch per DP rank) rather than silently mis-rebuilding. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Copilot AI review requested due to automatic review settings July 14, 2026 07:34

liuqi-1 requested review from Fridge003, Qiaolin-Yu, ShangmingCai, Ying1123, alphabetc1, ch-wan, fzyzcjy, hanming-lu, hnyls2002, hzh0425, ispobock, merrymercy, xiezhq-hermann and yizhang2077 as code owners July 14, 2026 07:34

github-actions Bot added the deepseek label Jul 14, 2026

Copilot started reviewing on behalf of liuqi-1 July 14, 2026 07:35 View session

gemini-code-assist Bot reviewed Jul 14, 2026

View reviewed changes

Copilot AI reviewed Jul 14, 2026

View reviewed changes

ShangmingCai self-assigned this Jul 14, 2026

liuqi-1 force-pushed the sglang-communiy-main-pp-mtp branch from 350fb8e to de9d76b Compare July 15, 2026 15:35

liuqi-1 changed the title ~~[PP&Spec]enable speculative decoding (eagle_worker_v2) under PP~~ [PP&Spec]enable speculative decoding (MTP&DSpark) under PP Jul 19, 2026

liuqi-1 changed the title ~~[PP&Spec]enable speculative decoding (MTP&DSpark) under PP~~ [PP&Spec] enable speculative decoding (MTP&DSpark) under PP Jul 19, 2026

liuqi-1 force-pushed the sglang-communiy-main-pp-mtp branch from de9d76b to c582834 Compare July 19, 2026 07:06

github-actions Bot added the speculative-decoding label Jul 19, 2026

liuqi-1 mentioned this pull request Jul 19, 2026

[Tracking] DSpark in SGLang: roadmap #30344

Open

liuqi-1 force-pushed the sglang-communiy-main-pp-mtp branch from c582834 to 95cc93a Compare July 19, 2026 09:08

liuqi-1 force-pushed the sglang-communiy-main-pp-mtp branch from 9657bee to 4fdd3d8 Compare July 29, 2026 11:13

ehuaa added a commit to ehuaa/sglang that referenced this pull request Jul 30, 2026

chore: restore schedule_batch.py to mode 644

8aa2e60

The sgl-project#31139 cherry-pick carried a 100644 -> 100755 mode change on this file; it is a plain module with no shebang and is never executed directly.

liuqi-1 force-pushed the sglang-communiy-main-pp-mtp branch 2 times, most recently from f4df794 to a747ca5 Compare August 17, 2026 03:50

liuqi-1 requested a review from huangtingwei9988 as a code owner August 17, 2026 03:50

liuqi-1 force-pushed the sglang-communiy-main-pp-mtp branch from a747ca5 to a741584 Compare August 18, 2026 06:03

liuqi-1 and others added 4 commits August 18, 2026 14:06

[PP&MTP]Enable speculative decoding (eagle_worker_v2) under pipeline …

c974536

…parallelism

perf(spec): relay PP verify_input_raw over the GPU channel

12abd13

liuqi-1 force-pushed the sglang-communiy-main-pp-mtp branch from a741584 to 12abd13 Compare August 18, 2026 06:07

style: format pool_configurator.py with pre-commit (black)

e7b9a49

liuqi-1 force-pushed the sglang-communiy-main-pp-mtp branch from fe979cb to 09f8b70 Compare August 20, 2026 14:50

amd-bot mentioned this pull request Aug 22, 2026

[CI Monitor] Daily Report - 2026-08-22 bingxche/sglang-ci-bot#174

Open

AliceChenyy mentioned this pull request Aug 27, 2026

Pipeline parallelism x speculative decoding (EAGLE/MTP) compatibility #30775

Open

YAMY1234 mentioned this pull request Sep 15, 2026

[Tracking] Pipeline parallelism x speculative decoding (EAGLE/MTP): converge prefill gate PRs, make the path model-agnostic, decide decode design #39634

Open

7 tasks

Uh oh!

Conversation

liuqi-1 commented Jul 14, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests

Checklist

Review and Merge Process

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dongyibo commented Jul 14, 2026

Uh oh!

enternal111 commented Jul 18, 2026

Uh oh!

liuqi-1 commented Jul 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

liuqi-1 commented Jul 29, 2026

Uh oh!

lmyybh commented Jul 29, 2026

Uh oh!

lmyybh commented Aug 14, 2026

Uh oh!

liuqi-1 commented Aug 14, 2026

Uh oh!

liuqi-1 commented Aug 18, 2026

Uh oh!

ziang663 commented Aug 18, 2026

Uh oh!

liuqi-1 commented Aug 18, 2026

Uh oh!

1e4ves commented Aug 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

liuqi-1 commented Jul 14, 2026 •

edited by github-actions Bot

Loading

liuqi-1 commented Jul 19, 2026 •

edited

Loading