Pipeline parallelism x speculative decoding (EAGLE/MTP) compatibility

[AliceChenyy]

Motivation

PP and speculative decoding are currently mutually exclusive (server_args.py asserts). On PCIe-only multi-GPU boxes (no NVLink), TP-only scaling is bottlenecked by per-layer all-reduces, and PP is the right scaling axis. Combining the two requires PP x spec compatibility, which vLLM has (vllm-project/vllm#16568) and SGLang doesn't.

This PR implements it behind SGLANG_ENABLE_PP_SPEC=1; default behavior is unchanged.

Dependencies: none outstanding. #26928 (SM120 GLM DSA) merged on 2026-07-28 and is in main. The former companion PRs #30774 (proxy buffer fixes) and #30994 (dense MLA kwargs) are closed — #30774's fixes are folded into this diff, #30994 is no longer needed on the sparse path.

Design

Key observation: non-spec PP decode already relies on a last->first output-dict relay (next sampled token becomes the next round's input, and every stage forwards the dict). Spec decoding is structurally identical — the relay payload just grows to {accept_lens, new_seq_lens, bonus_token, next draft tree}.

1. Draft worker only on the last stage (needs final hidden states + lm_head).
2. Tail-drafting: the last stage drafts for round r+1 at the end of round r (verify -> draft_extend -> draft). Round r+1's verify tokens must exist before stage 0 runs its half, so the single-stage head-of-iteration drafting cannot work under PP.
3. The relayed tree is a first-class SpecInput: PPSpecRelayInput carries the proposed tokens plus the topology they were arranged by (parent_list / top_scores_index), rides on ScheduleBatch.spec_info, and implements the filter_batch / merge_batch hooks, so the existing recomposition path carries it — a request that finishes, is retracted, or merges in from a just-finished prefill takes its row along. Every stage rebuilds an identical EagleVerifyInput from it via build_tree_kernel_efficient, so no draft model is needed off the last stage. The type is deliberately the algorithm-agnostic half: another speculative algorithm supplies only its own rebuild.
4. Composition-safe by rid: the relayed tensors are sized by the composition that ran the forward, which the live microbatch can have moved past by the time the result comes back around the ring, so rows are folded in by rid rather than by position.
5. KV bookkeeping is free: kv_committed_len advances in the shared result processor from the relayed CPU accept_lens; verify slots stay within eagle_prepare_for_decode's 2x reserve (same mechanism the overlap path uses).

A subtle but load-bearing detail: non-last stages must run eagle_prepare_for_verify inside scheduler._forward_isolation — it mutates ScheduleBatch.forward_mode to TARGET_VERIFY, and without the restore the next get_next_batch_to_run treats the decode batch as extend and re-merges it (duplicate requests).

Two hazards that only appear once the relay carries several GPU tensors, both fixed here:

• PP ring ordering. Device P2P work stays ordered on the stream, so enqueueing several 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. Adjacent stages are paired by pp_rank parity when the spec relay is active; non-spec PP keeps the existing behavior. Credit to [PP&Spec] enable speculative decoding (eagle_worker_v2) under PP #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, and storing it seeded the next verify round with a token the model never emitted.

Validation

4u8g-gen-0210, 8x RTX PRO 6000 Blackwell Server (SM120, PCIe, no NVLink), driver 595.58.03, image lmsysorg/sglang:nightly-dev-cu13-20260826-41e7612d, torch 2.13.0+cu130 / sgl-kernel 0.4.6.post1 / flashinfer 0.6.17 / sgl-deep-gemm 0.1.5.post3. GLM-5.2-NVFP4, FP8 KV, EAGLE (2 steps / topk 1 / 3 draft tokens). Plain NCCL all-reduce on both configs.

GSM8K 200 questions, 8-shot, greedy, two runs per config — see the caveat below on why one run is not enough.

	TP4xPP2+MTP (gate on)	TP8+MTP (gate off)
GSM8K x2	0.940 / 0.935	0.955 / 0.940
Invalid	0.000	0.000
accept_length	2.79 - 3.00	2.68 - 3.00
8K bs1 TTFT / TPOT	1721 / 9.87	1888 / 8.73
8K bs8 TTFT / TPOT / tok/s	5521 / 30.28 / 154.5	11032 / 30.55 / 108.7
8K bs32 TTFT / TPOT / tok/s	13189 / 95.3 / 218.0	35897 / 105.2 / 120.1
KV pool	389,184	241,728

Accuracy and accept length match; PP wins TTFT (2.0x at bs8, 2.7x at bs32), throughput (+42% / +81%) and KV capacity (1.61x).

Correction to the earlier revision of this description. It claimed PP beating TP8 by 3-6x across TTFT/TPOT/throughput. That was measured in July; re-running the same benchmark on today's main, TP8+MTP is now 13% faster than PP at bs1 TPOT (8.73 vs 9.87 ms) — TP8's bs1 TPOT improved 1.85x and its bs32 throughput 2.6x over that period, while the PP numbers barely moved. The cause looks structural: PP+spec forces --disable-overlap-schedule, and the overlap path is where much of that improvement landed, so PP cannot pick it up. Making PP+spec work under the overlap scheduler is the obvious follow-up and probably matters more than anything else in this PR. PP's remaining advantages — TTFT, throughput at batch, KV capacity — all grow with concurrency.

Measurement caveat. GSM8K on this stack has +-2-3% run-to-run variance: the same build gave 0.940 and 0.965 on consecutive 200-question runs, because batch composition follows arrival timing and the reductions follow batch composition. Greedy output is not token-reproducible here either. So single runs cannot resolve differences of this size, and the accuracy claim above is "no detectable regression", not bit-identical behavior.

Also validated on DeepSeek-V4-Flash FP8, TP2xPP2 (4x RTX 6000D): GSM8K 0.95-1.00, accept_length 2.44-2.72 — the PP+spec layer carried over with only two portability fixes (checkpoint embed tensor name, backend-agnostic context-len bound).

Tests

test/registered/pp/test_pp_spec.py:

• TestPPSpecConsistency runs GSM8K under plain spec and under PP2+spec and requires the PP run to match on both accuracy and accept length. Accuracy alone cannot distinguish a working relay from a broken one: if the relayed tree is lost, the bonus token is still force-accepted and the rest rejected, so the output stays correct and only acceptance collapses. Both a topk=1 chain and a topk=2 tree are covered.
• TestPPSpecGate pins that the gate is off by default (the existing PP+spec ban still fires) and that the combinations the relay cannot reproduce identically on every stage are rejected. It resolves argument state only, so it stays off the GPU.

Known gaps

• Non-overlap schedule only (PP itself requires that today). Per the correction above, this is now the main thing holding PP back at low batch, not a cosmetic limitation.
• PP + DP attention is rejected at startup: DP attention partitions the batch per DP rank, so the stages would no longer rebuild the same verify tree. [PP&Spec] enable speculative decoding (eagle_worker_v2) under PP #31139 handles this case; this PR does not yet.
• Adaptive spec (--speculative-adaptive), --speculative-token-map, and EAGLE3 drafts are rejected at startup (the relay/head-wiring paths don't cover them yet).
• topk>1 trees are implemented (the drafted topology is relayed) but are exercised only by the unit test above — the SM120 model family available to us runs the DSA backend, where upstream rejects topk>1 with page_size>1, and DSA forces page_size=64.
• One D2H sync per round in tail-drafting (optimizable).

Relationship to the other PP x spec PRs

#31139 (PP + EAGLE) reached essentially the same design independently — draft on the last stage, tail-draft, relay the raw tree, rebuild per stage, proxy buffers sized by token count. That convergence is probably the most useful signal here: the shape of the solution is not in dispute. Differences worth a maintainer's attention:

• This PR touches no model files; the draft embedding is read from the checkpoint by a candidate-name list, so a new model family needs no adaptation (validated on GLM-5.2 and DeepSeek-V4-Flash). [PP&Spec] enable speculative decoding (eagle_worker_v2) under PP #31139 reuses load_weights, which is the more idiomatic mechanism but carries per-model changes.
• This PR is gated default-off and is the only one of the group with tests.
• [PP&Spec] enable speculative decoding (eagle_worker_v2) under PP #31139 supports PP + DP attention; this PR asserts it off.

Whichever base the maintainers pick, happy to contribute the pieces that are specific to this one — the model-agnostic embed loading, the ring-ordering fix, the tests, and the SM120 / PCIe validation.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ⏳ Run #35013704529
Latest PR Test (Extra): 🚫 Run #35013704124
Latest PR Test (AMD ROCm 10): ❌ Run #35013704196

[gemini-code-assist]

Code Review

This pull request introduces support for Pipeline Parallelism combined with Speculative Decoding (PP+spec) in SGLang, ensuring draft models run only on the last PP stage and verify inputs are correctly rebuilt and relayed across stages. Feedback on the changes highlights several critical improvements: preventing a GPU memory leak and fragmentation by storing draft chains on the CPU, overriding PP parameters for the draft worker's ModelRunner to prevent embedding weights from being filtered out, avoiding potential KeyError and AttributeError crashes with safer attribute and dictionary accesses, and using get_embed_and_head() instead of fragile hardcoded attribute paths.

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.

[gemini-code-assist]

Storing slice views of the rows tensor (i.e., rows[i]) in the persistent dictionary self._pp_spec_chain_by_rid keeps the entire rows tensor's underlying storage alive in GPU memory. Since requests finish at different times, this will cause a severe GPU memory leak and memory fragmentation over time.

Additionally, storing these very small 1D tensors (size 4-16) on the GPU introduces unnecessary CUDA allocation overhead and fragmentation. Moving the tensors to CPU and cloning the slice completely resolves both issues.

Suggested change

	if chain_tokens is not None:
	rows = chain_tokens.to(torch.int64).reshape(
	len(batch.reqs), num_draft_tokens
	)
	else:
	rows = torch.zeros(
	(len(batch.reqs), num_draft_tokens),
	dtype=torch.int64,
	device=bonus_tokens.device,
	)
	rows[:, 0] = bonus_tokens.to(torch.int64)
	for i, req in enumerate(batch.reqs):
	if req.finished():
	self._pp_spec_chain_by_rid.pop(req.rid, None)
	else:
	self._pp_spec_chain_by_rid[req.rid] = rows[i]
	if chain_tokens is not None:
	rows = chain_tokens.to(device="cpu", dtype=torch.int64).reshape(
	len(batch.reqs), num_draft_tokens
	)
	else:
	rows = torch.zeros(
	(len(batch.reqs), num_draft_tokens),
	dtype=torch.int64,
	device="cpu",
	)
	rows[:, 0] = bonus_tokens.to(device="cpu", dtype=torch.int64)
	for i, req in enumerate(batch.reqs):
	if req.finished():
	self._pp_spec_chain_by_rid.pop(req.rid, None)
	else:
	self._pp_spec_chain_by_rid[req.rid] = rows[i].clone()

[gemini-code-assist]

The draft worker's ModelRunner is initialized with pp_size > 1 and pp_rank = pp_size - 1. This causes the weight loader to filter out the draft model's embedding weights (thinking they belong to stage 0), leaving them randomly initialized.

Since the draft worker runs entirely locally on the last stage and does not participate in pipeline parallel partitioning, we can elegantly solve this "known gap" by overriding self.pp_size = 1 and self.pp_rank = 0 early in ModelRunner.__init__ when self.is_draft_worker is True. This completely bypasses PP filtering for the draft worker.

Suggested change

	if self.pp_size > 1:
	assert (
	self.support_pp
	), "Pipeline Parallel is not compatible with this model."
	import os as _os
	if not (_os.environ.get("SGLANG_ALLOW_PP_SPEC") and self.is_draft_worker):
	assert (
	self.support_pp
	), "Pipeline Parallel is not compatible with this model."
	import os as _os
	if getattr(self, "is_draft_worker", False) and _os.environ.get("SGLANG_ALLOW_PP_SPEC"):
	self.pp_size = 1
	self.pp_rank = 0
	if self.pp_size > 1:
	assert (
	self.support_pp
	), "Pipeline Parallel is not compatible with this model."

[gemini-code-assist]

If a request is missing from self._pp_spec_chain_by_rid (e.g., due to preemption, rescheduling, or other state inconsistencies), accessing self._pp_spec_chain_by_rid[req.rid] directly will raise a KeyError and crash the scheduler. Using .get() with a safe default row of zeros prevents this crash.

Suggested change

	chain_rows = torch.stack(
	[self._pp_spec_chain_by_rid[req.rid] for req in batch.reqs]
	).to(device=device, dtype=torch.int64)
	default_row = torch.zeros(num_draft_tokens, dtype=torch.int64, device="cpu")
	chain_rows = torch.stack(
	[self._pp_spec_chain_by_rid.get(req.rid, default_row) for req in batch.reqs]
	).to(device=device, dtype=torch.int64)

[gemini-code-assist]

Hardcoding self.draft_runner.model.model.embed_tokens.weight is fragile and can break across different model architectures (e.g., those using different embedding attribute names). Using the standard get_embed_and_head() method is much more robust and maintainable.

Suggested change

	embed = self.draft_runner.model.model.embed_tokens.weight
	head = self.target_worker.model_runner.model.lm_head.weight
	self.draft_runner.model.set_embed_and_head(embed, head)
	embed, _ = self.draft_runner.model.get_embed_and_head()
	head = self.target_worker.model_runner.model.lm_head.weight
	self.draft_runner.model.set_embed_and_head(embed, head)

[enternal111]

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

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!