[Perf][Spec Decode] dflash/dspark: prefix-cache last-block-drop exemption + decode-path overhead reductions

[puririshi98]

Purpose

Batch-1 speculative decode with the dflash/dspark drafter family loses throughput in three independent places, each measured with per-span instrumentation over 1000 steady decode cycles on GB200 (Nemotron-3.5-Lightning-30B-A3B-NVFP4 + DSpark, production recipe). The worker host and the GPU are co-saturated on this path (~4.0 ms host vs ~4.1 ms GPU per cycle), so both GPU-side and host-side dead weight convert directly to inter-token latency:

1. Prefix caching drops a clean block on every hit. use_eagle() gates the EAGLE last-block drop, but dflash/dspark never cache lookahead-polluted KV (context KV is projected from target hidden states and positions only; the anchor token writes past the chunk end and is overwritten before its block can be hashed). The drop protects nothing for them and costs one full scheduler block (2,192 tokens, ~50 ms TTFT) per cache hit — on the mamba side too, since the align chunk-split backoff shares the predicate and the KV coordinator min()-reconciles.
2. The drafter's attention runs the serialized 2D kernel. The triton unified-attention 3D split-KV launch was gated to max_seqlen_q == 1; drafter/verify passes run small uniform q (1 bonus + N draft) that is still decode-shaped. On the drafter geometry (hq=32, hkv=2, hd=128, SW 1024 + sinks), the q=3 2D full-range pass takes 1011 us; the microbenchmark's q=1 3D full-range control on the same KV read measures 113.4 us at the shipped default segment count (NUM_PAR_SOFTMAX_SEGMENTS = 16, ~8.9x) and 61.9 us at the sweep's 32-segment point (~16.3x) — and a sliding window lights only ~1-2 of the equal-slice segments, so the segment layout is re-based on the visible window.
3. Dead host work every cycle. Under FULL cudagraph replay the propose path rebuilds draft attention metadata and slot mappings, then discards both (run_fullgraph reads only persistent input buffers) — 0.125 ms/cycle. The hybrid-mamba builder re-derives the single-token-prefill reclassification chain and the decode/prefill split from scratch each step, though both are value-determined constants in steady uniform decode.

Changes (four self-contained, independently revertable commits)

1. [Perf][Spec Decode] Exempt dflash/dspark from the EAGLE prefix-cache last-block drop — new SpeculativeConfig.prefix_cache_needs_last_block_drop() (True for eagle/eagle3/mtp, fail-closed for unknown eagle-family methods) wired into the KV-cache manager drop and the mamba-align backoff; num_prefill_lookahead unchanged.
2. [Perf][Kernel] Extend 3D split-KV to q <= 4, opt-in via a new max_query_len_3d argument (default None = exact previous behavior, batch-invariance still forces 2D), with per-query-token segm buffers and a window-based segment layout shared by main kernel and reduce. The window layout is gated on the same opt-in (max_query_len_3d > 1): every non-opted launch — including all existing q=1 sliding-window 3D decode (e.g. Gemma3 on TRITON_ATTN) — keeps the exact previous full-sequence segment layout (the USE_SW_SEGMENTATION=False branches are the unchanged upstream layout computation); extending the window layout to that default population is deferred to a follow-up with its own model eval. Method-gated fail-closed: only dflash/dspark raise the admission above q=1; eagle/eagle3/mtp and unknown speculative methods keep their verify passes byte-identical on the 2D launch and their q=1 3D decode on the previous segment layout.
3. [Perf][Spec Decode] Skip the dead rebuild on FULL replay behind a new fail-closed builder flag supports_skip_draft_rebuild (base False; Triton builder True iff R-SWA is inactive) AND cp_size == 1; any other builder/config keeps the rebuild for its side effects.
4. [Perf][Mamba] Fast uniform-decode metadata build: an all-decode shortcut that fires only when no row is prefilling and max_query_len <= decode_threshold (provably value-identical to the full reclassification + split chain), plus caching of the step-invariant gather-offset arange. No metadata value changes on any path.

Test Plan

Commands run (from the repo root, .venv per AGENTS.md; GPU suites inside the GB200 container):

# Commit 1 (CPU): 318 passed
pytest tests/v1/core/test_scheduler.py tests/v1/core/test_prefix_caching.py \
  tests/v1/core/test_mamba_align_chunk_split.py \
  tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
# Commit 2 (GPU, GB200): 24 new kernel-matrix tests + 1 gating test passed
# (25 in the selection; of 1609 passing overall; the 290 use_td-family
# failures reproduce identically on the unpatched container)
pytest tests/kernels/attention/test_triton_unified_attention.py
# Commit 2 (CPU): 2 passed (3D-admission gating + fail-closed method gate)
pytest tests/kernels/attention/test_triton_unified_attention.py -k "should_use_3d or method_gate"
# Commit 3 (CPU): 4 passed
pytest tests/v1/spec_decode/test_dflash_skip_rebuild.py
# Commit 4 (CPU): 31 passed (8 new + 23 pre-existing)
pytest tests/v1/attention/test_attention_splitting.py

• Commit 1: 318 passed (scheduler 153, prefix caching 96, mamba align chunk split 34, partial hits 35) incl. 3 new tests (predicate, ctor-spy wiring, align-split behavior).
• Commit 2: 26 new tests — 24 kernel-matrix tests on GB200 in-container (3D vs reference AND vs 2D, q 1-4 x SW {off,1024} x sinks x pow2/non-pow2 num_queries_per_kv, non-divisible segment tails, NaN-poison liveness; asserted via torch.testing.assert_close at atol 1.5e-2 / rtol 1e-2 against both references — the tighter maxdiff 8e-6 is from the standalone GB200 microbenchmark's q=1 3D-vs-2D full-range control, not from these tests) + 2 CPU gating tests (admission predicate; dflash/dspark-only method gate, fail-closed for eagle/eagle3/mtp/unknown). The GB200 GPU selection collected the 24 matrix tests plus the admission gating test: 25 passed.
• Commit 3: 4 new CPU tests (fail-closed gating: no builders, any unsafe builder, base default, R-SWA branch).
• Commit 4: 8 new CPU tests (fast path fires iff value-identical to the full chain; gather-offsets identity) alongside the 23 pre-existing.
• Series git am-clean on current main (re-verified 2026-08-27 on freshly-fetched origin/main fd57c4b7, after rebasing over [Model Runner V2][Spec Decode] Skip DP sync before EAGLE/MTP draft prefill #53694; git am --keep-non-patch of the whole 4-patch series in a throwaway worktree, all 17 resulting files byte-identical to branch HEAD, author/DCO preserved; no upstream commit in f25c580a..fd57c4b7 touches any changed file); ruff check/format clean; pre-existing GPU-test failures in our areas reproduce identically on the unpatched container (era drift, not introduced).

Test Result

A/B on GB200 and GB300, interleaved stock/patched arms on the same node with sha-verified overlays, aiperf 32K shared prefix / 2K in / 256 out, concurrency 1, temp 0, synthetic acceptance length 3. Two campaigns, labeled by code state: the validation jobs (GB200 + GB300) measured the pre-final code; because three post-validation changes touched functional lines (a mypy-driven attribute-form change, the rebase over #53694, and the split-KV opt-in gate in the hot kernel file), a shipped-HEAD recheck (GB200 job 1962848, interleaved stock/series 2v2, overlay regenerated in-job from the final patches and sha-verified = HEAD, failed_arms=0) re-ran the full A/B on exactly this branch HEAD. All pre-registered per-commit gates PASS on the recheck; it is the primary column below.

metric (attribution)	stock	this PR (shipped-HEAD recheck, GB200)	Δ	validation-phase value (pre-final code)
cache-hit TTFT (commit 1)	257.8 ms	203.4 ms	−54.4 ms (hit cached_tokens 28,496 → 30,688, +2,192 = exactly the block; bands exact in every arm)	−36.0 ms (252.1 → 216.1)
drafter FULL-graph GPU span (commit 2)	0.7986 ms/cycle	0.5923 ms/cycle	−0.206 ms	0.584 (0.582–0.584 across patched arms); standalone ITL −0.033 ms/tok
dead-rebuild host span (commit 3)	0.1249 ms/cycle	0.0071 ms/cycle	−0.118 ms	0.0073
target-side input-prep host span (commit 4)	1.7455 ms/cycle (instrumented baseline arm)	1.5766 ms/cycle	−0.17 ms (gate ≤ 1.60)	1.48–1.53 (−0.22..−0.29 within-arm vs 1.7455/1.7701 pre-commit-4 arms); ITL −0.030 ms/tok, no regression
decode tok/s, GB200 (series)	743.3 (ITL 1.3454)	776.6 (ITL 1.2877)	+4.5%, every patched rep > every stock rep (min separation 28.8 t/s; ITL fully separated)	732.4 → 801.4 (+9.4%, ITL 1.3655 → 1.2483; series n=3 vs interleaved stock n=2)
decode tok/s, GB300 (series, 2 reps — validation phase; the recheck was GB200-only)	716.3-724.3 (both clean stock arms; a third stock arm at 366.6 is excluded as a documented node-contention outlier)	765.7 / 763.5	+5.4-6.9% (vs the stock band)	(same — no GB300 recheck arm)
end-to-end request latency, GB200	600.9 ms	531.8 ms	−11.5%	600.3 → 534.5 (−11.0%)

Rows 2-4 are within-job attribution cells (per-span instrumentation) measured against the commit-1-instrumented baseline arm — commit 1 does not touch the decode path, so it stands in for stock on these spans; their stock column is that validation-phase baseline arm, their PR column the shipped-HEAD recheck cell.

Across the two GB200 campaigns the joint decode gain envelope is +33.3..+69.0 t/s (stock draw range across jobs 727–763 t/s, series 776–814): the recheck landed at the low end — reproducing an earlier same-day stock-vs-series cell (+36.6 t/s) — and the validation job at the high end. Three of six pre-registered aggregate recheck gates missed, each dispositioned in the artifact set, none implicating a commit: (a) joint decode +33.3 vs the ≥ +35 t/s gate and ITL −0.0577 vs ≤ −0.06 — missed by 1.7 t/s / 0.0023 ms under a uniform +1-6% host / +0.7-1.4% GPU drift measured on series-touched AND series-untouched spans alike (an environment signature; no series-specific span regressed); (b) a miscalibrated real-rejection acceptance gate — series 37.2% vs same-job stock control 52.0%, but the stock control's own run-to-run intervals span 28.4-52.0%, wider than the gate itself, and the series intervals (28.4-45.7%) sit inside the pooled stock envelope; (c) one greedy-argmax flip at a documented near-tie (details in the numerics paragraph below).

Memory: commit 2 sizes the three persistent split-KV softmax workspaces per admitted query token, so they grow only when a dflash/dspark drafter is configured, capped at 4x (MAX_QUERY_LEN_3D) on top of a base set by seq_threshold_3D. With decode cudagraphs disabled the threshold keeps its grid-derived value — at the production drafter geometry (hq=32, hkv=2, hd=128 → seq_threshold_3D=64) that is +48 MiB per builder instance (softmax_segm_output 16 MiB → 64 MiB; segm_max/expsum 0.125 → 0.5 MiB each). With decode cudagraphs enabled the threshold first snaps to the closest capture size, so the validated production recipe (cudagraph_mode FULL, capture sizes [1,2,4] → threshold 4) pays ~+3 MiB per builder (softmax_segm_output 1 → 4 MiB; segm_max/expsum +24 KiB each). Without a dflash/dspark drafter, sizes are unchanged.

Numerics gates, every production arm on both platforms and both campaigns: acceptance length pinned at 3.00 with per-position acceptance 1.000/1.000/0.000; with real rejection, acceptance parity within the measured small-sample floor (the series does not touch the sampler). Temp-0 stability: the patched hit-vs-miss top-20 logprob gap sits at/below each arm's own within-arm noise in every series arm — 0.25-0.32 in the three argmax-stable arms, matching the unpatched stock controls (0.24-0.30, one per EQ2 campaign). Greedy argmax was stable across all requests in those three arms; the fourth series EQ2 arm (shipped-HEAD recheck, job 1962848) flipped greedy argmax once, at the same documented 0.125-margin pos-0 near-tie the prior arm recorded — a flip at the greedy floor, not a numerics regression: that arm's gap (0.661) still sits below its own within-hit noise (0.666 = 5.3× the margin), and the unpatched stock controls themselves fail temp-0 self-consistency on the same probe (gen_text_a_eq_b = false in both stock runs). 2D-vs-3D is the same numerically-equivalent policy surface upstream already applies to all q=1 decode. Sub-fixes that did not reproduce their predicted span on hardware were dropped from the series before submission; every line above ships only what measured.

Related issues/PRs

Commit-1 area (searches dated 2026-08-26T09:42Z): not a duplicate. #53388 (draft) adds an opt-in family-wide disable_eagle_block_drop knob over the same scheduler sites; commit 1 changes the default by drafter-KV provenance, no knob, and the two compose. #47926 fixes the complementary dflash/dspark restored-region draft-KV defect (zero file overlap; acceptance-parity arms bound the interaction at nil). #53479, #50897 orthogonal; #53614/#53598 touch adjacent lines, mechanical rebase either order. Issues #51771, #47930. Re-swept 2026-08-27T08:56Z (the one area not refreshed alongside the 08-27 sweeps below), plus a 2026-08-27 audit pass that caught a miss of both sweeps — two additional hits in total, one from each. From the re-sweep, created after the original sweep: #53945 (OPEN, non-draft, created 2026-08-26T19:36Z) caches the Mamba state at the block-grid position of EAGLE resume and overlaps 4 of commit 1's 6 files (vllm/v1/core/sched/scheduler.py + all three shared test files); both of its scheduler hunks land inside Scheduler._mamba_block_aligned_split, the same function as commit 1's backoff hunk. Not a duplicate, but the composition is semantically non-trivial: its new tail_boundary/junction_stop logic is keyed on self.use_eagle with a rationale that only holds while the last-block drop applies — post-commit-1, use_eagle stays True for dflash/dspark while the drop no longer applies, so whichever lands second should re-key that logic on self.drop_prefix_cache_tail (commit 1's predicate); the shared test-file hunks are disjoint/adjacent, mechanical either order. From the audit, an original-sweep miss (created before that sweep) which the re-sweep also failed to surface: #51295 (OPEN, draft, created 2026-08-06T18:15:36Z) fixes a hybrid-attention cache miss caused by the EAGLE drop and overlaps the same 4 of commit 1's 6 files (vllm/v1/core/sched/scheduler.py + all three shared test files). Its single scheduler hunk (@@ -404) adds a self.use_eagle-keyed hash-unit backoff (tail_boundary = max(tail_boundary - self.hash_block_size, 0)) inside Scheduler._mamba_block_aligned_split — the same function as commit 1's backoff hunk (@@ -409), adjacent sub-blocks with no direct textual conflict, but the composition hazard is identical to #53945's: post-commit-1, use_eagle stays True for dflash/dspark while the drop no longer applies, so its tail_boundary backoff would fire for a drop that never happens; whichever lands second should re-key on self.drop_prefix_cache_tail. Not a duplicate (its mechanism is the hybrid-KV lookup alignment, not the drop exemption); the shared test-file hunks are mechanical either order.

Commit-2 area (searches dated 2026-08-27T02:41Z): #45450 (OPEN, non-draft, updated 2026-08-26) implements both of commit 2's mechanisms on the same three files — spec-decode 3D admission via a speculative_config-derived query length with per-query-token segment buffers, and window-relative 3D segmentation via a shared @triton.jit helper used by mainloop and reduce — measured on B300/NVFP4 with MTP over a v0.22.1 base (its body: not re-validated on main). Material differences here: (a) fail-closed method gate — only dflash/dspark opt in, eagle/eagle3/mtp verify numerics stay byte-identical, whereas #45450 admits MTP/EAGLE; (b) R-SWA and chunked-masking exclusions; (c) validated on current main-era containers with per-commit attribution cells. (Scope note: this PR's window layout is gated on the same opt-in, so existing q=1 sliding-window 3D decode is byte-unchanged here — matching #45450's q=1 exclusion rather than differing on it.) Also adjacent: #52879 (OPEN draft, same files, references #48076), #44652 (OPEN, relaxes the same 3D-launch gate for multi-query verify + non-causal support), and issue #48076 (batch>=12 3D->2D fallback — different trigger: the num_seqs threshold, which this series does not change). Re-sweep 2026-08-27T08:56Z surfaced a miss of the original sweep: #53930 (OPEN, non-draft, created 2026-08-26T17:51Z — about 9 h before the sweep above) adds a logger.warning_once in TritonAttentionMetadataBuilder.__init__ anchored exactly between the segm-buffer allocations commit 2 re-shapes and the rswa_window line commit 3 appends after, so it conflicts textually with this series whichever lands second (mechanical to resolve); more materially, its warning fires unconditionally whenever speculative_config is not None and asserts "Decode runs on the 2D path" (citing #48076) — factually wrong for dflash/dspark once commit 2 lands, so if it merges the warning needs scoping to speculative methods the 3D admission does not cover (spec_decode_max_query_len_3d() == 1).

Commit-3 area (searches dated 2026-08-27T04:10Z): not a duplicate. #53426 (OPEN, updated 2026-08-24) is the nearest neighbor — it also removes dead drafter-side work and touches the same dflash/speculator.py (plus commit 1's speculative.py; hunks disjoint from ours in both files), but skips a different thing under a different policy: an explicit opt-in, default-off knob to drop the K=0 draft sync forward when dynamic spec decode resolves K=0 (accepted for mtp/dflash only; dspark is refused), whereas commit 3 is a no-knob, fail-closed builder-flag skip of the draft metadata rebuild whose result FULL cudagraph replay discards — the draft forward still runs (as a replay) and the skip fires on ordinary K>0 steady decode for dflash AND dspark. Orthogonal dead work; the two compose, worst case a mechanical rebase in propose() (adjacent, non-overlapping hunks). #52782 (OPEN, NVTX/torch-profiler annotations to the model runner + DFlash speculator) and #53096 (OPEN, capture-log guard on needs_capture()) touch the same file with no semantic overlap — mechanical rebase either order; same class, added on the 2026-08-27T08:56Z re-sweep (a miss of the sweeps above, which it predates): #53978 (OPEN, non-draft, created 2026-08-27T02:32Z) hardens DFlash2 spec warmup against unfilled draft buffers in dflash/speculator.py hunks (__init__, _run_model) disjoint from commit 3's — no semantic overlap, mechanical rebase either order. Two more same-file opens the 08:56Z re-sweep missed (both added by a 2026-08-27 audit): #53970 (OPEN, non-draft, created 2026-08-27T01:12Z — inside the window that re-sweep demonstrably scanned) fixes the dflash batched token budget in dflash/speculator.py hunks (__init__ @@ -47, propose @@ -415) disjoint from commit 3's (@@ -31 / -184 / -456; its propose hunk sits ~35-40 lines above ours) — no semantic overlap, mechanical rebase either order; #53929 (OPEN, non-draft, created 2026-08-26T17:38Z) validates ragged GDN decode with Qwen3.5 DSpark and touches two series files — commit 1's vllm/config/speculative.py (@@ -844 vs c1's @@ -1795, disjoint) and commit 3's vllm/v1/attention/backend.py (@@ -569, in AttentionCGSupport, vs commit 3's @@ -597 in AttentionMetadataBuilder — disjoint, different classes) — no semantic overlap with either commit's mechanism, mechanical rebase either order. #53694 (MERGED 2026-08-27; re-swept 2026-08-27T04:56Z) landed in the same propose() lines commit 3 rewrites — it hoists the DP-sync num_tokens_across_dp read out of the dispatch call's return; orthogonal mechanism (DP-sync plumbing, not a metadata-rebuild skip), inert at DP=1 (dispatch_cg_and_sync_dp returns no sync state when dp_size == 1, the measured recipe), and this series is rebased over it (commit 3's only conflict, context-only). No open PR implements a replay-time metadata-rebuild skip (dflash-speculator and draft-metadata sweeps).

Commit-4 area (searches dated 2026-08-27T04:10Z): not a duplicate. #39936 (OPEN, last updated 2026-06-04) cheapens split_decodes_and_prefills itself (env-gated CPU-launch-overhead reduction inside the function, plus block_table.py slot-mapping changes, Blackwell-targeted); commit 4 instead avoids re-running the reclassification chain + split at all on provably all-decode steady steps, and its utils.py change is confined to mamba_get_block_table_tensor — zero hunk overlap with #39936's split_decodes_and_prefills hunks, and #39936 does not touch mamba_attn.py. Different mechanism, composes: #39936 speeds the general path, commit 4 removes it from the steady-decode loop while keeping it as the fail-closed fallback. Also in the area, an enumeration miss of the 04:10Z sweep (added by a 2026-08-27 audit): #48188 (OPEN, non-draft, created 2026-07-09) speeds up _compute_chunk_metadata internals by ~6x in the same mamba_attn.py — its single hunk (@@ -281,43, spanning lines 281-323) is disjoint from all of commit 4's hunks (26 / 197 / 450 / 529 / 815), and its mechanism (cheapening the chunk-metadata computation itself) is different from commit 4's (skipping the reclassification chain + split entirely on provably all-decode steady steps); the two compose — commit 4's fallback path would simply run #48188's faster computation — mechanical rebase either order. The split_decodes_and_prefills / mamba-metadata-builder sweeps surfaced no open PR with commit 4's shortcut (that statement stands; #48188 is the nearest same-file neighbor, not a duplicate).

AI assistance disclosure

Developed with AI assistance (Claude Code): profiling, root-cause tracing, the fixes, and the hardware A/Bs. I have reviewed every changed line and can defend the series end-to-end.

Signed-off-by: Rishi Puri riship@nvidia.com

🤖 Generated with Claude Code