[DSpark] Support pipeline-parallel targets in aggregated serving - #56956
Open
lucifer1004 wants to merge 1 commit into
Open
lucifer1004 wants to merge 1 commit into
lucifer1004 wants to merge 1 commit into
Conversation
Draft tokens under PP: the last stage runs the speculator and now broadcasts the fresh drafts to earlier stages (whose next verification step would otherwise embed stale buffer contents), with the double-post on the pp_broadcast group gated out for the speculator-less diffusion path. Non-last stages JIT-compile the deferred post-update kernel during warmup so its first compile cannot deadlock the pipeline mid-serving. Draft embedding under PP: the target's embedding table lives on the first stage, so DeepSeek-V4/Kimi-K3 DSpark drafters load their own copy from the checkpoint (loads_own_embed_under_pp) instead of aliasing. Padded graph batch safety: the DFlash prepare-inputs kernel now clears input_ids/positions and sets is_padding on CUDA-graph padding rows, and the DSv4 top-k router zeroes padded-row selections instead of reading uninitialized state. Co-authored-by: Kimi Code <noreply@moonshot.cn> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
lucifer1004
requested review from
AndreasKaratzas,
DarkLight1337,
WoosukKwon,
mgoin,
njhill,
pavanimajety,
tlrmchlsmth,
yewentao256,
ywang96 and
zyongye
as code owners
September 15, 2026 06:07
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
mergify
Bot
added
deepseek
GirasoleY
reviewed
Sep 15, 2026
Comment on lines
+163
to
+168
| if _target_pp_world_size() > 1: | ||
| self.embed_tokens = VocabParallelEmbedding( | ||
| self.config.vocab_size, | ||
| self.config.hidden_size, | ||
| prefix=maybe_prefix(prefix, "embed_tokens"), | ||
| ) |
Contributor
There was a problem hiding this comment.
#50514 added token embedding at pp last stage, I think we don't need this anymore (same for the dsv4 path)
| # window and restore it before serving. | ||
| pp_handler = getattr(self.model_runner, "pp_handler", None) | ||
| if pp_handler is not None: | ||
| pp_handler.set_disabled(True) |
Contributor
There was a problem hiding this comment.
Could you help me understand how the deadlock happen? If it's about collective running on multiple streams, I'd prefer enforce collective ordering instead of skip warm up here.
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
2 participants
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.
What
DSpark on a pipeline-parallel target in aggregated (non-PD) serving. The drafter runs wholly on the last PP stage; this PR makes that work correctly:
PPHandler.broadcast_drafts), and the broadcast is gated out of the no-op second call site when a speculator ran (double-posting misaligned the recv FIFO and hung the pipeline).warmup_pp_decode_updatecompiles it during capture, and the sampled-token broadcast is disabled for the warmup window.loads_own_embed_under_pp) instead of aliasing.input_ids/positionsand setsis_paddingon CUDA-graph padding rows, and the DSv4 top-k router zeroes padded-row selections instead of reading uninitialized state.Not a duplicate
Split from #53577 (reviewer request): this PR is the aggregated-serving half; the PD prefill + KV-transfer half follows separately.
Tests
tests/v1/worker/test_pp_utils.py,test_spec_decode_embed_sharing_pp.py,tests/v1/spec_decode/test_dflash_prepare_inputs.py,tests/models/kimi_k3/test_eagle3.py,test_dspark_mla.py,tests/kernels/moe/test_topk_softplus_sqrt.py.pytest tests/v1/worker/test_pp_utils.py tests/v1/worker/test_spec_decode_embed_sharing_pp.py tests/models/kimi_k3/test_eagle3.py tests/models/kimi_k3/test_dspark_mla.py tests/v1/spec_decode/test_dflash_prepare_inputs.py tests/v1/worker/test_gpu_warmup_blocks.py tests/v1/worker/test_gpu_model_runner_v2.py tests/kernels/moe/test_topk_softplus_sqrt.py -q→ 1744 passed.Model evaluation
DeepSeek-V4-Flash (0731) DSpark K=5, IFB PP2×TP2, GSM8K strict-match: 0.9545 (matches the target-only arm within noise).
Kimi-K3 DSpark IFB PP additionally needs #51065 (without it, the draft's non-causal flag raises the shared KV group's TritonMLA reorder threshold and misroutes the causal target's verification blocks).
AI assistance was used in preparing this PR.