Repository navigation
[DRAFT] Captured spec-decode: byte-identity achievable & capture-safe; speedup fundamentally acceptance-bound - #1006
Draft
justinchuby wants to merge 13 commits into
Draft
justinchuby wants to merge 13 commits into
justinchuby wants to merge 13 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>
Keep cold native generation workspace preparation prompt-shaped so speculative verify width does not leak into later M=1 fallback decode. Add qwen W=7/W=9 regression coverage for captured speculative decode byte identity. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Phase-1 diagnosis of why captured speculative decode diverges from plain M=1 greedy, plus a proven byte-identity reference path (opt-in, default OFF so there is zero production change vs the aceed71 baseline). Findings: - The prior byte-identity reference (plain greedy with Marlin prefill ON, ONNX_GENAI_MARLIN_M_GT_1=1) is non-deterministic run-to-run; the production default (Marlin prefill OFF) is deterministic and is the correct reference. This is why the three prior attempts, validated against MARLIN=1, were unwinnable. - The verify divergence is a K-reduction reassociation: any batched M=W int4 GEMM (Marlin fp16-accumulate, or the portable fp32 tiled GEMM) sums K in a different order than the M=1 GEMV greedy uses, flipping the argmax at sensitive positions. Not a margin/tie problem. - Routing every verify row through the exact M=1 GEMV (per-row decomposition) is byte-identical by construction: qwen2.5-14b 5/5 (W=5..9), glm-4-9b 4/5 (W=5..8; W=9 flips at token 157 in the q_seq=W block ATTENTION path -- a separate root cause needing fp32 verify attention). The per-row path re-reads the int4 weights W times per verify step (~6x slower than greedy), so it is committed OPT-IN only (ONNX_GENAI_SPEC_PERROW_VERIFY=1, default OFF) as a validated byte-identity reference oracle. The speed-preserving completion (batched M=W fp32-accumulate GEMV that reads weights once + fp32 verify attention for glm W=9) is scoped in .squad/decisions/inbox/leon-spec-verify.md. Every per-row launch is static-grid capture-safe (verify graph still captures/replays). Oracle 25+8 matmul_nbits tests green; fmt clean; no new clippy. 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 #1006 +/- ##
==========================================
- Coverage 78.99% 78.78% -0.22%
==========================================
Files 356 364 +8
Lines 150211 146798 -3413
Branches 150211 146798 -3413
==========================================
- Hits 118661 115654 -3007
+ Misses 26897 26528 -369
+ Partials 4653 4616 -37
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
|
…/lm_head) Single-launch batched M=W verify for the General single-warp int4 GEMV: hoists weight load/dequant/scale once and applies per-row hfma2, so each row is byte-identical to M=1 greedy while weights are streamed once. Covers q/o/lm_head; down-projection, gate_up SwiGLU and split-K (k/v) still fall back to the per-row oracle. Opt-in via ONNX_GENAI_SPEC_BATCHED_VERIFY. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…capturability diagnosis
Adds single-launch batched byte-identical captured-verify GEMV kernels for the
tall-skinny down-projection and the fused gate_up SwiGLU (both plain and the
rmsnorm-prologue variant qwen2.5-14b uses), routing each M=W verify GEMM to the
batched sibling of the exact M=1 GEMV greedy uses. Each row's K reduction is
bit-identical to M=1 by construction, so the verify logits are byte-identical to
plain greedy; the weights are streamed once per launch and every launch is
static-grid, allocation-free and CUDA-graph capture-safe (advertises
last_call_capture_safe), so the composed verify graph captures whole-graph
(segments=1) instead of segmenting on MatMulNBits[KernelCaptureUnsupported].
Also fixes run_f16_per_row_verify to dispatch gate_up nodes through
run_f16_gate_up_swiglu (they are dispatched in run(), not run_f16()).
Validation (qwen2.5-14b-int4, deterministic Marlin-off reference):
* G1 byte-identity: qwen 5/5 (W=5..9), glm 3/5 (W=5,7,8; W=6,9 diverge in the
attention path, not the GEMV — a separate q_seq=W flash-attn reassociation).
* G3 capture: verify graph is whole-graph (segments=1), captures>0, zero
capture-time device alloc.
* G4: f64 oracle green; fmt clean; only the pre-existing allowed
platform_capacity.rs clippy casts.
Diagnosis instrumentation: spec_capture debug now reports captured segment
count; the leverb phase0 probe can route its M=K capture attempt through the
batched verify path (ONNX_GENAI_LEVERB_BATCHED) to measure the whole-graph
captured M=W replay wall.
DECISIVE finding (leverb phase0, qwen, whole-graph captured replay walls):
M=1 = 7.9 ms; batched byte-identical M=8 = 65 ms (8.3x); Marlin M=8 = 36 ms
(4.5x). An M=W verify forward is inherently ~4.5x (Marlin) / ~8x (byte-
identical) an M=1 step on this model/HW, so captured spec-decode beats greedy
only when mean_accept exceeds that ratio (>4.5 Marlin / >8.2 byte-identical),
whereas prompt-lookup on normal prompts yields mean_accept 2.4-3.3. G2
(spec faster than greedy) is therefore not achievable together with byte-
identity at realistic acceptance; see decision drop.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… G2 fundamentally acceptance-bound Records the decisive leverb-phase0 captured replay walls (M=1=7.9ms, byte-identical M=8=65ms/8.3x, Marlin M=8=36ms/4.5x), the break-even acceptance math (Marlin>4.5, byte-identical>8.2 vs measured 2.4-3.3), the shared graph-slot re-warm thrash, and the recommended contract relaxation (ship byte-identity as an opt-in determinism mode; re-scope the speed gate to high-acceptance regimes; give verify its own graph slot). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
changed the title
This was referenced Aug 15, 2026
Owner
Author
|
Triage on this draft, which has been open since 2026-08-15. At 16 files and +4008/-165 this is the largest of my outstanding drafts, and its headline is a negative that matters: byte-identity under capture is achievable and capture-safe, but the speedup is fundamentally acceptance-bound. That conclusion is worth publishing regardless of whether the implementation lands -- it tells anyone considering speculative decoding where the ceiling actually is, which is a design fact, not an implementation detail. I cannot validate the implementation half here. This machine no longer has a CUDA runtime (no Proposed disposition, in order of preference:
Either way this should stop being a draft. A 4000-line branch ages badly against a tree that has taken several merges to the same decode path since it was opened. |
This branch has not been deployed
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.
[DRAFT] Captured spec-decode: byte-identity achievable & capture-safe; speedup is fundamentally acceptance-bound
Status: DRAFT — reports a fundamental finding, does NOT claim all gates green. Per the mission's "stop and report with evidence rather than ship a hollow fix" clause.
Summary
Byte-identity to plain M=1 greedy is achievable and CUDA-graph capture-safe. But spec-decode faster-than-greedy (G2) cannot coexist with byte-identity at the acceptance rates prompt-lookup produces on normal prompts — and this is not a GEMM/kernel problem: even the fast, non-byte-identical Marlin verify loses to greedy at these acceptance rates.
What works (committed, validated)
Single-launch batched byte-identical captured-verify GEMV kernels (General q/o/lm_head, tall-skinny down-projection, fused gate_up SwiGLU incl. the rmsnorm-prologue form qwen uses). Each M=W row's K-reduction is bit-identical to the M=1 GEMV greedy uses -> byte-identical by construction. All advertise
last_call_capture_safe-> verify graph captures whole-graph (segments=1), zero capture-time alloc.Gates
platform_capacity.rsclippy casts.Decisive evidence (qwen, whole-graph captured replay walls)
An M=W verify forward is inherently ~4.5x (Marlin) / ~8x (byte-identical) an M=1 step (attention, 152k-vocab lm_head, MLP all scale ~M). Break-even (no thrash): spec beats greedy only when
verify_ms / mean_accept < 7.9ms-> Marlin needs mean_accept >4.5, byte-identical >8.2; measured mean_accept on normal prompts is 2.4-3.3. Additionally the verify graph shares the single EP device-graph slot with M=1 decode, so miss-path decodes evict it (re-warm thrash), pushing end-to-end spec to 0.04-0.21x.Recommendation
ONNX_GENAI_SPEC_BATCHED_VERIFY), not a speed feature.Full analysis + repro:
.squad/decisions/inbox/leon-spec-verify.md.