feat: add lag confidence to decouple adaptive speculative pruning from critical path.

[weizhehuang0827]

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_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.

This PR adds enable_lag_confidence (default false, 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 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), which is a separate relay from the always-on token future map.

Building on that, the second commit takes the same lag-1 signal one step further. Once the prune decision is known before the draft forward, the pruned varlen target batch can also be built in the draft-overlap window rather than on the critical path. Today run_decode_draft already prepares a dense validate batch there — but on a prune step apply_per_seq_varlen_prune throws it away and rebuilds the true varlen batch inside run_validate, on the critical path. This commit builds the pruned batch directly in the overlap slot instead, so the host rebuild loop overlaps the in-flight draft. No new kernels — strictly a reorder, gated on the same flag.

Why not SGLang's device-side prune? SGLang lets the host decide only a scalar budget and does the gather/scatter on-device. That relies on attention metadata (seq_lens/cu_seqlens/block_tables) being device-resident. On xLLM's NPU path it is not: the ATB attention kernel reads seq_lens/q_lens as host scalar arrays via .hostData pointers for tiling, and ACL-graph bucket selection reads host kv_seq_lens. Per-seq attention metadata is intrinsically host-bound on NPU, so the device-prune cannot be ported as-is. Reordering the host rebuild into the overlap window is the NPU-appropriate equivalent.

What changed

• EmbeddingCache: DecodeState.confidence stores this step's per-draft confidence per slot at validate end — one D2H past the validate sync, off the critical path. read_lagged_confidence reads it next step, reusing the existing request_id freshness gate (slot reuse / first step -> fall back to full width, no prune).
• DFlashWorkerImpl (8bb4cecd3, lag confidence): 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).
• Overlap the varlen rebuild (fb3ee7d50, perf): under lag confidence, run_decode_draft builds the pruned varlen validate batch in its existing draft-overlap slot (replacing the dense prep that apply_per_seq_varlen_prune used to discard). run_validate then consumes the pre-built batch and only injects draft tokens (fill_varlen, which genuinely depends on this-step draft output). A shared entry prepare_overlap_validate_input covers both DFlashWorkerImpl and DSparkWorkerImpl (each has its own run_decode_draft override). Flag off = current ordering, byte-for-byte.
• New flag plumbed through speculative_config -> common Options -> runtime Options.

Measurements (Qwen3 DSpark, single die, ShareGPT, conc=32)

1. Lag confidence — controller decision cost (commit 8bb4cecd3)

Controller decision cost still on the critical path (inside run_validate), measured with a temporary probe (not shipped), per call:

SL	this-step	lagged
7	63.85 us/call	0.31 us/call
16	77.62 us/call	0.14 us/call

The ~52 us of real decision work relocates into step_decode, overlapping the in-flight draft.

Throughput / acceptance for the lag decision alone (num_prompts=32; 32/32 completed, coherent, no crashes):

SL	arm	out tok/s	acc_rate
7	static	613	0.576
7	this-step adaptive	985	0.803
7	lagged	1109	0.780
16	static	349	0.239
16	this-step adaptive	893	0.803
16	lagged	766	0.709

2. Overlap the varlen rebuild — critical-path cost (commit fb3ee7d50)

Temporary probe measuring the host prune-block cost that remains on run_validate's critical path, per decode step:

SL	before (rebuild in run_validate)	after (built in overlap window)
7	873 us/step	285 us/step
16	1045 us/step	285 us/step

The residual 285 us is fill_varlen (draft-token scatter) — it depends on this-step draft output, so it must stay serial. The varlen metadata rebuild (~588 us at SL7, ~760 us at SL16) relocates into the draft-overlap window. The residue is now decoupled from block size (285 at both SLs), confirming only the fill remains on the critical path. This probe is the trustworthy signal for this commit.

Three-arm end-to-end (num_prompts=48, single die, single run; nolag=this-step adaptive, lag-only=8bb4cecd3, lag+overlap=fb3ee7d50):

SL	arm	out tok/s	TPOT ms	acc_rate
7	nolag (this-step)	648	15.7	0.783
7	lag-only	631	19.3	0.743
7	lag+overlap	681	19.3	0.760
16	nolag (this-step)	616	17.9	0.827
16	lag-only	681	21.0	0.726
16	lag+overlap	691	20.1	0.735

lag+overlap leads both SLs (+7.9% at SL7, +1.5% at SL16 over lag-only). Caveat: single-die E2E throughput here is noise-dominated — over 5 reps the per-arm CV is ±6–15% and the overlap-vs-lag-only delta flipped sign between runs (−8.6% median in one 5-rep pass, +7.9% in this single run). So the E2E numbers are consistent with the improvement but do not statistically resolve it; the critical-path probe above is the load-bearing evidence. The acc_rate pattern is stable and expected: nolag (0.78–0.83) > lag arms (0.73–0.76) because lag prunes from stale (previous-step) confidence, and lag+overlap ≈ lag-only because the reorder is behavior-preserving — the same tokens are validated, only the batch is built earlier.

Test plan

• Builds clean in the CANN container; shipping binary verified free of the temporary probes.
• DSpark SL7/SL16 x {static, this-step adaptive, lagged}, 32/32 completed, no fatals, coherent output.
• Decision-cost probe confirms the controller decision leaves the critical path (~64–78 us -> ~0.1–0.3 us/call).
• Rebuild-cost probe confirms the varlen metadata rebuild leaves the critical path (~873–1045 us -> 285 us/step; residue = draft-token fill only).
• Three-arm E2E (nolag / lag-only / lag+overlap) SL7/SL16, 48/48 completed, coherent; lag+overlap acc_rate ≈ lag-only (behavior-preserving reorder confirmed).
• Flag off = byte-for-byte current behavior (this-step, dense-prep path unchanged).
• DFlash lag+overlap arm not yet benchmarked — DFlash uses proposal probs (weaker signal than a confidence head); worth a separate run to confirm acc_rate/throughput hold on the shared prepare_overlap_validate_input path.