diff --git a/.squad/decisions/inbox/gaff-mtp-graph-retain-capture-unsafe-blocker.md b/.squad/decisions/inbox/gaff-mtp-graph-retain-capture-unsafe-blocker.md new file mode 100644 index 0000000000..775a47a78b --- /dev/null +++ b/.squad/decisions/inbox/gaff-mtp-graph-retain-capture-unsafe-blocker.md @@ -0,0 +1,26 @@ +### 2026-08-21: MTP graph retention across verify+rewind is capture-unsafe — root cause localized, dormant seam landed, no speedup number + +**By:** Gaff + +**What:** +Investigated the "make MTP `replays>0`" speedup lever (turn the 78% acceptance into a real net win vs the 62.56 tok/s greedy baseline). Landed the two-site invalidation seam as a **documented dormant flag** (`retain_decode_graph_across_spec`, default OFF) plus test-only setters, mirroring the existing dormant `retain_graph_on_rewind`/option-c convention. **No speedup number exists** — naive retention is both insufficient and capture-unsafe (GPU-verified). Behavior is byte-inert vs origin/main `843b0bf7d` (flag never enabled). + +Root cause of MTP `replays=0` (empirically localized, GPU, 10 verify-steps steady): +- The captured M=1 decode graph is torn down TWICE per spec step: **30× at `cuda.rs` `run_cuda_eager_rows_owned`** (the eager M>1 verify forward calls `invalidate_graph`) + **21× at `backend.rs:145` `rewind_inner`** (commit rewind). The capture state machine needs ≥2–3 consecutive same-shape M=1 steps to reach a replay; the per-step invalidations tear it down first. `binding_signature` (bindings.rs:1090) keys only on physical_shape+device_ptr (not logical length/data), which is why greedy replays fine as KV grows but MTP never does. + +Why naive retention fails (GPU-verified via temporary A/B env toggles, since reverted): +- Retain-across-rewind-only (safe, replays=0, ~15 tok/s) and retain-across-verify-only (safe, replays=0, ~15 tok/s) each leave the *other* site invalidating, so no graph ever survives to replay. +- Retain across **BOTH** — the only config that lets a graph survive a full verify→rewind→re-advance cycle and replay — produces **non-finite logits (exit 1, finite-guard caught)**. Cause: the eager M=K verify reserves a larger StepScoped `step_workspace` that `release_step_workspace` (bindings.rs:793) frees after the run; the captured M=1 graph baked the old workspace pointer, so the next M=1 replay reads a stale/moved address → NaN. Greedy is immune because every step is the same M=1 shape → arena returns the same address. + +Deeper structural blocker (the real reason there's no easy win): +- Even with the M=1 replay workspace issue solved, the **M=K verify forward stays eager (un-captured)** and pays full per-op launch overhead — exactly what graphed greedy avoids (55.91 vs eager). A real MTP speedup requires **capturing the verify itself** into a fixed-shape replayable graph: option-c "padded verify capture" (pad every forward to maxK rows, capture once at maxK, replay for base and verify), plus a pinned/isolated StepScoped workspace and a shape-keyed graph slot (the EP currently holds a single `device_graph_signature`, can't hold M=1 base + M=K verify at once). This is a multi-turn executor workstream; dormant scaffolding already exists (`configure_padded_verify_capture`, `padded_query_capacity`, `retain_graph_on_rewind`, all `#[cfg(test)]`). + +**Why:** +The coordinator's fallback clause: if graph retention across differing shapes proves infeasible in one turn, land the root-cause analysis + any partial capture-safety improvement and report precisely — do not fabricate a speedup. Retention is not just infeasible in one turn, it's **wrong** as scoped (option-b eager verify can never beat graphed greedy no matter how the M=1 graph is retained). Landed the documented two-site seam + flag as the low-risk building block the future option-c work will flip on, keeping the exact GPU evidence inline at both invalidation sites so the next agent has the map. Recommend the next turn be scoped explicitly to **option-c padded verify capture** (capture the verify), not further M=1-retention tuning. + +**Validation (GPU, H200 CUDA_VISIBLE_DEVICES=5, all 8 idle 0MiB/0%; build `--features bench-native,native-cuda,cuda-13000`; ORT 1.28 cuda13 `.ort-cuda-1.28/root`; int4 block-32; harness `profile_native`; branch `squad/mtp-retain-graph-on-rewind` off origin/main `843b0bf7d`):** +- Full lib suite `cargo test -p onnx-genai-engine --no-default-features --features native-backend --lib`: **575 passed, 0 failed, 1 ignored** (greedy inert). +- Inertness on real Qwen3.8-27B int4 hybrid artifact `/home/justinchu/qwen38-27b-int4-mtp-cuda` (short window, flag OFF = origin behavior): + - MTP (`--steady`): 14.45 tok/s, acceptance 78.9%, `cuda_graph enabled=true captures=16 replays=0 fallbacks=0 invalidations=99`, no NaN, tokens_per_verify_step 2.58. + - Greedy (plain): 55.39 tok/s, `cuda_graph captures=2 replays=92 fallbacks=0 invalidations=1` (healthy, matches origin ~55.91). +- **No MTP speedup number is reported — none exists** (no config safely produced replays>0). MTP stays ~3.8× slower than greedy under eager verify. diff --git a/crates/onnx-genai-engine/src/native_decode/backend.rs b/crates/onnx-genai-engine/src/native_decode/backend.rs index 84f7cc0ae9..30f6fdc11b 100644 --- a/crates/onnx-genai-engine/src/native_decode/backend.rs +++ b/crates/onnx-genai-engine/src/native_decode/backend.rs @@ -137,11 +137,22 @@ impl NativeDecodeSession { if let Some(state) = &mut self.cuda { // Option (b) default: invalidate the captured decode graph before the // KV roll-back (the eager verify path captures nothing, and the plain - // M=1 path re-warms cleanly). Option (c) (dormant until WP4) retains - // 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 { + // M=1 path re-warms cleanly). Two dormant seams (both OFF by default) + // would retain the graph across a contents-only rewind instead — the + // rewind only zeros the mask tail + truncates the KV logical length, + // leaving every binding's physical_shape/device_ptr fixed, so the + // captured M=1 graph's replay signature stays valid: `retain_graph_on_rewind` + // (option (c) padded verify capture) and `retain_decode_graph_across_spec` + // (spec-decode retention). Retention on rewind alone is capture-safe, + // but NOT sufficient for a speedup: the eager M>1 verify forward tears + // the graph down every step regardless, and retaining across BOTH + // sites is capture-unsafe until the verify workspace is pinned (see + // the verify site + decision note). A full reset to `target_len == 0` + // (between generations) always invalidates so a stale graph never + // leaks into the next generation. + let retain = state.retain_graph_on_rewind + || (state.retain_decode_graph_across_spec && target_len != 0); + if !retain { state.invalidate_graph(&mut self.session)?; } state.rewind(target_len)?; diff --git a/crates/onnx-genai-engine/src/native_decode/cuda.rs b/crates/onnx-genai-engine/src/native_decode/cuda.rs index 0b404b802c..a23fc81e52 100644 --- a/crates/onnx-genai-engine/src/native_decode/cuda.rs +++ b/crates/onnx-genai-engine/src/native_decode/cuda.rs @@ -731,6 +731,26 @@ pub(crate) struct DecodeCudaState { /// captured graph already tolerates on the M=1 replay path. Kept dormant /// (default `false`) until WP4 graduates verify to the captured path. pub(crate) retain_graph_on_rewind: bool, + /// **Dormant seam (default `false`).** When `true`, the captured M=1 decode + /// graph would be retained across a speculative verify+commit cycle instead + /// of being torn down twice per step (once by the eager M>1 verify forward, + /// once by the commit rewind), the two invalidations that pin MTP at + /// `replays=0` (empirically 30 verify + 21 rewind invalidations over 10 + /// verify steps). + /// + /// Kept OFF because enabling it is **not capture-safe** as-is: the eager M>1 + /// verify reserves a larger StepScoped `step_workspace` that is freed after + /// the run, so a later M=1 replay reads the captured graph's now-stale + /// workspace pointer and yields non-finite logits (GPU-verified; the finite + /// guard catches it — no silent corruption). Retaining across the rewind + /// alone, or across the verify alone, is safe but useless (the other site + /// still tears the graph down, so `replays` stays 0); retaining across both — + /// the only config that actually replays — is what corrupts. A real speedup + /// requires the M=K verify itself captured into a fixed-shape replayable + /// graph with a pinned workspace (option-c padded verify capture), since the + /// eager verify otherwise pays the full per-op launch overhead that graphed + /// greedy avoids. See the decision note for the GPU evidence and plan. + pub(crate) retain_decode_graph_across_spec: bool, /// Dormant option (c) scaffolding: the fixed query-row capacity (M=maxK) a /// padded single-capture verify graph would be captured at. `None` today — /// the eager verify path (option (b)) captures nothing. Set only by the @@ -1251,7 +1271,21 @@ impl NativeDecodeSession { .cuda .as_mut() .context("CUDA decode state is not initialized")?; - state.invalidate_graph(&mut self.session)?; + // An eager M>1 verify/prefill forward tears down the captured M=1 decode + // graph here so the plain M=1 hot path re-warms cleanly. The dormant + // `retain_decode_graph_across_spec` seam (default OFF) would skip this to + // let the base-decode graph survive a spec verify+commit cycle — but that + // is NOT capture-safe yet: this forward reserves a larger StepScoped + // `step_workspace` that is freed (`release_step_workspace`) after the run, + // so a subsequent M=1 replay reads the captured graph's now-stale + // workspace pointer and produces non-finite logits (GPU-verified; the + // finite-check guard catches it). Graduating this to a real speedup needs + // the M=K verify itself captured into a fixed-shape replayable graph with + // a pinned workspace (option-c padded verify capture), not just retaining + // the M=1 graph. Kept as a documented seam; see the decision note. + if !state.retain_decode_graph_across_spec { + state.invalidate_graph(&mut self.session)?; + } // Auxiliary graph outputs (e.g. the MTP hidden-state seed // `hidden_states.63`) get a persistent device binding whose symbolic // query-seq axis is collapsed to `1` for the captured decode step @@ -4459,6 +4493,11 @@ impl DecodeCudaState { graph_fallback_report: None, auxiliary_bind_declines: declined_auxiliary, retain_graph_on_rewind: false, + // Dormant seam (default off): retaining the M=1 graph across a spec + // verify+commit cycle is not capture-safe until the M=K verify is + // itself captured with a pinned workspace. See the field docs and the + // decision note for the GPU evidence. + retain_decode_graph_across_spec: false, #[cfg(test)] padded_query_capacity: None, device_token_loop_k, @@ -5469,6 +5508,20 @@ impl DecodeCudaState { self.retain_graph_on_rewind = retain; } + /// Toggle the dormant spec-decode graph-retention seam + /// (`retain_decode_graph_across_spec`). Kept off in production because it is + /// not capture-safe until the M=K verify is captured with a pinned workspace; + /// exposed for the option-c work and its tests. See the field docs. + #[cfg(test)] + pub(crate) fn set_retain_decode_graph_across_spec(&mut self, retain: bool) { + self.retain_decode_graph_across_spec = retain; + } + + #[cfg(test)] + pub(crate) fn retain_decode_graph_across_spec(&self) -> bool { + self.retain_decode_graph_across_spec + } + /// Fixed query-row capacity (M=maxK) of the dormant padded verify capture, or /// `None` while the eager (option (b)) verify path is in force. #[cfg(test)]