[DSpark] Support pipeline-parallel targets in disaggregated and aggregated serving (+ padded graph batch safety) - #53577
lucifer1004 wants to merge 25 commits into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
mergify
Bot
added
deepseek
|
This pull request has merge conflicts that must be resolved before it can be |
Keep the None guard around get_mtp_target_hidden_states() alongside the new GPUWatermarkSampler preparation from upstream. Co-authored-by: Kimi Code <noreply@moonshot.cn> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Adapt to the DFlash JIT-warmup migration (vllm-project#56323): prepare_dflash_inputs now returns a DispatchSpec, with our out_is_padding_ptr output added to the launch dict and warmup inputs. Sampler init moves under the jit_warmup_registry block, keeping the dspark_prefill_only gate on the rejection sampler. Co-authored-by: Kimi Code <noreply@moonshot.cn> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Follow the revert of the Triton JIT warmup migration (vllm-project#56654): prepare_dflash_inputs launches the kernel directly again, with our out_is_padding_ptr argument kept; sampler init dedents back out of the jit_warmup_registry block, keeping the dspark_prefill_only gate. Co-authored-by: Kimi Code <noreply@moonshot.cn> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Take upstream's use_confidence_head attribute declaration in DSparkSpeculator.__init__ alongside our confidence/adaptive-verification and fixed-graph-batch setup. Co-authored-by: Kimi Code <noreply@moonshot.cn> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
|
/ci run |
|
❌ This PR is 4 commits behind upstream |
|
/ci run |
|
✅ Triggered Buildkite CI #88978 for commit |
- warmup: short-circuit on is_last_pp_rank before reading pp_handler so SimpleNamespace runner stubs without the attribute keep working. - test_gpu_model_runner_v2: the QSA runner stub bypasses __init__, so set dspark_prefill_only explicitly. Co-authored-by: Kimi Code <noreply@moonshot.cn> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
GirasoleY
left a comment
GirasoleY
left a comment
There was a problem hiding this comment.
The changes feel a bit wider than expected for supporting PP + Dspark. Could we maybe break it down to a few isolated PRs for easy review? High level it seems like we can have
- DSpark + PP aggregated serving
- DSpark + PP prefill + KV transfer
- K3 cache fix
This comment was marked as duplicate.
This comment was marked as duplicate.
Sorry, something went wrong.
|
Per @GirasoleY's review, this PR is now split into three isolated pieces:
Each carries its own tests and e2e evidence (GSM8K 0.9545 on the IFB PP arm). Closing this one in favor of the series; thanks for the review! |
4 participants
[DSpark] Support pipeline-parallel prefill in disaggregated serving (+ padded graph batch safety)
Two stacked commits:
feat(dspark): support PP prefill in disaggregated servingfix(dspark): make padded graph batches safePurpose
PD-disaggregated serving with DSpark previously forced the prefill (producer)
side to run without speculation: the DFlash/DSpark drafter does not support
pipeline parallelism, so a PP>1 producer could not warm the draft context, and
the decode-side drafter started cold (measured per-token acceptance dropped
from ~45% to ~35% in PD vs aggregated serving).
This PR materializes the DSpark draft's context KV on PP prefill workers
without running the decode-only draft path:
SpeculativeConfig.is_dspark_prefill_only()detects the producer role(dspark + PP>1 + kv_producer, fed by the new
target_kv_transfer_config).context_kv_onlymode: onlythe context-KV projection (main_proj/main_norm + per-layer fused wq_a/wkv,
kv_norm, rotary, SWA cache layer) is constructed and loaded; embedding,
lm_head, Markov/confidence heads are skipped; drafting tokens raises.
DFlashSpeculator.propose()is split somaterialize_context_kv()runs thecontext-KV projection without token generation; the prefill model runner
calls only this path and skips the rejection sampler, draft CUDA graphs,
EPLB registration, and adaptive verification. The draft parallel config is
forced to PP=1, and all auxiliary hidden states must live on the last
pipeline stage (explicitly validated).
connectors that register and align named per-layer KV regions (verified with
MooncakeConnector, which aligns producer/consumer regions by layer name and
pulls from every producer PP stage). The NIXL connector does not yet describe
packed KV regions per producer stage, so it raises a loud
NotImplementedErrorfor this configuration instead of silentlymis-transferring. NIXL packed-region support is left to the [KVConnector][NIXL] Support attention-HMA layouts in pipeline-parallel push prefill #50494/[KVConnector][NIXL] Support packed MLA KV layouts in pipeline-parallel push prefill #50499
line of work.
The second commit makes padded CUDA-graph batches safe by tracking and masking
graph-padding rows through DFlash input preparation and DSV4 routing,
preserving compact request-slot mappings, and allowing an optional fixed DSpark
graph batch size.
Relationship to existing work
speculators. This PR targets DSpark (non-causal SWA draft caches) and does
not conflict with it.
this PR does not change the NIXL connector; if/when they land, the
NotImplementedErrorguard can be revisited.Test plan
pytest tests/config/test_dspark_prefill_only.py tests/models/test_deepseek_v4_mega_moe.py tests/v1/spec_decode/test_dflash_prepare_inputs.py tests/kernels/moe/test_topk_softplus_sqrt.py— all pass.pre-commit run --from-ref upstream/main --to-ref HEAD— all hooks pass.PP2xTP2 prefill (node A) + TP4 decode (node B) over MooncakeConnector,
DSpark draft-context materialization enabled: both roles reach readiness,
CUDA-graph capture succeeds, and an 8K/1K random benchmark across six
concurrency levels completes with zero failed requests and active draft
acceptance (36% accepted/drafted tokens this run). Note: current
mainadditionally needs [Bugfix][SM120] DSv4: pass contiguous C128A decode topk indices on SM120 #53574 (one-line C128A contiguity fix) to boot on SM120
at all; the E2E above was run on this PR branch stacked with that fix.
aggregated parity at 41-48% under an 18K/3K workload) was taken on our full
integration tree, which carries the same two commits plus unrelated
SM120-enablement work; the mechanism under test is identical.
Notes
AI assistance (Kimi Code / OpenAI Codex) was used in preparing this change;
every line was reviewed and tested by the submitting human.
Update 2026-09-03: extended to aggregated (IFB) serving + PP
The branch now carries a fresh merge with
main(ee17d0d8) plus three newcommits:
feat(dspark): support pipeline-parallel targets in aggregated servingfix(pp): make warmup deadlock-free under pipeline parallelismfix(pp): complete the sampled-token broadcast contract for spec decodingWhat changed:
SpeculativeConfig.use_dspark_last_stage_drafter()replaces thekv_producer-only check: a DSpark drafter under a PP>1 target always runs
with a draft-local PP=1 config, in aggregated serving as well as in PD.
The PD prefill-only paths are unchanged (all new logic is gated on
num_speculative_steps > 0, which is 0 for that mode).spinning, unmatched NCCL kernel blocking the host-side CUDA module load of
a first-time triton compile:
discarded there anyway), and the deferred post-update kernel is
explicitly pre-warmed on non-last ranks;
[N, 1]payloads arepadded to
[N, max_sample_len], and the(2, N)counts buffer is paddedso both views stay 16-byte aligned (triton specializes on pointer
alignment; a misaligned view compiled a second kernel variant at serving
time).
propose()andadopted into their
req_states.draft_tokens. Previously those stagesembedded token id 0 for every draft slot, which would have produced garbage
verification logits in aggregated serving (PD prefill-only never needs
drafts on the producer).
Validation (8x RTX 6000 Pro SM120, DeepSeek-V4-Flash-0731 NVFP4, aggregated
IFB PP2xTP2 + DSpark, b12x MoE backend):
pytest tests/v1/worker/test_pp_utils.py tests/config/test_dspark_prefill_only.py tests/models/test_deepseek_v4_mega_moe.py tests/v1/spec_decode/test_dflash_prepare_inputs.py— 45 pass;pre-commit run— all hooks pass.target-only tp4 = 0.9469 — within noise, distribution preserved.
zero failed requests): PP costs ~22-35% at C1-C16, reaches parity at C32,
and leads at C64 (508 vs 416 tok/s). As expected, PP in aggregated serving
is a capacity play, not a decode-throughput one.
AI assistance (Kimi Code) was used for this update as well; all changes were
reviewed and tested by the submitting human.
Review follow-ups (2026-09-04)
Addressing review findings with
0f851e2c5d:KVTransferConfig.has_connector(), catching NixlConnector nested in aMultiConnector.
use_eagle_block_dropnow follows the narroweduse_eagle, so a DSpark prefill-only producer keeps its trailingprefix-cache block.
DFlashSpeculator.proposerebuilds the reused DP-sync token count whenthe fixed DSpark graph batch re-pins the dispatch shape (dp_size > 1
assertion path; dp_size = 1 behavior unchanged and re-covered by the IFB
PP2xTP2 smoke run).
Update 2026-09-04 (later): Kimi-K3 cross-stage aux taps
New commit 6.
feat(dspark): support Kimi-K3 targets with cross-stage aux hidden tapsAt the previous head, K3 + PP + DSpark crashed at startup: the aux-tap check
in
load_dspark_modelreaddspark_target_layer_idsdirectly, whichspeculators-format draft configs (K3) do not define, and it required every
tap on the last stage while K3's five taps ([24, 48, 72, 88, 92]) span the
whole model. The check now resolves taps via the shared
get_eagle3_aux_layers_from_configfallback chain (value-identical forDeepSeek-V4's
dspark_target_layer_ids), and targets may opt intocross-stage transport with
supports_pp_aux_hidden_state_transport;KimiLinearModelpacks earlier stages' taps intoIntermediateTensorsandunpacks them downstream (stage-entry capture stays first-stage-only so
boundary taps are not captured twice).
Validation (8x RTX 6000 Pro SM120):
pytest tests/config/test_dspark_prefill_only.py tests/models/kimi_k3/test_eagle3.py— 21 pass;prek runon changed files — all hooks pass.tests/evals/gsm8kshape): DSpark PP2xTP4 = 0.2570 vs TP8 = 0.2684 (McNemar z = 1.08, ns); target-only PP2xTP4 = 0.3336 vs TP8 = 0.3351 (z = 0.13, ns). Draft acceptance length 5.3-6.3 (gate 4.8). PP has no measurable quality impact on K3 DSpark.Two honest notes: (1) on this pruned K3 checkpoint DSpark scores ~7pp below
target-only at temp 0 under both TP8 and PP2xTP4 — speculator-side, not PP;
flagged separately. (2) K3 PD with PP (KDA state transfer) is not covered by
this change. AI assistance (Kimi Code); all changes reviewed and tested by
the submitting human.
Update 2026-09-06: rebased on main (52358e6) + four fixes
Upstream #50514 landed a generic EAGLE3-over-PP aux-hidden-state relay, so the
merge adopts that mechanism wholesale and drops this branch's earlier custom
transport. Both DSv4 and K3 opt in via
supports_aux_hidden_states_over_pp.Four fixes on top of the merge:
fix(dspark): let PP drafters load their own embedding when the target's is stranded(4728759) — under PP the target's embedding table lives on thefirst stage, so a drafter on a later stage cannot alias it. K3/DSv4 DSpark
drafters now load the checkpoint's own embed copy
(
loads_own_embed_under_pp) instead; without this, DSv4 DSpark + PP raisesat load time on stock main.
fix(dspark): address review findings on topk uint32 dispatch and context-KV construction(1bb49fb) — uint32+padding is excluded fromthe dsv4_topk fast path (its -1 sentinel requires signed indices);
context-KV-only draft layers build the attention submodule directly instead
of a whole decoder layer, removing a transient MoE allocation that could
OOM on memory-tight GPUs.
fix(pp): do not double-post draft broadcasts when a speculator ran(5b572da) — the merge kept two
broadcast_draftscall sites on the lastPP rank (one in the propose() path, one next to
set_draft_tokens()from[Core][MRV2] Support eagle3 spec decode with pipeline parallel #50514). With a speculator present both fired per step while earlier stages
post one recv: the extra send shifted the pp_broadcast FIFO, the next
step's sampled-token recv paired with a leftover draft broadcast of a
different size, and the mismatched NCCL collective spun on-device, blocking
any later context-level CUDA call and deadlocking spec-decode + PP
deterministically. Reproduced on DSv4 pp2tp2 and K3 pp2tp4 (non-spec PP
unaffected); fixed by broadcasting from the handler path only when there
is no speculator (e.g. diffusion-style drafts).
fix(k3): keep the DSpark draft's marker from flagging the target's KV group(0a2f859) — the K3-native DSpark draft'snon_causal_multi_token_decode MLA layers merged with the target's
identical-geometry MLA layers (OR-merge since [Bugfix][MLA] Restore DSpark cache-group capability under optimized Python #55234), flagging the
target group and misrouting its short prefills and causal verification
blocks into TritonMLA's single-row decode path: NaN logits from the
first token, acceptance 0, and an IMA under graph capture. The draft's
spec now carries a distinct model_version so the groups stay separate.
Verified on Kimi-K3-pruned75 + Inferact/Kimi-K3-DSpark (SM120): tp8
eager and pp2tp4 graphs both sane, draft per-position acceptance 0.85.
Validation on this branch (RTX PRO 6000 Blackwell, SM120):
pytest tests/v1/worker/test_pp_utils.py tests/v1/worker/test_spec_decode_embed_sharing_pp.py tests/v1/worker/test_mixed_warmup_gate.py tests/models/kimi_k3/test_eagle3.py tests/models/kimi_k3/test_dspark_mla.py— all pass (incl. new mapper andembed-under-PP tests);
tests/kernels/moe/test_topk_softplus_sqrt.pyincl.the new uint32+padding fallback case passes on GPU;
pre-commit runclean.exact_match (strict) = 0.9568.
(1319 questions): accuracy 0.2798, invalid 0.0015 — in line with the
pre-merge 0.2570 reference on this weak pruned checkpoint.
target's triton_mla decode kernel hits an illegal memory access on SM120
during profiling, also reproducible with TP8 alone, independent of PP and
of this PR's changes; tracked separately.
AI assistance (Kimi Code) was used for this update as well; every line was
reviewed and the validation above run by the submitting human's setup.