feat: add lag confidence to decouple adaptive speculative pruning from critical path. - #2280
Open
weizhehuang0827 wants to merge 2 commits into
Open
feat: add lag confidence to decouple adaptive speculative pruning from critical path.#2280weizhehuang0827 wants to merge 2 commits into
weizhehuang0827 wants to merge 2 commits into
Conversation
…m the critical path. The adaptive speculative controller (DFlash/DSpark) prunes per-sequence how many draft tokens to validate from per-step confidence. Today that decision sits on the serial decode critical path: run_validate reads THIS step's confidence (produced by the draft forward that just finished), does a blocking D2H, and only then can the target forward proceed. Add enable_lag_confidence (default false, zero regression): the controller instead prunes from the PREVIOUS step's confidence (lag-1), read from the embedding cache. The decision no longer data-depends on this step's draft, so it runs in step_decode while the async draft forward is still in flight, leaving the critical path. This aligns with SGLang's independently-gated ConfidenceRelay channel (needs_confidence_relay). - EmbeddingCache: store this step's per-draft confidence per slot at validate end (one D2H past the validate sync, off the critical path); read it next step via read_lagged_confidence with the existing request_id freshness gate (slot reuse / first step -> fall back to full width, no prune). - DFlashWorkerImpl: split compute_prefix_lengths_from_probs out of compute_adaptive_prefix_lengths (linear cost model reused verbatim); add decide_lagged_prefix_lengths in step_decode; run_validate consumes the precomputed lagged decision when the flag is on. One site covers both DFlash and DSpark. Signal source mirrors the existing controller: ConfidenceHead output (DSpark) or proposal probs (DFlash). Measured (Qwen3 DSpark, single die): the controller decision on the critical path drops from ~64-78 us/call to ~0.1-0.3 us/call (the ~52 us of work moves into step_decode, overlapping the in-flight draft). SL7 lagged throughput 1109 tok/s vs 985 this-step adaptive (+12.6%), acc_rate 0.780 vs 0.803.
weizhehuang0827
requested review from
DongheJin,
DragonFive,
JimHsiung,
Kang-Meng,
liujinguang0125,
liutongxuan,
ustcfy,
xiao-yu-chen,
yingxudeng,
yinjiawei01 and
zhang-minchao
as code owners
August 20, 2026 12:23
weizhehuang0827
changed the title
Under adaptive speculative decode, when the controller prunes, run_validate rebuilt the target batch as a true varlen [Σ(prefix_i+1)] batch via a host-side per-seq loop (7 host vectors + heap allocs + H2D copies). That rebuild sat on the serial decode critical path: ~2.2% of a decode step, growing with the speculative block size. Under lag confidence the per-seq prune decision comes from the previous step's confidence, so it is known BEFORE this step's draft forward and has no data dependency on it. run_decode_draft already builds a validate batch on the host in the draft-overlap window, but it built the dense full-width batch that run_validate then discarded on every prune step. Build the pruned varlen batch there instead: the whole rebuild now overlaps the in-flight draft and leaves run_validate's critical path. - Hoist decide_lagged_prefix_lengths above run_decode_draft in step_decode and thread the decision in; the draft launch is async so this costs nothing. - prepare_overlap_validate_input: shared entry both DFlash and DSpark run_decode_draft call; builds the pruned varlen batch (recording the prune onto DraftBlock) when the lagged decision prunes, else the dense batch as before. - run_validate consumes the pre-built batch (DraftBlock.varlen_prebuilt) and skips its own rebuild; only fill_validate_input_from_draft_outputs_varlen — which needs this step's draft tokens — stays on the critical path. - Factor prefix_lengths_to_val_tokens out of the old apply_per_seq_varlen_prune, shared by the overlap path and the legacy this-step fallback. Derive effective_prefix from per_seq_val_tokens so it is correct on both paths. Flag off (default) is byte-for-byte the previous ordering: dense prep in run_decode_draft, this-step decision + rebuild in run_validate. Critical-path probe (DSpark, single die): host prune-block cost on run_validate's path drops SL7 873->285 us/step, SL16 1045->285 us/step — the metadata rebuild leaves the critical path and the residue (fixed ~285 us, the draft-token fill) is decoupled from block size. acc_rate unchanged (reorder is behavior-preserving).
ustcfy
approved these changes
Aug 22, 2026
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
2 participants
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.
Summary
The adaptive speculative controller (DFlash/DSpark) prunes per-sequence how many draft tokens to validate, based on per-step confidence. Today that decision sits on the serial decode critical path:
run_validatereads this step's confidence — produced by the draft forward that just finished — does a blocking D2H, and only then can the target forward proceed.This PR adds
enable_lag_confidence(defaultfalse, zero regression). When on, the controller prunes from the previous step's confidence (lag-1), read from the embedding cache. The decision no longer data-depends on this step's draft, so it runs instep_decodewhile the async draft forward is still in flight — leaving the critical path. This aligns with SGLang's independently-gatedConfidenceRelaychannel (needs_confidence_relay), which is a separate relay from the always-on token future map.What changed
DecodeState.confidencestores this step's per-draft confidence per slot at validate end — one D2H past the validate sync, off the critical path.read_lagged_confidencereads it next step, reusing the existingrequest_idfreshness gate (slot reuse / first step -> fall back to full width, no prune).compute_prefix_lengths_from_probsout ofcompute_adaptive_prefix_lengths(linear cost model reused verbatim); adddecide_lagged_prefix_lengthsinstep_decode;run_validateconsumes the precomputed lagged decision when the flag is on. One site covers both DFlash and DSpark. Signal source mirrors the existing controller: ConfidenceHead output (DSpark) or proposal probs (DFlash).speculative_config-> commonOptions-> runtimeOptions.Measurements (Qwen3 DSpark, single die, ShareGPT, conc=32)
Overlap — controller decision cost still on the critical path (inside
run_validate), measured with a temporary probe (not shipped):The ~52 us of real decision work relocates into
step_decode, overlapping the in-flight draft.Throughput / acceptance (32/32 completed, coherent, no crashes):
SL7 lagged is a net win (+12.6% over this-step adaptive, acceptance barely moves). SL16 lagged regresses vs this-step (stale confidence prunes less aggressively at the larger block) but still far beats static and stays coherent.
Test plan