Repository navigation
fix(spec-decode): capture-stable graph slot + GQA workspace + M=1 base contract (#984 re-review) - #988
Closed
justinchuby wants to merge 8 commits into
Closed
fix(spec-decode): capture-stable graph slot + GQA workspace + M=1 base contract (#984 re-review)#988justinchuby wants to merge 8 commits into
justinchuby wants to merge 8 commits into
Conversation
…e decode Extract the greedy speculative accept decision into a pure, unit-tested `greedy_accept` with a `TieGuard`. The default (TieGuard::STRICT, eps=0) is byte-identical to plain greedy: commits are always the strict argmax. When ONNX_GENAI_SPEC_TIE_EPS>0 the guard detects draft tokens that are near-tie co-winners of the argmax (logit within eps of the row max) and either counts them (diagnostic, default) or, with ONNX_GENAI_SPEC_TIE_TOLERANT=1, accepts them to recover acceptance on tie-prone (block-32) models. Surface `near_tie_rejections` in SpeculativeStats and profile_native observability. Refs #957 (Stage 4 speculative-decode e2e). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…e decode Add the CUDA-graph-captured verify path (WP4 of #957): a persistent padded [1, width] token/position/logits binding set and a warm->arm->capture->replay state machine (DecodeCudaState::verify_capture + decode_verify_captured), modelled on run_one_token and the Increment-0 capture probe. Graceful eager fallback on NotCapturable keeps correctness independent of capture. Session surface: set_retain_graph_on_rewind / is_cuda / reset_captured_verify; rewind(0) tears down the retained verify graph before freeing its buffers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… wire fused driver Re-adds the CudaKvDebugStats struct header dropped when inserting logits_bytes_to_rows, converts the read_bytes SessionError into anyhow, and wires the prompt-lookup driver's fused captured-verify branch behind ONNX_GENAI_SPEC_CAPTURED_VERIFY (propose-before-forward, retained graph, row0=base / rows1..=k=verify, rewind to accepted prefix). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Completes the Stage-4 captured M=width verify path for native speculative decode (env ONNX_GENAI_SPEC_CAPTURED_VERIFY, requires ONNX_GENAI_MARLIN_M_GT_1): * Fixed-width fused forward: propose-before-forward, run one [1,1+draft_width] captured graph (row 0 = base distribution, rows 1..=k = verify rows), then rewind to the accepted prefix. A single constant-signature graph occupies the shared session slot so it CAPTURES once and REPLAYS across steps. * Empty-draft (lookup miss) runs a cheap M=1 eager base forward (decode_base_eager_keep_verify_graph) instead of paying the ~2.2x width-W verify cost for a single token. Running an eager forward alongside an INSTALLED device graph corrupts memory (illegal address), so the miss path invalidates and the next hit re-warms — the single-graph-slot cost that caps realistic prompts (a dedicated second graph slot is the scoped follow-on). * Threaded a preserve_graph flag through the eager row helpers. Correctness (byte-identical to plain greedy) verified: glm-4-9b-int4 and qwen2.5-14b-int4 captured-spec token streams == non-spec, CUDA smoke test green, 6 accept-guard unit tests pass. GPU7 (H200), steady tokens>=160. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…RE/AFTER Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The fused captured-verify path won 1.85-1.93x on high-acceptance prompts
but REGRESSED generic prose to ~0.60x: on realistic prompts most steps
miss the prompt-lookup, and an isolated hit re-warms the width-W verify
graph (a fresh capture, not a replay) to commit only B tokens at ~2.2x
the M=1 cost -- a net loss that never amortizes.
Two fixes, both preserving byte-identity:
1. HitDensityGate: the n-gram proposal is a cheap CPU search, so always
propose, but only ENGAGE the width-W captured verify once recent
would-hit density (sliding 16-step window, >=50% default) predicts the
warmed graph will replay repeatedly. Sparse hits stay on the plain
decode floor; clustered hits get the full 1.9x win. Env-tunable
(ONNX_GENAI_SPEC_GATE{,_WINDOW,_MIN_HITS}); GATE=0 restores always-engage.
2. Gated-out / miss steps now run the normal M=1 CAPTURED decode
(device-graph replay = full plain baseline speed) instead of an M=1
EAGER forward, then reset the verify phase so the next engaged step
re-warms cleanly. `decode` invalidates the slot on signature mismatch
before launching any kernel, so it is safe alongside a stale width-W
graph. This removes the residual eager-vs-replay tax on miss-heavy
prompts (generic prose 81 -> 104 tok/s, vs 111 non-spec baseline).
Result (glm-4-9b-int4, H200 GPU7, steady): generic prose 66.8 -> 104
tok/s (no catastrophic regression); repetitive 347 tok/s (beats ORT 250);
gate costs ~2% on fully-favorable prompts. glm spec == plain greedy
byte-identical; qwen control byte-identical on non-degenerate prompts
(pre-existing Marlin-row0-vs-M1-GEMV near-tie divergence in qwen's
degenerate repetition loop is unchanged by this commit).
Removes the now-obsolete decode_base_eager_keep_verify_graph and the
always-false preserve_graph plumbing on run_cuda_eager_rows. Adds 6
HitDensityGate unit tests; CUDA smoke test disables the gate to exercise
the engage path on its short fixture.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… near-tie note Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…+ M=1 base contract (#984 re-review) Fixes the two capture-stability bugs and the greedy contract that Gaff and Chew rejected in #984, on the opt-in captured-verify speculative path. - BUG 1 (Gaff): the M=1 base-decode graph and the M=width captured-verify graph share the session's single device-graph slot. Track prev_engaged and invalidate the slot (both graph phases + verify phase -> NeedsWarmup) on an engage<->miss TRANSITION, before the incoming forward -- so an M=1 decode can never replay a stale width-W verify graph (glm "invalidated graph replay" / qwen CUDA_ERROR_ILLEGAL_ADDRESS). Consecutive same-mode steps still replay at full speed. New NativeDecodeSession::invalidate_graph_for_mode_switch(). - BUG 2 (Chew): the persistent GroupQueryAttention workspace is sized by query width and prepared only at prefill for q_seq=prompt_len. When prompt_len < W the width-W verify entered undersized ("workspace invariant mismatch"). Reserve the workspace for q_seq=W in run_verify_captured_cuda before warming/capturing. - CONTRACT (Chew, binding): spec output must equal plain greedy (base token from the M=1 GEMV), not greedy-under-Marlin. Add a row-0 near-tie classifier (ONNX_GENAI_SPEC_ROW0_TIE_EPS, default 1.0): on a base-row near-tie, rewind the fused fold and recompute the base token from a fresh M=1 GEMV decode, so the committed base token is byte-identical to plain greedy. Chose the coordinator-sanctioned invalidate-both fallback over a dedicated second graph slot (documented as scoped perf-only follow-on). Validation (GPU7 H200): f64 oracle 7/7; qwen degenerate no crash 160 steps/79 verify; spec==plain byte-identical on glm generic+repetitive and qwen degenerate; glm favorable 326 tok/s (2.88x, past ORT); generic no re-warm crash; default-OFF byte-identical. New CUDA regressions (engage/miss/re-engage; qwen workspace growth) + 6 row-0 classifier unit tests. fmt+clippy clean (except 2 pre-existing platform_capacity casts). Feature stays opt-in (ONNX_GENAI_SPEC_CAPTURED_VERIFY=1), default-OFF. Refs #957, #984. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #988 +/- ##
==========================================
+ Coverage 78.99% 79.47% +0.47%
==========================================
Files 356 356
Lines 150211 150211
Branches 150211 150211
==========================================
+ Hits 118661 119375 +714
+ Misses 26897 26179 -718
- Partials 4653 4657 +4
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
This was referenced Aug 15, 2026
justinchuby
added a commit
that referenced
this pull request
Aug 15, 2026
Scribe consolidation: merge 7 decision drops (#984/#986/#988 reviews + wide-load design) into decisions.md, refresh stale now.md to current base-decode-192/ORT-1.30x reality, add session+orchestration logs. .squad/-only. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This was referenced Aug 15, 2026
Owner
Author
|
Closing as part of the spec-decode wrap-up. The captured-speculative-decode investigation reached a definitive, evidence-backed conclusion (see PR #1006 and
Spec-decode is therefore parked, not abandoned: it ships as an opt-in determinism mode in #1006, and would only pay off as a speed feature with a real draft model or batch>1 (plus its own device-graph slot — the idea in this PR). The branch is preserved for that future revival. Full rationale is in the decision log. |
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
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.
Speculative-decode capture-stability fix (#984 re-review)
Fixes the two capture-stability bugs and the greedy contract that Gaff and Chew rejected in #984, on the opt-in captured-verify speculative path. The glm path already worked; this is a targeted fix, not a rewrite. Feature stays opt-in (
ONNX_GENAI_SPEC_CAPTURED_VERIFY=1), default-OFF byte-identical to main. Refs #957, #984.Bugs fixed
invalidate_graph_for_mode_switch()drops the slot + resets both graph phases and the verify phase toNeedsWarmup; the driver tracksprev_engagedand invalidates only on an engage↔miss transition, before the incoming forward. No path can replay a stale/foreign graph; consecutive same-mode steps still replay at full speed. Kills the glm "invalidated graph replay" + qwenCUDA_ERROR_ILLEGAL_ADDRESS.q_seq=prompt_len. Whenprompt_len < W, the width-W verify entered undersized (workspace invariant mismatch: requires 184320, prepared 163840).run_verify_captured_cudanow reserves the workspace forq_seq=Wbefore warming/capturing.ONNX_GENAI_SPEC_ROW0_TIE_EPS, default1.0): on a base-row near-tie, rewind the fused fold and recompute the base token from a fresh M=1 GEMV decode → byte-identical to plain greedy.Approach: invalidate-both (mode-transition), not a second graph slot
Landed the coordinator-sanctioned fallback. A real second device-graph slot needs changes to the shared
ExecutionProvidergraph trait (implemented by the CPU EP too), the runtime's singleCudaGraphLifecycle, graph-exec-handle multiplexing, and per-slot session state — deep, cross-cutting, capture-stability-risky. The reviewers required correctness, which the mode-transition fix delivers with a small auditable blast radius. Second slot is documented as a scoped, perf-only follow-on (removes the miss↔engage re-warm tax).Validation (GPU7, H200, verified idle)
matmul_nbits_marlin_numerics-D warnings/ fmtplatform_capacity.rscasts)Tests added
native_captured_verify_engage_miss_reengage_no_stale_replay_cuda— BUG 1 (engaged, byte-identical, no illegal-address).native_captured_verify_short_prompt_grows_gqa_workspace_cuda— BUG 2 (qwen-14b, 160 tok, 79 verify, no workspace mismatch).Decision drop:
.squad/decisions/inbox/deckard-specfix-graphslot.md.