[DSpark] Support PP prefill with KV transfer (disaggregated serving) - #56957
Draft
lucifer1004 wants to merge 3 commits into
Draft
lucifer1004 wants to merge 3 commits into
lucifer1004 wants to merge 3 commits into
Conversation
mergify
Bot
added
deepseek
lucifer1004
force-pushed
the
pr/dspark-pp-pd
branch
from
September 15, 2026 13:25
7bd3602 to
16523b6
Compare
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
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>
On current main the target model retains its embedding table on the last PP rank when speculative decoding is active (spec_decode_needs_target_embed covers dspark), so the last-stage drafter aliases it via maybe_share_target_embed. loads_own_embed_under_pp is redundant. Co-authored-by: Kimi Code <noreply@moonshot.cn> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Stacked on the aggregated-serving PP support. Under PD disaggregation with a pipeline-parallel producer, the DSpark drafter runs in prefill-only mode: it materializes the draft's context KV on the producer stages and transfers named per-layer KV regions, instead of running propose() inline. - SpeculativeConfig.target_kv_transfer_config + is_dspark_prefill_only() (and a NIXL guard: it cannot describe packed per-layer KV regions). - The scheduler narrows use_eagle (and eagle block-drop) for the prefill-only producer so it never does the EAGLE-style read-ahead. - DFlash gains materialize_context_kv(): precompute and store the draft context KV without generating draft tokens; the model runner calls it in place of propose() on the producer and skips draft-side rejection sampling, drafter graph capture, and EPLB registration there. - The DSv4 draft builds only the attention submodule needed for prefill (context_kv_only), avoiding transient MoE expert allocation, and loads just the context-KV weights. Following the aggregated-serving branch, the drafter aliases the target embedding via vllm-project#50514 on the last PP stage instead of loading its own copy. Co-authored-by: Kimi Code <noreply@moonshot.cn> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
lucifer1004
force-pushed
the
pr/dspark-pp-pd
branch
from
September 16, 2026 01:01
16523b6 to
23913fb
Compare
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
1 participant
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 prefill on a pipeline-parallel producer with KV transfer (PD-disaggregated serving). The drafter runs in prefill-only mode on the producer: it materializes the draft's context KV and lets the KV connector ship named per-layer regions, instead of proposing draft tokens inline.
SpeculativeConfig.target_kv_transfer_config+is_dspark_prefill_only(); NIXL is rejected for this mode (it cannot describe packed per-layer KV regions; Mooncake works).use_eagle(and eagle block-drop) for the prefill-only producer so it never does the EAGLE-style read-ahead.DFlashSpeculator.materialize_context_kv(): precompute and store draft context KV without drafting; the model runner calls it in place ofpropose()on the producer and skips draft-side rejection sampling, drafter graph capture, and EPLB registration there.context_kv_only) — avoiding transient MoE expert allocation — and loads just the context-KV weights.Depends on #56956 (aggregated-serving PP support; only the last commit here is new). Kimi-K3 DSpark additionally needs #51065 (TritonMLA causal multi-token fix; verified on SM120).
Split from #53577 (reviewer request).
Tests
tests/config/test_dspark_prefill_only.py(13 tests: mode detection, connector guard, last-stage drafter helper).pytest tests/config/test_dspark_prefill_only.py -q→ 13 passed; the aggregated-serving suites from [DSpark] Support pipeline-parallel targets in aggregated serving #56956 stay green on the combined tree.Model evaluation
Validated end-to-end as part of #53577: DeepSeek-V4-Flash DSpark PD with PP producer (Mooncake), GSM8K strict-match 0.9568; Kimi-K3 DSpark PD smoke. Will re-run on the rebased stack once #56956 lands.
AI assistance was used in preparing this PR.