[Spec] DFlash: remove per-step host syncs so the CPU runs a full step ahead (spec-v2 overlap) by thanhhao98 · Pull Request #31468 · sgl-project/sglang

thanhhao98 · 2026-07-16T14:02:49Z

Motivation

Under spec-v2 overlap scheduling, DFlash stalls the host every decode step: nsys shows a large cudaStreamSynchronize between the draft and verify cuda graphs, so cudaGraphLaunch for step N+1 is only issued after step N has finished (launch lead ≈ −50 µs). EAGLE3 spec-v2 runs a full step ahead; DFlash should too. This PR removes every per-step host↔device sync on that path. Result: the CPU queues step N+1 a full step early (launch lead +19 ms, run-ahead fraction 0.88 → 0.995), worth +12–13% tok/s/user at low concurrency and +3–5% at high, with GSM8K and accept-length unchanged.

Modifications

Each item removes one per-step host block:

Compact draft-cache rebuild → one triton kernel (rebuild_compact_draft_req_to_token). The old rebuild synced three times per step (lengths.max().item(), a boolean-mask gather's implicit nonzero D2H, and a blocking .to("cpu")), each waiting on the previous verify step. The kernel writes the suffix window + verify-block slots directly into the draft req_to_token row with no data-dependent shapes. The host seq_lens_cpu planning bound becomes min(len, window+page) — deliberately an envelope, not a mirror of the exact page-align math, which is a non-monotonic sawtooth and would under-shoot when fed the reserved over-estimate.
HybridAttnBackend: OR-delegate needs_cpu_seq_lens to the sub-backends. The missing attribute defaulted to True, forcing FutureMap.resolve_seq_lens_cpu's pinned D2H + stream.synchronize() before every step's launches — zero run-ahead on hybrid arms even when both sub-backends opt out.
Fold the TP>1 draft greedy sampler into the draft cuda graph (tp=1 fold already existed): per-rank shard (max, global id) → two in-graph NCCL all-gathers → first-max select. Bit-exact vs a full-vocab argmax including shard-boundary ties. SGLANG_DFLASH_EAGER_DRAFT_SAMPLER=1 restores the eager path.

Correctness

Unit tests in this PR (test/registered/unit/spec/test_dflash_overlap_hostsync.py): kernel bit-exact vs the legacy path (incl. no writes past the verify block), sampler == full-vocab argmax incl. a shard-boundary tie, the sawtooth-envelope bound with its counterexample, hybrid flag delegation, filter_batch host-list equivalence.
8×B200 probes at production dims (tp=8, vocab 163840): fused rebuild and folded sampler bit-exact, cross-rank identical.
GSM8K 5-shot: 0.965 == 0.965 (main vs main+PR); accept-length 2.47 / 2.51.

Benchmarks (Kimi-K2.6-NVFP4 + DFlash draft, 8×B200 TP=8, aiperf 50K-prefix)

Same-base A/B, identical stack, only this PR's commits differ:

cc	main `40a3bd76`	main + PR	Δ
1	256.9	288.0	+12.1%
2	215.0	243.6	+13.3%
4	169.7	174.4	+2.8%
8	125.2	131.8	+5.3%
16	80.5	84.5	+5.0%

nsys (cc16, same pair): launch lead median −51 µs → +19.1 ms; run-ahead fraction 0.883 → 0.995; per-thread cudaStreamSynchronize 336×/3816 ms → 9×/11 ms. The only remaining per-step host sync is the intended copy_done in result processing (identical to EAGLE3 spec-v2). Same shape on the tokenspeed_mla arm (K2.7: +11.5% cc1, lead −57 µs → +17.1 ms); details in the PR comments.

Checklist

Format your code according to the Code Formatting with Pre-Commit.
Add unit tests as outlined in the Running Unit Tests (test/registered/unit/spec/test_dflash_overlap_hostsync.py, 7 tests, validated on B200).
Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.

🤖 Generated with Claude Code

CI States

Latest PR Test (Base): ⏳ Run #29632179798
Latest PR Test (Extra): ❌ Run #29632179752

gemini-code-assist · 2026-07-16T14:03:04Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

…to the draft graph Two per-step host-stall fixes for spec-v2 run-ahead (the CPU must launch step N+1's cuda graphs while step N still executes): 1) Compact draft req_to_token rebuild: replace the masked segment gather (+ 2x assign_req_to_token_pool) with one fused fixed-grid triton kernel. The old path blocked the host twice per decode step: lengths.max().item() and the boolean-mask packed gather (data-dependent shape => implicit nonzero D2H), serializing the CPU behind the previous verify step. The draft-forward host planning bound (seq_lens_cpu) is now computed on CPU from the scheduler-maintained host lens via the monotonic envelope min(len, window+page) instead of a blocking D2H. The exact page-align mapping is a non-monotonic sawtooth, so the envelope (not a mirror of the arithmetic) is required to stay an upper bound when fed the reserved overlap over-estimate. Legacy path kept only for non-GPU-triton platforms. 2) Extend _DflashDraftSampler to the vocab-parallel (TP>1) head and lift the tp>1 eager bail: per-rank shard max/argmax -> global id, two NCCL all-gathers (captured into the draft graph like the draft model's own collectives), first-max select. Tie resolution is bit-exact vs a full-vocab argmax (contiguous ascending shards + first-max argmax). A/B escape hatch: SGLANG_DFLASH_EAGER_DRAFT_SAMPLER=1 forces eager. Removes the eager between-graphs sampler cluster (matmul + reduce + 2x all_gather + gather) from the host launch path.

decide_needs_cpu_seq_lens ORs the flag across attn backends; the hybrid prefill/decode wrapper had no attribute so the getattr default (True) forced FutureMap.resolve_seq_lens_cpu into its per-step pinned D2H + fwd_prepare_d2h_stream.synchronize() even when both sub-backends opted out (trtllm_mla prefill + cutedsl_mla decode). That host sync sits before any of step N's launches and blocks on step N-1's verify completing on GPU, forcing zero CPU run-ahead on the hybrid arm.

DFlashDraftInputV2.filter_batch sliced reserved_seq_lens_cpu with new_indices.cpu(), a blocking D2H of the GPU index tensor on the scheduler thread every time batch composition changed. ScheduleBatch already holds the host-side keep list; pass it through and index the CPU tensor directly. EAGLE/ngram accept the kwarg unchanged.

…prep churn - init_forward_metadata: never derive max_seq from the GPU tensor (the .max().item() drains the whole stream backlog under spec-v2 overlap); fall back to the static context bound when no host mirror exists. - DCP eager paths: set dcp_local_max_seq_len / dcp_prefix_local_max host-side (ceil(global/world) bound covers get_dcp_lens on every rank), so forward_decode/_forward_verify_dcp never hit the per-layer local .max().item() device sync (61 syncs/step on Kimi) off-graph. - _apply_cuda_graph_metadata target-verify: fuse the seq_lens+T add straight into the captured int32 seq_lens_k buffer (out= same-kind downcast) and feed that buffer to the block-table kernel: one kernel, no per-step temp alloc. - WAR fastpath: reuse a 2-deep event ring instead of allocating a fresh cuda Event per graph replay (DFlash records twice per step).

…comments further, add unit tests - SGLANG_DFLASH_EAGER_DRAFT_SAMPLER moves to the typed envs registry. - Second comment pass: keep only the load-bearing invariants (sawtooth envelope, tie-safety, needs_cpu_seq_lens OR-default trap, re-record safety, legacy-path sync rationale). - test_dflash_overlap_hostsync.py: fused rebuild bit-exact vs legacy path (+ no write past the verify block), vocab-parallel sampler == full-vocab argmax incl. shard-boundary ties (simulated TP group, single GPU), host-bound envelope >= exact incl. the sawtooth counterexample, HybridAttnBackend needs_cpu_seq_lens delegation, filter_batch host keep-list equivalence.

thanhhao98 · 2026-07-17T08:31:16Z

nsys before/after — the overlap improvement, Kimi-K2.6-NVFP4 cuteDSL, cc16

BEFORE (main 40a3bd7):

AFTER (this PR):

hnyls2002 · 2026-07-18T04:59:03Z

/rerun-test registered/unit/spec/test_dflash_overlap_hostsync.py registered/spec/dflash/test_dflash.py registered/core/test_basic_sanity_dflash.py registered/spec/eagle/test_spec_eagle.py registered/spec/test_spec_ngram.py registered/attention/test_hybrid_attn_backend.py registered/attention/unittests/mla/test_trtllm_mla.py

hnyls2002 · 2026-07-18T04:59:21Z

/tag-and-rerun-ci

github-actions · 2026-07-18T04:59:28Z

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit a5c0b94034.

Re-run was not dispatched. What to do:

Rebase your branch onto the latest main and push again
Follow issue #21065 for context
CI-fix PRs may request the bypass-maintenance label to skip this check

github-actions · 2026-07-18T05:00:11Z

Results for /rerun-test registered/unit/spec/test_dflash_overlap_hostsync.py registered/spec/dflash/test_dflash.py registered/core/test_basic_sanity_dflash.py registered/spec/eagle/test_spec_eagle.py registered/spec/test_spec_ngram.py registered/attention/test_hybrid_attn_backend.py registered/attention/unittests/mla/test_trtllm_mla.py:

🚀 1-gpu-5090 (4 tests): ✅ View workflow run

cd test/ && python3 registered/unit/spec/test_dflash_overlap_hostsync.py
cd test/ && python3 registered/spec/dflash/test_dflash.py
cd test/ && python3 registered/core/test_basic_sanity_dflash.py
cd test/ && python3 registered/spec/eagle/test_spec_eagle.py

🚀 1-gpu-h100 (3 tests): ✅ View workflow run

cd test/ && python3 registered/spec/test_spec_ngram.py
cd test/ && python3 registered/attention/test_hybrid_attn_backend.py
cd test/ && python3 registered/attention/unittests/mla/test_trtllm_mla.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/attention/unittests/mla/test_trtllm_mla.py

hnyls2002 · 2026-07-18T05:00:32Z

@thanhhao98 I reverted the read_done event ring on your branch (commit 4a488a5) — it is orthogonal to this PR's goal: it removes a per-step allocation, not a host sync, and the same per-replay Event() pattern exists elsewhere (e.g. copy_done in scheduler result processing), so it should be done uniformly rather than for one site.

Could you resubmit the event-ring change as a separate PR covering all the per-step Event() sites? Please also drop the event-ring bullet from this PR's description (item 4) since it is no longer part of the diff.

hnyls2002 · 2026-07-18T05:04:58Z

/rerun-failed-test

thanhhao98 · 2026-07-19T06:29:10Z

Thanks @hnyls2002 @kpham-sgl for pushing this. I removed event-ring in PR description.
I will create separate PR for it.

… ahead (spec-v2 overlap) (sgl-project#31468) Co-authored-by: Hao Phan <htphan@nvidia.com>

github-actions Bot added the blackwell SM100/SM120 label Jul 16, 2026

Hao Phan added 4 commits July 17, 2026 10:33

thanhhao98 force-pushed the htphan/dflash-overlap-hostsync branch from 677d454 to 361a7e8 Compare July 17, 2026 03:33

Hao Phan added 2 commits July 17, 2026 14:15

[Spec] DFlash overlap: trim comments to the load-bearing invariants

bfd7263

thanhhao98 marked this pull request as ready for review July 17, 2026 08:33

thanhhao98 requested review from Fridge003, HaiShaw, Qiaolin-Yu, Ying1123, hebiao064, hnyls2002, ispobock, kpham-sgl, merrymercy and xiezhq-hermann as code owners July 17, 2026 08:33

kpham-sgl self-assigned this Jul 17, 2026

hnyls2002 added 4 commits July 17, 2026 16:51

drop getattr in hybrid needs_cpu_seq_lens

1ff2ff9

trim test module docstring; fix eager sampler env comment

fbf13a7

revert read_done event ring (out of scope)

4a488a5

use CustomTestCase in dflash hostsync test

ec753e7

hnyls2002 approved these changes Jul 18, 2026

View reviewed changes

github-actions Bot added the run-ci label Jul 18, 2026

Merge branch 'main' into htphan/dflash-overlap-hostsync

ec3bb5e

fix fake backend attr in attention setup test

c13f728

kpham-sgl approved these changes Jul 18, 2026

View reviewed changes

hnyls2002 merged commit 72c4ed1 into sgl-project:main Jul 18, 2026
97 of 121 checks passed

amd-bot mentioned this pull request Jul 18, 2026

[CI Monitor] Daily Report - 2026-07-18 bingxche/sglang-ci-bot#139

Open

hnyls2002 mentioned this pull request Jul 18, 2026

[Spec] Extract DFlash compact draft-cache rebuild helpers #31677

Merged

amd-bot mentioned this pull request Jul 19, 2026

[CI Monitor] Daily Report - 2026-07-19 bingxche/sglang-ci-bot#140

Open

thanhhao98 mentioned this pull request Jul 19, 2026

[Scheduler] Reuse per-step cuda events uniformly (WAR read_done + copy_done) #31698

Open

5 tasks

thanhhao98 deleted the htphan/dflash-overlap-hostsync branch July 19, 2026 06:31

Qiaolin-Yu mentioned this pull request Jul 24, 2026

Fix stale flashinfer-MLA fallback poisoning spec verify capture (trtllm_mla + tc_piecewise) #32288

Merged

tanth47 mentioned this pull request Jul 25, 2026

[Spec][DSpark] Remove overlap scheduler host syncs for pure TP #32374

Draft

3 tasks

Zhylkaaa pushed a commit to Zhylkaaa/sglang that referenced this pull request Jul 29, 2026

[Spec] DFlash: remove per-step host syncs so the CPU runs a full step…

befde24

… ahead (spec-v2 overlap) (sgl-project#31468) Co-authored-by: Hao Phan <htphan@nvidia.com>

jinzhenfan pushed a commit to jinzhenfan/sglang that referenced this pull request Jul 29, 2026

[Spec] DFlash: remove per-step host syncs so the CPU runs a full step…

2fbfe8a

… ahead (spec-v2 overlap) (sgl-project#31468) Co-authored-by: Hao Phan <htphan@nvidia.com>

kpham-sgl mentioned this pull request Aug 13, 2026

Add thanhhao98 to CI_PERMISSIONS.json #34649

Merged

1 task

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

[Spec] DFlash: remove per-step host syncs so the CPU runs a full step…

a96bb8c

… ahead (spec-v2 overlap) (sgl-project#31468) Co-authored-by: Hao Phan <htphan@nvidia.com>

jakki-amd pushed a commit to jakki-amd/sglang that referenced this pull request Sep 9, 2026

[Spec] DFlash: remove per-step host syncs so the CPU runs a full step…

23348d3

… ahead (spec-v2 overlap) (sgl-project#31468) Co-authored-by: Hao Phan <htphan@nvidia.com>

Navigation Menu

Uh oh!

[Spec] DFlash: remove per-step host syncs so the CPU runs a full step ahead (spec-v2 overlap) - #31468

[Spec] DFlash: remove per-step host syncs so the CPU runs a full step ahead (spec-v2 overlap)#31468
hnyls2002 merged 12 commits into
sgl-project:mainfrom
thanhhao98:htphan/dflash-overlap-hostsync

thanhhao98 commented Jul 16, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Jul 16, 2026

Uh oh!

thanhhao98 commented Jul 17, 2026 •

edited

Loading

Uh oh!

hnyls2002 commented Jul 18, 2026 •

edited

Loading

Uh oh!

hnyls2002 commented Jul 18, 2026

Uh oh!

github-actions Bot commented Jul 18, 2026

Uh oh!

github-actions Bot commented Jul 18, 2026 •

edited

Loading

Uh oh!

hnyls2002 commented Jul 18, 2026 •

edited

Loading

Uh oh!

hnyls2002 commented Jul 18, 2026

Uh oh!

Uh oh!

thanhhao98 commented Jul 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

thanhhao98 commented Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Correctness

Benchmarks (Kimi-K2.6-NVFP4 + DFlash draft, 8×B200 TP=8, aiperf 50K-prefix)

Checklist

CI States

Uh oh!

gemini-code-assist Bot commented Jul 16, 2026

Uh oh!

thanhhao98 commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

nsys before/after — the overlap improvement, Kimi-K2.6-NVFP4 cuteDSL, cc16

Uh oh!

hnyls2002 commented Jul 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hnyls2002 commented Jul 18, 2026

Uh oh!

github-actions Bot commented Jul 18, 2026

⚠️ Rebase Required Before Re-run

Uh oh!

github-actions Bot commented Jul 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hnyls2002 commented Jul 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hnyls2002 commented Jul 18, 2026

Uh oh!

Uh oh!

thanhhao98 commented Jul 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

thanhhao98 commented Jul 16, 2026 •

edited

Loading

thanhhao98 commented Jul 17, 2026 •

edited

Loading

hnyls2002 commented Jul 18, 2026 •

edited

Loading

github-actions Bot commented Jul 18, 2026 •

edited

Loading

hnyls2002 commented Jul 18, 2026 •

edited

Loading