Repository navigation
[Perf] [NVIDIA] dflash: Fill decode-prep length buffers in bulk - #47678
Conversation
Root cause: DFlashDraftInputV2.prepare_for_decode writes the committed, current and next KV lengths with one CPU-tensor __setitem__ per request and field (3 x bs element writes each step). DEP16 c2884 nsys: prepare_for_decode 1.0-1.4 ms every step, on the pre-agreement host path of the last DP rank to arrive. Fix: SGLANG_OPT_DFLASH_BULK_PREPARE_DECODE (defaults to SGLANG_ENABLE_NV_MEGA_RUNTIME_OPTIMIZATION) gathers the committed lengths and top-k values into lists and writes each buffer once through its numpy view; sums, max top-k and the uniform top-k value are computed from the lists with the same results as the loop. Validation (GB300, DSV4.1-Flash, on a 2026-10-30 nightly): - DEP16 nsys: prepare_for_decode median 1162 -> 774 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}
The bulk fill is an exact CPU-side rewrite of the per-element loop (same buffers, sums, max and uniform top-k), so it needs no opt-in switch. Drop SGLANG_OPT_DFLASH_BULK_PREPARE_DECODE and the loop.
|
/rerun-test test/registered/core/test_basic_sanity_dflash.py test/registered/spec/dflash/test_dflash.py test/registered/unit/spec/test_decode_bookkeeping_ownership.py test/registered/spec/test_spec_mixed_chunk.py test/registered/e2e/speculative/test_dflash_domino.py test/registered/e2e/speculative/test_dflash_dp_attention.py test/registered/spec/test_gemma4_dflash_31b_extra.py test/registered/e2e/models/test_nvidia_nemotron_3_nano.py test/registered/spec/dflash/test_muse_glimmer_dflash_assistant_gsm8k.py test/registered/e2e/disaggregation/test_glm53_cp_bcg_dflash.py test/registered/e2e/models/test_muse_spark_12_b300_pd.py |
|
🚀 🚀 🚀 🚀 🚀 🚀 🚀 |
|
Tiny change. DFlash tests passes. Merging now |
Motivation
DFlashDraftInputV2.prepare_for_decodewrites the committed, current and next KV lengths with one CPU-tensor__setitem__per request and field, i.e. 3 x bs element writes every decode step. On DSV4.1-Flash GB300 disagg DEP16 (40 running per rank),prepare_for_decodetakes 1.0-1.4 ms per step, and these element writes are about 0.4 ms of it. That time sits on the pre-agreement host path of the last DP rank to arrive.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
The per-element loop is replaced by an exact rewrite, so there is no switch:
All runs: GB300, DSV4.1-Flash, on a 2026-10-30 nightly.
Accuracy Tests
Speed Tests and Profiling
prepare_for_decodeNVTX range, median over ranksCombined DEP32 / DEP16 / TP4 numbers: see #47676. They come from the same runs, and those numbers are combined across the three changes.
The teardown note from that PR applies here as well: a MegaMoE NVLink barrier timeout after the benchmark abort also occurs without these changes; serving-phase results are unaffected.
Checklist
CI States
Latest PR Test (Base): ❌ Run #41821589951
Latest PR Test (Extra): ❌ Run #41821589332
Latest PR Test (AMD ROCm 10): ❌ Run #41821589526
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: