fix(spec-decode): capture-stable graph slot + GQA workspace + M=1 base contract (#984 re-review)

.squad/decisions/inbox/deckard-specfix-graphslot.md

@@ -0,0 +1,103 @@
+# Decision drop — Deckard: speculative-decode capture-stability fix (#984 re-review)
+
+**Author:** Deckard (Systems Dev, CUDA/decode-perf)
+**Branch:** `squad/spec-decode-graphslot` (based on #984 head `82cde423`)
+**Supersedes review verdicts:** gaff-pr984-quality.md, chew-pr984-correctness.md
+**Feature:** opt-in (`ONNX_GENAI_SPEC_CAPTURED_VERIFY=1`), default-OFF; ⚠️ needs re-review.
+
+## What was wrong (both reviewers, both confirmed)
+- **BUG 1 (Gaff) — miss-path graph-slot hazard.** The M=1 base-decode graph and
+  the M=width captured-verify graph share the session's *single* device-graph
+  slot. The old miss branch set `retain_graph_on_rewind` then ran an M=1 decode,
+  and `reset_captured_verify()` ran *after* it — so an M=1 decode could replay a
+  stale width-W verify graph → glm "invalidated graph replay" + qwen
+  `CUDA_ERROR_ILLEGAL_ADDRESS`.
+- **BUG 2 (Chew) — GQA workspace under-prep.** The persistent GroupQueryAttention
+  workspace is sized by query width and prepared once at prefill for
+  `q_seq = prompt_len`. On a short degenerate-repetition prompt
+  (`prompt_len < W`, e.g. qwen 哈×20 → 8 tokens, W=9) the width-W verify entered
+  with an undersized workspace → executor rejected: *"GroupQueryAttention
+  workspace invariant mismatch: requires 184320, prepared 163840."*
+- **CONTRACT (Chew, binding).** Spec output MUST equal *plain greedy* (base token
+  from the **M=1 GEMV**), not "greedy-under-Marlin." Fusion sourced row 0 from
+  Marlin M=W, so at a genuine logit near-tie the base argmax could flip vs the
+  M=1 GEMV (the qwen one-token divergence).
+
+## Approach chosen: invalidate-both (mode-transition) — NOT the second graph slot
+I implemented the coordinator-sanctioned **fallback**: a correct mode-transition
+invalidation of the shared slot, plus the workspace re-prep and the row-0 M=1
+contract guard. I deliberately did **not** land a dedicated second device-graph
+slot in this increment.
+
+**Why (scoped rationale):** a real second slot requires changing the shared
+`ExecutionProvider` graph trait (`begin/end/abort/replay/reset_device_graph` take
+no slot param and are implemented by the CPU EP too), the runtime's single
+`CudaGraphLifecycle` field, CUDA graph-exec-handle multiplexing, and per-slot
+`device_graph_signature`/`capture_schedule` in the session executor — deep,
+cross-cutting, and capture-stability-risky. The reviewers required correctness
+(no crash, byte-identical), which the mode-transition fix delivers with a far
+smaller, auditable blast radius. **Second slot is documented as the scoped
+follow-on** (removes the miss↔engage re-warm tax; theoretical ceiling
+~111×W/B*). It is NOT needed for correctness and is a perf-only optimization.
+
+## The fix (three parts, all in the opt-in captured path)
+1. **BUG 1 — mode-transition slot invalidation** (`native_speculative.rs`,
+   `native_decode/mod.rs`). New `NativeDecodeSession::invalidate_graph_for_mode_switch()`
+   drops the EP slot and resets *both* the M=1 `graph_phase`/`inline_graph_phase`
+   and the retained `verify_capture.phase` to `NeedsWarmup` (verify bindings
+   kept). The driver tracks `prev_engaged` and calls it **only on an
+   engage↔miss transition, before** the incoming forward — so consecutive
+   same-mode steps still replay their captured graph at full speed (the miss-path
+   M=1 replay is preserved), but no path can ever replay a foreign/stale graph.
+   Invariant restored: `NeedsWarmup ⇔ slot empty`. The eager branch resets
+   `prev_engaged=false` (it invalidates internally).
+2. **BUG 2 — width-sized workspace reservation** (`native_decode/cuda.rs`,
+   `run_verify_captured_cuda`). Before warming/capturing the verify graph (phase
+   == `NeedsWarmup`), reserve the persistent GQA workspace for `q_seq = W` via
+   `prepare_cuda_prefill_workspace_with_step_inputs`, so the (possibly larger)
+   reservation is baked into the captured graph and every replay is valid.
+   Idempotent + only on (re)warm; a replay carries the reservation.
+3. **CONTRACT — row-0 near-tie → M=1 GEMV fallback** (`native_speculative.rs`).
+   After the fused forward, if the base row's top-1/top-2 margin ≤
+   `ONNX_GENAI_SPEC_ROW0_TIE_EPS` (default **1.0** logit), undo the fused KV fold
+   (`rewind(past)`), switch the slot to the base graph, and recompute the base
+   token from a fresh **M=1 GEMV** `decode` — so the committed base token is
+   byte-identical to plain greedy. Confident prompts sit far above eps (no
+   fallback, full speculative win); degenerate near-tie loops simply park on the
+   correct plain-decode floor.
+
+## Validation (GPU7, H200, verified idle before each run)
+- **f64 oracle** `matmul_nbits_marlin_numerics`: **7/7** (Marlin untouched).
+- **BUG 2 crash-fix:** qwen2.5-14b degenerate prompt (哈×20), captured verify +
+  Marlin ON, **ran to completion 160 tokens, 79 verify steps — no illegal
+  address, no workspace mismatch.**
+- **Byte-identical spec == plain greedy** (`generated_token_ids` diff):
+  - glm-4-9b generic prose: PASS (gate-off force-engage AND gate-on adaptive)
+  - glm-4-9b repetitive: PASS
+  - qwen2.5-14b degenerate (the previously-diverging near-tie): **PASS** —
+    identical stream (the row-0 M=1 fallback closes the flip).
+- **Favorable win:** glm repetitive **326 tok/s** (2.88× plain 113.5; past ORT
+  250), 97% acceptance, byte-identical.
+- **Generic no-regression:** gate correctly disengages (0 verify steps),
+  byte-identical, no crash/re-warm; 106 vs 112 plain (~5% CPU proposer tax,
+  inherent to running prompt-lookup — no longer a warmup crash).
+- **Default-OFF:** eager path (feature env unset) byte-identical to plain @160.
+- **Regression tests (both reviewer-required):**
+  (a) `native_captured_verify_engage_miss_reengage_no_stale_replay_cuda` —
+      engaged 2 verify steps, byte-identical, no illegal-address. PASS.
+  (b) `native_captured_verify_short_prompt_grows_gqa_workspace_cuda` — qwen-14b
+      short prompt (prompt_len < W), 160 tokens, 79 verify steps, no workspace
+      mismatch. PASS.
+  Plus 6 new fast unit tests for the row-0 near-tie classifier.
+- **fmt** clean; **clippy** `-p onnx-genai-engine --features cuda,native-backend
+  --all-targets -D warnings` clean except the 2 pre-existing `platform_capacity.rs`
+  u64 casts (untouched, as instructed).
+
+## New env knobs
+- `ONNX_GENAI_SPEC_ROW0_TIE_EPS` (default 1.0) — row-0 M=1-fallback margin; 0
+  disables (A/B only; default MUST stay on for the contract).
+
+## Follow-on (not in this PR)
+- Dedicated **second device-graph slot** (M=1 base + M=W verify both installed,
+  replay independently): eliminates the miss↔engage re-warm tax → generic prompts
+  also win, ceiling ~111×W/B*. Structural EP/session change; perf-only.

.squad/decisions/inbox/sebastian-specdecode-e2e.md

@@ -0,0 +1,103 @@
+# Decision drop — Stage-4 speculative decode with CAPTURED fused verify (native CUDA EP)
+
+- **Author:** Sebastian (Performance Engineer, CUDA & Perf)
+- **Date:** 2026-08-15
+- **Branch:** `squad/spec-decode-e2e` — head `a3fbd41c`
+- **PR:** draft (references #957 Stage 4)
+- **GPU:** single idle H200, `CUDA_VISIBLE_DEVICES=7` (re-checked `nvidia-smi` idle before every run); `source .cudaenv.sh`
+- **Feature gate:** `ONNX_GENAI_SPEC_CAPTURED_VERIFY=1` **and** `ONNX_GENAI_MARLIN_M_GT_1=1` (Marlin M>1 makes the M=width verify a single capturable graph). OFF by default → zero risk to the shipped M=1 decode path.
+
+## TL;DR
+The captured fused verify path is **built, correct (byte-identical to plain greedy on glm AND qwen), and delivers a real multiplicative decode win on high-acceptance prompts** — up to **215 tok/s on glm-4-9b-int4 (1.93× the 111 tok/s non-speculative baseline)**. It does **not** yet beat stock ORT (~250 tok/s) even on favorable prompts, and it **regresses on generic prose** (67 tok/s). Both the ORT gap and the generic-prompt regression trace to a single root cause: **the session has ONE device-graph slot**, so every lookup *miss* must invalidate + re-warm the width-W verify graph (running an eager forward while a graph is installed corrupts device memory — `CUDA_ERROR_ILLEGAL_ADDRESS`). A **dedicated second graph slot** (EP work) removes the re-warm penalty and is the scoped follow-on that would let realistic prompts win and clear ORT.
+
+## Design (what shipped)
+Fused, prompt-lookup-only, steady-state-only path in `native_speculative.rs::generate`:
+1. **Propose before the forward** (prompt-lookup needs only context, not base logits).
+2. **One fused `[1, 1+draft_width]` captured forward** over `[bonus ⊕ draft]`, padded to a **constant width** so a single constant-signature CUDA graph replays across steps: row 0 = base next-token distribution, rows `1..=k` = verify rows (numerically identical to the eager base+verify forwards).
+3. **Rewind to the accepted prefix** with `retain_graph_on_rewind=true` so the graph survives the per-step rewind and *replays* instead of re-capturing.
+4. **Empty-draft (lookup miss) → cheap M=1 eager base forward** (`decode_base_eager_keep_verify_graph`) rather than paying the ~2.2× width-W cost to commit a single token. This invalidates the graph (eager-alongside-installed-graph is unsafe), so the next hit re-warms.
+5. **Numerically-safe near-tie acceptance guard** (`greedy_accept`/`TieGuard`, from the earlier layer) — accepted stream is byte-identical to plain greedy; `near_tie_rejections` observes acceptance (not correctness) lost to tie-prone logits.
+
+## BEFORE/AFTER measurements — glm-4-9b-int4 (H200 GPU7, `--steady --tokens 200 --decode-skip 40 --runs 3 --warmups 1`, `ONNX_GENAI_MARLIN_M_GT_1=1`)
+
+| Prompt regime | Mode | K | mean-accepted B | acceptance | **decode tok/s** | vs non-spec |
+|---|---|---|---|---|---|---|
+| — | non-spec (plain M=1) | — | — | — | **111.2** | 1.00× |
+| Verbatim copy | captured-spec | 4 | 4.87 | 98.1% | 157.7 | 1.42× |
+| Verbatim copy | captured-spec | 8 | 7.92 | 89.7% | **205.9** | **1.85×** |
+| Near-pure repetition | captured-spec | 8 | 6.70 | 72.8% | **215.4** | **1.93×** |
+| Generic prose | captured-spec | 8 | 3.40 | 30.0% | 66.8 | 0.60× (regress) |
+| — | stock ORT-genai (prior session) | — | — | — | ~250.3 | — |
+
+## qwen2.5-14b-int4 control (block-32) — correctness + no-regression
+| Mode | K | B | acceptance | tok/s | tokens vs non-spec |
+|---|---|---|---|---|---|
+| non-spec | — | — | — | 126.7 | — |
+| captured-spec | 8 | 5.00 | 100% | 130.6 | **byte-identical ✓** |
+
+qwen's M>1 forward is costlier relative to M=1 (14B, block-32), so even at 100% acceptance it is ~flat (no regression). Marlin M>1 correctness holds across both block sizes.
+
+## Correctness (non-negotiable gate) — PASS
+- **glm** captured-spec token stream **== non-spec** on a mixed/generic prompt (heavy empty-draft interleave): byte-identical.
+- **qwen** captured-spec token stream **== non-spec** on a copy prompt (100% acceptance): byte-identical.
+- CUDA smoke test `native_prompt_lookup_matches_plain_greedy_cuda` (with `ONNX_GENAI_SPEC_CAPTURED_VERIFY=1`) — **green**.
+- 6 GPU-free `greedy_accept`/tie-guard unit tests — **green**.
+- `cargo fmt --all -- --check` clean; `clippy -p onnx-genai-engine --features cuda,native-backend --all-targets -D warnings` clean **except 2 pre-existing `platform_capacity.rs` `u64→u64` casts unrelated to this change** (last touched #947/#952).
+
+## The #1 lever to go PAST ORT (recommendation)
+A **second dedicated device-graph slot** in the session/EP so the M=1 base decode graph and the M=width verify graph can BOTH stay installed and be replayed independently. Today the shared single slot forces an invalidate + 2-forward re-warm on every miss→hit transition, which (a) caps the favorable-prompt win below ORT and (b) turns low-acceptance prompts into a regression. With no re-warm, the theoretical ceiling at 100% acceptance is ~`111 × W / B*` ≈ `111 × 9 / 2.2 ≈ 450` tok/s for K=8 — comfortably past ORT's 250. This is measurement-backed and is the logical next increment.
+
+Interim guidance: keep the feature **opt-in** (env-gated OFF). It is a clear win for high-acceptance / repetitive / copy / code-continuation workloads and must not be enabled globally until the second-slot work removes the generic-prompt regression.
+
+## Reproduce (exact commands)
+```bash
+source /home/justinchu/onnx-genai/.cudaenv.sh
+export CUDA_VISIBLE_DEVICES=7   # re-check nvidia-smi idle first
+cd .worktrees/sebastian-specdecode-e2e
+cargo build --release -p onnx-genai-bench --features bench-native,cuda --bin profile_native
+GLM=/home/justinchu/glm-e2e-artifacts/glm-4-9b-int4-cuda
+
+# non-spec baseline
+ONNX_GENAI_MARLIN_M_GT_1=1 ./target/release/profile_native --model "$GLM" --ep cuda \
+  --steady --tokens 200 --decode-skip 40 --runs 3 --warmups 1 --speculative none --prompt "<prompt>"
+
+# captured speculative (favorable prompt, K=8)
+ONNX_GENAI_MARLIN_M_GT_1=1 ONNX_GENAI_SPEC_CAPTURED_VERIFY=1 ./target/release/profile_native \
+  --model "$GLM" --ep cuda --steady --tokens 200 --decode-skip 40 --runs 3 --warmups 1 \
+  --speculative prompt-lookup --spec-ngram 3 --spec-tokens 8 --prompt "<repetitive/copy prompt>"
+
+# correctness: token-identity vs non-spec + CUDA smoke
+ONNX_GENAI_RUN_CUDA_SMOKE=1 ONNX_GENAI_MARLIN_M_GT_1=1 ONNX_GENAI_SPEC_CAPTURED_VERIFY=1 \
+  cargo test -p onnx-genai-engine --features cuda,native-backend \
+  --test native_speculative_driver native_prompt_lookup_matches_plain_greedy_cuda -- --nocapture
+```
+`ONNX_GENAI_SPEC_DEBUG=1` prints the per-step capture state machine (`phase`, `captures`, `replays`, `fallback`) for verification.
+
+---
+
+## UPDATE (2026-08-15) — adaptive hit-density gate: generic-prose regression removed (head `9b8776d4`)
+
+The single-graph-slot re-warm tax that caused the generic-prose regression is now **mitigated in the driver** (the second-slot EP work below is still the way to fully clear ORT everywhere, but is no longer required to make the feature safe).
+
+### What changed
+1. **`HitDensityGate`** (native_speculative.rs): the n-gram proposal is a cheap CPU search, so the driver **always proposes** but only **engages** the width-W captured verify once recent would-hit density (sliding 16-step window, ≥50% by default) predicts the warmed graph will *replay* repeatedly. Sparse hits stay on the plain-decode floor; clustered hits get the full ~1.9–3.1× win. Env-tunable: `ONNX_GENAI_SPEC_GATE` (`0`=always-engage A/B baseline), `ONNX_GENAI_SPEC_GATE_WINDOW`, `ONNX_GENAI_SPEC_GATE_MIN_HITS`.
+2. **Miss/gated-out steps run the normal M=1 *captured* decode** (device-graph replay = full baseline speed) instead of an M=1 *eager* forward, then reset the verify phase so the next engage re-warms cleanly. `decode` invalidates the slot on signature mismatch before any kernel launch → safe alongside a stale width-W graph. Removed the obsolete `decode_base_eager_keep_verify_graph` + always-false `preserve_graph` plumbing.
+
+### AFTER numbers (glm-4-9b-int4, H200 GPU7, steady tokens=200, decode-skip 40, 3 runs, CUDA-graph on)
+| prompt regime | gate | K | mean-B | accept | **tok/s** | vs non-spec 111 |
+|---|---|---|---|---|---|---|
+| generic prose | **ON (default)** | 8 | — | 0% (verify_steps=0) | **104.1** | 0.94× (was **0.60× / 66.8** always-engage) |
+| generic prose | OFF (always-engage) | 8 | 3.4 | 7% | 75.8 | 0.68× |
+| repetitive | **ON** | 8 | 8.64 | 95.5% | **347.3** | **3.12× — beats ORT 250** |
+| repetitive | OFF | 8 | 8.57 | 94.6% | 355.8 | 3.19× |
+
+Net: the gate **eliminates the catastrophic generic-prose regression** (−40% → −6%) at a **~2% cost on fully-favorable prompts**, and repetitive/copy/code-continuation prompts now **clear ORT 250** (347 tok/s).
+
+### Correctness (re-verified on cleaned head `9b8776d4`)
+- **glm** captured-spec (gate ON) token stream **== non-spec byte-identical** on generic prose (heavy miss path exercising the new fast gated decode). ✓
+- CUDA smoke `native_prompt_lookup_matches_plain_greedy_cuda` green (gate disabled in-test to exercise engage). ✓
+- **12** GPU-free unit tests green (6 `greedy_accept`/tie-guard + **6 new `HitDensityGate`**). ✓
+- fmt clean; engine clippy `--all-targets` clean except the 2 pre-existing `platform_capacity.rs` casts.
+
+### KNOWN pre-existing numerics boundary (flag for Chew — NOT introduced here)
+On **qwen2.5-14b** with a prompt that drives a **degenerate repetition loop**, captured-spec diverges from non-spec by **one token at a genuine logit near-tie**: the fused **Marlin M=W row-0** base distribution and the **M=1 GEMV** base distribution pick different argmax at the tie (both streams remain the same degenerate loop, offset by one repetition — benign). **Verified pre-existing**: the pre-gate committed head (`c55de4ef`) diverges identically on the same prompt; non-spec qwen is deterministic run-to-run. This is the same Marlin-vs-GEMV near-tie class handled by #960/#961's f64-oracle + `classify_parity`, surfacing at the base-token selection (which the draft-acceptance `TieGuard` does not cover). Recommend Chew decide whether the spec contract is "== plain greedy" (needs base token derived from the M=1 GEMV kernel) or "== greedy-under-Marlin" (accept the boundary).

crates/onnx-genai-bench/src/bin/profile_native.rs

@@ -507,11 +507,12 @@ fn print_speculative_observability(stats: &SpeculativeStats) {
     };
     println!(
         "speculative_stats: verify_steps={} proposed={} accepted={} acceptance={acceptance:.1}% \
-         multi_token_accepts={} tokens_per_verify_step={tokens_per_step:.2}",
+         multi_token_accepts={} near_tie_rejections={} tokens_per_verify_step={tokens_per_step:.2}",
         stats.verification_steps,
         stats.proposed_tokens,
         stats.accepted_tokens,
         stats.multi_token_accepts,
+        stats.near_tie_rejections,
     );
 }
 

crates/onnx-genai-engine/src/native_decode/backend.rs

@@ -70,9 +70,17 @@ impl NativeDecodeSession {
             // the single fixed-topology M=maxK graph and rewinds contents only —
             // `state.rewind` mutates just the mask tail + KV logical length, the
             // same data-driven mutation the captured graph already tolerates.
-            if !state.retain_graph_on_rewind {
+            //
+            // A full rewind to 0 (generation boundary / session reuse) always
+            // tears the graph down and drops the WP4 padded verify bindings, even
+            // under `retain_graph_on_rewind`: the next generation re-warms and the
+            // captured graph must not outlive the buffers it recorded addresses of.
+            if !state.retain_graph_on_rewind || target_len == 0 {
                 state.invalidate_graph(&mut self.session)?;
             }
+            if target_len == 0 {
+                state.verify_capture = None;
+            }
             state.rewind(target_len)?;
             self.current_len = target_len;
             return Ok(());