perf(mtp): per-slot executor host capture state (foundation for two-slot verify replay)

.squad/decisions/inbox/gaff-mtp-perslot-capture-shared-buffer-arena-blocker.md

@@ -0,0 +1,63 @@
+### 2026-08-21: Per-slot executor host capture state landed; verify replays>0 blocked by SHARED interior buffer arena
+
+**By:** Gaff
+
+**What:**
+Made the `Executor`'s host-side CUDA-graph capture state per-slot so the M=1 `Primary` decode
+graph and the M=k+1 `Verify` speculative-verify graph can coexist on ONE executor (branch
+`squad/mtp-perslot-capture` off `6923a016b`, PR to follow). Concretely:
+- `DeviceGraphSlot::{COUNT=2, index()}` in ep-api (Primary=0, Verify=1).
+- Extracted the 8 capture fields (`device_graph_signature`, `capture_schedule`,
+  `capture_segmentation`, `capture_cf_shapes`, `capture_warm_signature`, `capture_warm_shapes`,
+  `capture_warm_seeded`, `capture_quarantine_ops`) into `SlotCaptureState`; `Executor` now holds
+  `slot_capture: [SlotCaptureState; 2]` with `cap()`/`cap_mut()` accessors indexed by
+  `graph_slot`. `buffer_shapes`/`if_last_predicate`/`control_flow_output_values` stay shared
+  per-run scratch (per coordinator scope).
+- `set_graph_slot` is now a pure retarget (no longer resets the other slot).
+- Dropped the `enable_decode_inline()` gate on `configure_verify_capture`; removed the permanent
+  `set_main_exec_graph_slot(Verify)` (main exec stays on Primary for M=1 decode; `run_verify_captured`
+  flips to Verify around the verify forward and back to Primary after).
+- Removed the dead `set_retain_decode_graph_across_spec`/`retain_decode_graph_across_spec` accessors.
+- Added regression test `set_graph_slot_is_non_resetting_and_per_slot_isolated` (greedy-uses-only-Primary
+  invariant; Primary marker survives a Verify round-trip, no cross-slot bleed).
+
+**Validation:**
+- Engine lib suite `--features native-backend`: **579 passed / 0 failed** (greedy inert, oracle
+  `native_verify_logits_require_restored_recurrent_state` green).
+- ep-cuda `graph::tests` `--features cuda,cuda-13000`: **8/8** (incl.
+  `primary_and_verify_graph_slots_are_independent`).
+- Session crate tests green. Release CUDA bench builds.
+- GPU (H200 ordinal 5, CUDA_VISIBLE_DEVICES=5, int4 block-32, ORT 1.28 cuda13, build
+  `--features bench-native,native-cuda,cuda-13000`, harness `profile_native --steady --tokens 128
+  --warmups 3 --runs 1`, median-of-5): MTP clean, **no NaN**, coherent output, **83.3% acceptance**
+  (2.67 tok/verify-step), **median 14.67 tok/s**.
+  - `cuda_graph:` (Primary) captures=64 replays=0 fallbacks=0 invalidations=195
+  - `cuda_graph_verify:` (Verify) captures=0 replays=0 **fallbacks=1** invalidations=0
+  No regression vs gated #1652 (~14.6 tok/s); verify machinery now ARMS+DECLINES gracefully
+  instead of being inert.
+
+**Why / remaining gap (GPU-proven):**
+Per-slot host capture state is **necessary but NOT sufficient** for verify replays>0 on this
+sibling-less GDN artifact. The Verify capture DECLINES because the interior device
+`buffers`/`buffer_shapes` arena is SHARED (per-run scratch, deliberately not per-slot). The M=1
+decode JIT-sizes interior values to `[1,1]`; during the M=2 verify capture pass the interior
+buffer for `model/Slice_node_10` is still `[1,1]`, so the kernel errors
+`output shape [1,2], expected [1,1]` → capture declines → graceful eager fallback (correct, no replay).
+Even if capture succeeded, the next M=1 decode would resize those shared interior buffers,
+clobbering the M=2 graph's baked interior pointers.
+
+**Exact next lever:** give the M=k+1 verify its OWN device buffer arena (a verify-dedicated
+executor / buffer arena, analogous to the decode-inline sibling but for M=K) so its interior
+JIT-sized scratch is independent of the interleaved M=1 decode. That is the true precondition for
+verify replays>0 on this artifact — a larger architectural change than the per-slot host-state
+refactor.
+
+**Blocker B (separate/pre-existing):** Primary M=1 decode never replays in MTP mode
+(`replays=0 invalidations=195`) even before this change — the interleaved
+`commit_recurrent_state_to_accepted` (snapshot→restore→re-advance accepted tokens, variable width
+M=1 or 2) runs on the Primary slot sharing the same persistent decode bindings, resizing them and
+invalidating the M=1 graph each spec step.
+
+**Honest status:** No net speedup yet (14.67 tok/s MTP vs ~56 greedy). The per-slot foundation is
+done, correct, tested, and greedy byte-identical; replays>0 is blocked by the shared interior
+buffer arena + blocker B, both requiring the verify-dedicated-arena/executor lever.

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

@@ -1353,27 +1353,21 @@ impl NativeDecodeSession {
             .as_mut()
             .context("CUDA decode state is not initialized")?;
         // 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.
+        // graph here so the plain M=1 hot path re-warms cleanly — EXCEPT once
+        // option-c verify capture is active, where each graph lives in its own
+        // per-slot host capture state and must not be disturbed by the other.
         if state.verify_width.is_some() {
-            // Option-c verify capture is active: the M=1 decode graph lives in
-            // the Primary slot on the decode-inline sibling, and the captured
-            // M=K verify graph lives in the main exec's Verify slot (owned by
-            // `run_verify_captured`). This eager main-exec forward is either the
-            // one-time prefill or a tail verify at width != M; in neither case
-            // may it invalidate the sibling's decode graph (that is what pinned
-            // MTP at replays=0 before option-c). A plain eager run leaves the
-            // installed Verify-slot graph untouched, so the next full-width step
-            // still replays it.
+            // Option-c verify capture is active. With per-slot host capture state
+            // the main executor holds the M=1 decode graph in its `Primary` slot
+            // and the fixed-M verify graph in its `Verify` slot; `Primary` is the
+            // current slot on this eager path (the base decode / prefill runs
+            // here), so we must NOT invalidate — that would tear down the M=1
+            // decode graph every verify step, which is exactly what pinned MTP at
+            // replays=0 before option-c. The captured verify graph is owned by
+            // `run_verify_captured` on the `Verify` slot and is untouched here.
+            // This eager forward is either the one-time prefill or a tail verify
+            // at width != M; a plain eager run leaves both installed graphs intact
+            // so the next step still replays the matching one.
         } else if !state.retain_decode_graph_across_spec {
             state.invalidate_graph(&mut self.session)?;
         }
@@ -1490,28 +1484,17 @@ impl NativeDecodeSession {
         if width < 2 || self.has_eager_step_inputs() || !self.has_recurrent_state() {
             return Ok(());
         }
-        // Capturing the M=K verify into the Verify graph slot is only sound when
-        // the interleaved M=1 base decode / commit re-advance run on a DIFFERENT
-        // executor+slot — the decode-inline sibling (Primary slot). This
-        // executor keeps a single host capture signature; if the M=1 decode also
-        // runs on this (main) executor it clobbers the verify signature every
-        // step (GPU-confirmed: `run_one_token` captures M=1 into the Verify slot
-        // between verify capture and replay), pinning verify replays at 0. Two
-        // coexisting replayable graphs on one executor need per-slot host capture
-        // state (device_graph_signature + warm-seed/schedule), a large executor
-        // refactor tracked in the decision note. Until then, only arm verify
-        // capture when a decode-inline sibling exists to own the M=1 decode; on
-        // models whose recurrent Scan is not single-trip-inlineable (e.g. this
-        // artifact) the sibling cannot be built, so leave verify eager (correct,
-        // token-identical, no capture churn — the pre-option-c behavior).
-        if !self.session.enable_decode_inline()? {
-            tracing::info!(
-                "native MTP verify capture stays inert: no decode-inline sibling to own the M=1 \
-                 decode, so a single-executor Verify-slot capture would be clobbered by the \
-                 interleaved M=1 decode capture (needs per-slot host capture state)"
-            );
-            return Ok(());
-        }
+        // Capturing the M=K verify into the `Verify` graph slot is now sound even
+        // when the interleaved M=1 base decode / commit re-advance run on the
+        // SAME (main) executor. The executor holds per-slot host capture state
+        // (device_graph_signature + schedule + warm-shape/quarantine, indexed by
+        // graph_slot), so an M=1 decode captured into the `Primary` slot no
+        // longer clobbers the `Verify` slot's signature between the verify's
+        // capture and its replay. Both graphs coexist and each replays by shape
+        // key. This lifts the earlier `enable_decode_inline()` gate: on models
+        // whose recurrent Scan is not single-trip-inlineable (no decode-inline
+        // sibling — e.g. this artifact) the M=1 decode simply runs on the main
+        // executor's `Primary` slot while the verify runs on its `Verify` slot.
         // Gather the padded binding plan from the live decode bindings (immutable
         // borrow), then allocate (session borrow), then store (state borrow).
         struct VerifyBindingPlan {
@@ -1629,13 +1612,13 @@ impl NativeDecodeSession {
         state.verify_width = Some(width);
         state.verify_graph_phase = DecodeCudaGraphPhase::NeedsWarmup;
 
-        // Route the main executor to the Verify slot and pin its StepScoped
-        // workspace at the M=K peak: the captured verify graph then replays
-        // against a stable scratch pointer even though the M=1 decode step (on
-        // the inline sibling, Primary slot) reserves a smaller scratch between
-        // verify steps (#1647). Both are inert for every non-verify path.
-        self.session
-            .set_main_exec_graph_slot(DeviceGraphSlot::Verify)?;
+        // Pin the main executor's StepScoped workspace at the M=K peak so the
+        // captured verify graph replays against a stable scratch pointer even
+        // though the interleaved M=1 decode step reserves a smaller scratch
+        // (#1647). The graph slot is NOT set here: the main executor stays on
+        // `Primary` for M=1 decode and `run_verify_captured` flips it to `Verify`
+        // only around the verify forward, then back — safe because the executor's
+        // per-slot host capture state keeps both graphs' signatures independent.
         self.session.set_main_exec_pin_step_workspace(true);
         Ok(())
     }
@@ -1700,14 +1683,21 @@ impl NativeDecodeSession {
             }
         }
 
-        // Swap the padded bindings into the persistent vector, run the verify
-        // graph phase (which reads the logits while swapped in), then swap the
-        // M=1 decode bindings back — unconditionally, even on error.
+        // Swap the padded bindings into the persistent vector, retarget the main
+        // executor to the Verify slot, run the verify graph phase (which reads
+        // the logits while swapped in), then restore the Primary slot + swap the
+        // M=1 decode bindings back — unconditionally, even on error. The slot
+        // flip is a pure retarget (per-slot host capture state), so the M=1
+        // decode graph in Primary is untouched and still replays next step.
         self.cuda
             .as_mut()
             .context("CUDA decode state is not initialized")?
             .swap_verify_bindings();
+        self.session
+            .set_main_exec_graph_slot(DeviceGraphSlot::Verify)?;
         let outcome = self.run_verify_graph_phase(m);
+        self.session
+            .set_main_exec_graph_slot(DeviceGraphSlot::Primary)?;
         self.cuda
             .as_mut()
             .context("CUDA decode state is not initialized")?
@@ -1795,19 +1785,17 @@ impl NativeDecodeSession {
                     }
                 }
                 DecodeCudaGraphPhase::Ready => {
-                    // A single main executor drives BOTH the M=1 decode (Primary
-                    // slot) and this M=K verify (Verify slot) when the model has
-                    // no decode-inline sibling. The executor keeps ONE host
-                    // capture signature, so the interleaved M=1 decode capture
-                    // clobbers the verify signature and this replay sees a
-                    // mismatch (hard `Err` from `replay_device_graph`, which also
-                    // resets the slot). Two coexisting replayable graphs on one
-                    // executor need per-slot host capture state (see the decision
-                    // note). Until that lands, degrade gracefully instead of
-                    // aborting the generation: count the invalidation, run the
-                    // forward eagerly to produce this step's logits, and after a
-                    // couple of clobbers stop re-capturing (permanent eager verify
-                    // = the pre-option-c behavior, correct + no capture churn).
+                    // The main executor now holds PER-SLOT host capture state, so
+                    // the interleaved M=1 decode (Primary slot) no longer clobbers
+                    // this M=K verify graph (Verify slot): each slot keeps its own
+                    // signature/schedule and replays independently. A replay can
+                    // still legitimately retire the graph — e.g. the pinned
+                    // StepScoped workspace grew on the first verify and moved its
+                    // scratch pointer, or a control-flow branch flip — so we keep
+                    // the graceful path: count the invalidation, run eagerly to
+                    // still produce this step's logits, re-warm once (self-
+                    // stabilizing after the workspace settles at the M=K peak),
+                    // and only latch to permanent-eager if churn persists.
                     match self.session.replay_device_graph(&mut state.bindings[..]) {
                         Ok(true) => {
                             state.verify_graph_replays += 1;
@@ -6027,26 +6015,6 @@ 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.
-    ///
-    /// `allow(dead_code)`: this pair is the seam itself, landed ahead of the
-    /// WP4 tests that drive it. The sibling `#[cfg(test)]` accessors above and
-    /// below are already called, so only these two are unreachable today.
-    #[cfg(test)]
-    #[allow(dead_code)]
-    pub(crate) fn set_retain_decode_graph_across_spec(&mut self, retain: bool) {
-        self.retain_decode_graph_across_spec = retain;
-    }
-
-    #[cfg(test)]
-    #[allow(dead_code)]
-    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)]

crates/onnx-runtime-ep-api/src/provider.rs

@@ -69,6 +69,24 @@ pub enum DeviceGraphSlot {
     Verify,
 }
 
+impl DeviceGraphSlot {
+    /// Number of distinct captured-graph slots. Used to size per-slot host
+    /// capture-state arrays on the executor so `Primary` (M=1 decode) and
+    /// `Verify` (M=k+1) graphs can coexist without clobbering each other.
+    pub const COUNT: usize = 2;
+
+    /// Dense array index for this slot (`Primary` = 0, `Verify` = 1). `Primary`
+    /// is index 0 so the historical single-slot code path — which only ever
+    /// touches `Primary` — maps to slot 0 and stays byte-identical.
+    #[inline]
+    pub const fn index(self) -> usize {
+        match self {
+            DeviceGraphSlot::Primary => 0,
+            DeviceGraphSlot::Verify => 1,
+        }
+    }
+}
+
 /// Opaque, namespaced configuration passed to [`ExecutionProvider::initialize`].
 #[derive(Clone, Debug, Default)]
 pub struct EpConfig {