Repository navigation
Conversation
liuqi-1
requested review from
Qiaolin-Yu,
Ying1123,
alphabetc1,
fzyzcjy,
hnyls2002 and
merrymercy
as code owners
July 24, 2026 04:07
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
liuqi-1
requested review from
Fridge003,
ShangmingCai,
ch-wan,
hanming-lu,
hzh0425,
ispobock,
xiezhq-hermann and
yizhang2077
as code owners
July 24, 2026 04:07
5 tasks
7 of 18 tasks
liuqi-1
force-pushed
the
sglang-community-main-pp-dspark
branch
from
August 17, 2026 06:50
a17d374 to
379f1b6
Compare
liuqi-1
requested review from
HaiShaw,
hebiao064,
huangtingwei9988,
iforgetmyname,
ping1jing2,
whybeyoung and
yuan-luo
as code owners
August 17, 2026 06:50
liuqi-1
requested review from
BBuf,
DarkSharpness,
HydraQYH and
celve
as code owners
August 17, 2026 06:50
liuqi-1
changed the title
liuqi-1
changed the title
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.
…ommits Organically squash-merged from the kimi_k3_rebase_731 work (cd71a396, b6da9197, 3c04aff8, 12904266) and adapted to this branch's PP&MTP + GPU-relay baseline: - DSpark under PP: the last rank relays DSparkPPVerifyInputRaw around the PP ring so every rank rebuilds an identical verify context each iter; the draft model is built only on the last rank while non-last ranks relay the target verify forward; DSpark aux hidden states ride the PP proxy buffers (token-major dspark_aux). - Kimi-K3 hybrid-mamba state commit after PP DSpark verify: non-last ranks commit accepted mamba states via commit_pp_mamba_states_after_ verify, passing explicit state / conv source indices through the fused scatter kernel (src_indices_tensor), bypassing the conv multi-launch fast path only when scratch rows are explicit. - Mamba state commit for PP MTP (Eagle): commit_mamba_states_after_verify gains state_indices_tensor / scratch_source_indices_tensor and supports linear-chain verify (accept_index=None) via per-req step positions; Eagle keeps the GPU-tensor relay path and the module-level move_accept_tokens_to_target_kvcache. - Fix ScheduleBatch NameError in eagle_info.py under PP2+DSPARK via a TYPE_CHECKING import.
Co-Authored-By: Claude <noreply@anthropic.com>
…ll ranks Extracted from f2fb9f5e7b (speculative: support separate draft KV cache dtype). - _profile_available_bytes: query free memory with distributed=False so each rank budgets from its own local free memory instead of an all-reduced cluster-wide figure. - _apply_token_constraints: sync the resolved token capacity across every model rank (world group), not only PP ranks, so pure TP/DP workers also settle on one shared capacity.
Adapted from 3928e9b93 to this branch's DSparkPPVerifyInputRaw. DSparkPPVerifyInputRaw now carries GPU tensors (bonus_tokens [bs], draft_tokens [bs, gamma], accept_lens [bs], confidence [bs] f32, cap_trim_lens [bs] i32, verify_lens [bs] i64) so the PP ring transmits them over the GPU channel via the nested-dict flattening already provided by _split_tensor_dict (a747ca5). new_seq_lens stays a CPU list because it feeds only the last rank's sampling cache keyed by Python ints. - DSparkPPVerifyInputRaw: tensor fields; build_dummy_for_decode / filter_batch / merge_batch operate on tensors (torch.cat / tensor index). - DSparkWorkerV2: build the relay raw from accept/proposal tensors instead of .tolist(); rebuild verify candidates from pp_raw tensors directly. - batch_result_processor / scheduler_pp_mixin: accept_lens handled as a tensor; DSpark and Eagle now share the same tensor path.
The DSpark PP check used bare enable_dp_attention, which is forced on by DeepSeek-V4 prefill context parallel (dp_size==1) as an execution flag reuse. That topology has no attention data-parallel ranks, so it should not be rejected. Align the check with the DP-only restriction (enable_dp_attention and dp_size>1) already used earlier in _handle_dspark.
liuqi-1
force-pushed
the
sglang-community-main-pp-dspark
branch
from
August 18, 2026 07:21
379f1b6 to
4b63351
Compare
…k_draft_sampler dspark_worker_v2 imported resolve_greedy_mask from dspark_draft_sampler, but the function is defined in dspark_draft. Split the import: keep maybe_build_draft_sampler from dspark_draft_sampler and take resolve_greedy_mask from dspark_draft, matching the pp_spec branch layout. This fixes ImportError at draft worker init under DSpark.
liuqi-1
changed the title
Contributor
Contributor
|
I would like to test it on H200 with dsv4-0731, but i can not patch both this pr and 33614 into nightly image. Without 33614 crash will happen. Could u deal with it so i can test @liuqi-1 |
…k) only The unconditional even-rank-sends-first ordering was introduced for PP+spec, where the spec worker's 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 and DSPARK, the algorithms the PP support in eagle_worker_v2.py / dspark_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.
Author
3 of 5 tasks
10 tasks
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
roadmap: #11857, #30344, #32607
Our previous PR (#31139) enabled PP + EAGLE support. Testing showed DSpark speculation performs far better than EAGLE, so this PR builds on it to add DSpark + PP compatibility. This PR depends on the previous one and should be merged afterwards.
DSpark's verify path is structurally different from EAGLE's (per-layer aux hidden capture, linear block verify, confidence-driven budget), so it could not reuse the EAGLE PP machinery verbatim. This PR lifts that ban and makes DSpark run under PP, targeting DeepSeek-V4-Flash + the bundled DSpark draft checkpoint.
Modifications
The core idea mirrors the EAGLE PP design — the draft model lives only on the last PP rank, every rank runs its own target verify forward, and the draft proposal is relayed across stages so each rank rebuilds an identical verify context. The DSpark-specific work is making the per-layer aux-hidden capture and the linear block verify cross the PP boundary correctly.
Per-rank draft proposal carrier (
DSparkPPVerifyInputRaw). A new spec input type carries the last rank's draft / confidence / accounting fields (bonus tokens, draft tokens, new seq lens, confidence, accept/cap-trim/verify lens) across the PP ring. Every rank rebuilds an identical verify layout from the same source each iteration, so all PP ranks produce byte-identical verify windows. It mixes in the sharedDFlashDecodePrepareMixinso the non-last rank's verify forward (which readsverify_cache_locfromreq_to_token) pre-allocates KV headroom exactly like the last rank. Linear verify reuses the formalreq_to_tokenslots, soaccept_indexis alwaysNoneand the token-move step inbatch_result_processoris skipped automatically.Per-layer aux hidden capture across PP stages (
deepseek_v4.py). DSpark's draft needs per-layer aux hidden states (dspark_aux_hidden_states), but each PP rank only executes layers in its own[start_layer, end_layer)range. The aux list is now stacked as[L, num_tokens, hidden], relayed downstream under the"dspark_aux"key inPPProxyTensors, and rebuilt in global layer order on the last rank (upstream layers first, then local captures).set_dspark_layers_to_captureno longer early-returns on non-last ranks, so every rank registers its own capture range.PP-aware verify forward.
pp_proxy_tensorsis threaded throughTargetVerifyExecutor.run/run_non_compactso the target verify forward relay hidden states along the PP ring on every rank.TargetVerifyResultcarriespp_hidden_states_proxy_tensorsback to the scheduler. Logits adjustments (penalizer / vocab mask / logit bias) are now gated onresult.logits_output is not None, since non-last ranks never sample and their logits output isNone.Embed / weight loading split (
deepseek_v4_dspark.py). The DSpark draft used to borrowembed_tokensfrom the target. Under PP the last rank builds its ownVocabParallelEmbeddingand onlylm_headis shared from the target;attach_shared_modulesnow takesembed_tokens: Optionalso the PP path passesNoneand the draft keeps its locally-built embedding. Weight-name remapping keepsembed_tokens.weighton the last PP rank.PP proxy buffer sizing (
base_runner.py). The staticpp_proxy_tensorsbuffers are allocated by max token count, not batch size, so the dummy / capture buffer fits the verify-token-expanded shape[num_tokens, *hidden]produced under speculative decoding.Planner / draft fixes for the relayed layout.
alloc_verify_windowandDraftBlockProposerderiveprefix_lensfrom the verify window's anchor position instead ofbatch.seq_lens, because the last rank's re-propose pass builds the next window from post-iter seq lens whilebatch.seq_lensis still the pre-iter value. The planner forcesRAGGED_VERIFY_MODE=STATICand non-compact (eager) verify under PP, since the confidence-driven ragged scheduler and the cuda-graph compact path only run on the last rank and cannot keep cross-rank layout identity.Supporting changes:
draft_worker_common.build_draft_tp_workeraccepts app_global_random_seedso the draft sampler on the last rank is seeded consistently with the target;speculative_hookreplaces the blanketpp_size != 1ban with a narrower "PP + DP-attention unsupported" check (DP-attention partitions the batch, which conflicts with the PP requirement that every rank see an identical batch).Limitations asserted now
--dp-size 1is permitted under PP.Accuracy Tests
Model: DeepSeek-V4-Flash-FP8
GSM8K accuracy (200 examples, 5-shot, temperature=0):
Speed Tests
Setup: DeepSeek-V4-Flash-FP8, single 8*H20 node; three PP2+TP4 configs:
speculative-num-steps=3,speculative-eagle-topk=1,speculative-num-draft-tokens=4)bench_serving:
--dataset-name randomwith local ShareGPT path (--dataset-path /datasets/ShareGPT_V3_unfiltered_cleaned_split.json),--random-range-ratio 1.0,--num-prompts 5,--max-concurrency 1,--random-output-len 1024,--seed 42; 3 rounds per point, reporting the min.TPOT (ms)
Average accepted length (tokens / iter)
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Missing
run-cilabel -- add it to run CI tests.Latest PR Test (Extra): ❌ Blocked --
run-ciis required first.Latest PR Test (AMD ROCm 7.2): ➖ No AMD PR run found for this commit.