Repository navigation
[Perf] [NVIDIA] dspark: Stage draft DP-MoE token counts through persistent buffers - #47677
weireweire wants to merge 3 commits into
Conversation
Root cause: under DP-MoE sync, DraftBlockProposer._fill_dp_moe_sync_metadata builds the draft's global_num_tokens / for_logprob tensors from fresh pinned torch.tensor calls each step, plus two fresh num_token_non_padded H2Ds (one at ForwardBatch build, one in the fill). DEP16 c2884 nsys: P.dp_meta 0.23-0.26 ms and P.fwdbatch 0.21 ms on the post-agreement host tail. Fix: SGLANG_OPT_DSPARK_PERSISTENT_DP_META (defaults to SGLANG_ENABLE_NV_MEGA_RUNTIME_OPTIMIZATION) writes both count rows into a two-slot pinned [2, 2, dp] host buffer via numpy and issues one H2D into a matching device buffer; global_num_token_non_padded reuses the already-filled _num_token_non_padded. Each slot records an event after its H2D and synchronizes it before the slot is rewritten, so reuse does not depend on the scheduler's overlap cadence (pipeline microbatches or a second fill per step cannot overwrite a pending copy). Validation (GB300, DSV4.1-Flash, on a 2026-10-30 nightly): - DEP16 nsys: P.dp_meta median 256 -> 125 us, P.fwdbatch 214 -> 112 us. - Combined with the other two host changes (same runs; numbers are for all three together): DEP32 c5768 decode step p50 15.2 -> 13.9 ms; steady-window per-GPU throughput (mean of minutes 8-20 of the aiperf run, x0.96, both arms) 660.6k vs 664.4k tok/s/GPU at p90 interactivity 219.8 vs 200.8 tok/s/user (vLLM-frontier gap -26.3% -> -20.8%); DEP16 c2884 720.1k vs 719.0k tok/s/GPU (same normalisation; no regression); GSM8K 1P1D with real acceptance 96.82% vs 96.89%, accept 3.50; TP4 c1 agg T=1 GSM8K 97.33% vs 97.33%, long-context needle retrieval at 20K-160K tokens 48/48 vs 48/48, decode step neutral (3.514 vs 3.501 ms).
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
SGLANG_OPT_DSPARK_PERSISTENT_DP_META now defaults to on instead of following SGLANG_ENABLE_NV_MEGA_RUNTIME_OPTIMIZATION, which DEP decode deployments do not set. It stays as a kill switch because the path keeps staging buffers and CUDA events across steps. test_dspark_dp_tier builds the proposer without __init__; give it the persistent-meta fields so it runs the default path.
Root cause: SGLANG_OPT_DSPARK_PERSISTENT_DP_META kept the per-step fresh pinned-tensor path alive as a kill switch. The persistent path writes the same values: the device slots are rewritten only in stream order after their readers, the pinned slots only after their recorded H2D has completed, and the aliased global_num_token_non_padded holds the count the fill already writes, with every later writer rebinding it or rewriting the same value. The switch only doubled the paths to test. Fix: remove the env var, the fresh-tensor branch and _make_num_token_non_padded; the draft always stages through the persistent buffers. _run_forward's device argument, used only by that helper, goes too. The DP-coordination unit test returns to a single path and keeps the staged device-count assertions. Validation: all GPU results on the PR already ran the persistent path (DEP16 / DEP32 with the switch on, TP4 agg with the mega flag on). pre-commit passes.
| # Rewrite a pinned slot only after its previous H2D has run. | ||
| copied = self._dp_meta_copied[slot] | ||
| if copied is not None: | ||
| copied.synchronize() |
There was a problem hiding this comment.
Why does here need sync? Could we remove this?
There was a problem hiding this comment.
This is a double buffer (note self._dp_meta_slot ^= 1), so it waits for the copy from two iterations ago. That copy has typically finished already, so it doesn't actually wait, but we still need the barrier to guarantee it.
There was a problem hiding this comment.
note here we are reusing the two slots, so we need to make sure its data is transferred already and then rewrite again
nvpohanh
left a comment
nvpohanh
left a comment
There was a problem hiding this comment.
[by Claude Code] This review is done by agents. Please feel free to disagree if they do not make sense.
| forward.global_num_tokens_for_logprob_cpu, [0, 3 * width, 0] | ||
| ) | ||
| self.assertEqual( | ||
| forward.global_num_tokens_gpu.tolist(), [0, 3 * width, 0] |
There was a problem hiding this comment.
[by Claude Code] Severity: style | Confidence: High
Issue
These new GPU-count assertions cannot catch a row mix-up in _stage_dp_meta. Here both counts are [0, 3 * width, 0], so the two rows hold the same values. Each test also builds a fresh proposer and fills once, so the two-slot reuse never runs.
Example
I changed _stage_dp_meta to return dev[0], dev[0]. All 13 tests in the three test files listed in the PR still pass.
Fix
In test_preserves_unscaled_request_counts_for_cuda_graph_admission, the counts differ ([6, 18, 0, 12] vs [1, 3, 0, 2]), so assert both GPU tensors there. Then fill three times with different counts on one proposer. Check that each call's tensors hold its own counts and that the third call reuses the first call's storage.
|
@Qiaolin-Yu could you review again? Thanks! |
Motivation
Under DP-MoE sync,
DraftBlockProposer._fill_dp_moe_sync_metadatabuilds two fresh pinnedtorch.tensors each decode step and copies them H2D: the draft'sglobal_num_tokens_gpuandglobal_num_tokens_for_logprob_gpu. The draft ForwardBatch also gets a freshglobal_num_token_non_paddedH2D twice, once when it is built and once in the fill.On DSV4.1-Flash GB300 disagg DEP16, this host work costs
P.dp_meta0.23-0.26 ms andP.fwdbatch0.21 ms. It sits on the host path between the DP agreement and the draft graph launch.Related PRs
One of four DSV4.1 attn-DP decode host-path changes that cut the host work between the DP agreement and the draft graph launch. They are independent in code and were measured together: #47676 (padded metadata glue), #47677 (persistent DP-MoE counts), #47678 (bulk decode-prep fill) and #47685 (KV page top-up, measured on top of the other three).
Modifications
[2, 2, dp]host buffer, written through its numpy view, and one H2D copies them into a matching device buffer.global_num_tokens_gpuandglobal_num_tokens_for_logprob_gpuare views into the current device slot.num_token_non_padded(EP > 1, DP gather, a custom A2A backend, DP attention with TP MoE) passes the persistent_num_token_non_paddedasglobal_num_token_non_padded, instead of a fresh H2D of the same value. The fill already writes that count into it. Later writers rebind the field, except the idle dummy-extendfill_inprepare_mlp_sync_batch, which writes the same value (0) that is already there._make_num_token_non_paddedare removed. There is no env switch: the new path writes the same values and has no configuration in which the old one is better.Why reusing the slots is safe:
prepare_mlp_sync_batchstill writesglobal_num_tokens_gpuin place, now into the current slot.All runs: GB300, DSV4.1-Flash, on a 2026-10-30 nightly.
Accuracy Tests
Speed Tests and Profiling
P.dp_meta, median over ranksP.fwdbatch, median over ranksCombined DEP32 / DEP16 / TP4 numbers: see #47676. They come from the same runs and are combined across the three changes.
The teardown note from that PR applies here as well: a MegaMoE NVLink barrier timeout after the benchmark aborts also happens without these changes. Serving-phase results are unaffected.
Checklist
test/registered/spec/dspark/test_dspark_dp_tier.py,test/registered/unit/spec/test_dspark_dp_spec_prefill_coordination.py(now also checks the staged device counts),test/registered/unit/spec/test_dspark_draft_block.pyCI States
Latest PR Test (Base): ❌ Run #41413829634
Latest PR Test (Extra): ❌ Run #41413829668
Latest PR Test (AMD ROCm 10): ❌ Run #41413829660
Tip
Base CI only covers the core cases, and its workload may not exercise your change. Use
/rerun-testto run the tests your change touches, or copy this prompt for an agent: