perf(mtp): approach-B (skip re-advance on full accept) + device-to-device recurrent snapshot — 20.68→34.34 tok/s, token-identical

.squad/decisions/inbox/gaff-mtp-approach-b-snapshot-d2d.md

@@ -0,0 +1,80 @@
+### 2026-08-21: MTP approach-B (skip re-advance on full accept) + device-to-device recurrent snapshot — 20.68 → 34.34 tok/s, token-identical, both slots replay
+
+**By:** Gaff
+
+**What:**
+Two committed, GPU-validated, token-identical speedups on the real Qwen3.8-27B
+int4 hybrid (GDN+GQA) MTP self-spec decode (k=1, 83.3% accept, 2.67 tok/step,
+128-token run on an idle H200, median-of-5). Branch `squad/mtp-approach-b` off
+origin/main `66e19c2b5`.
+
+1. **approach-B — eliminate the redundant re-advance (commit `6c7f71650`).**
+   When ALL draft tokens are accepted (~67% of steps at 83.3% accept, k=1), the
+   eager M=K verify forward has ALREADY advanced KV + GDN/conv recurrent state to
+   exactly `base+accepted` (== committed length). The snapshot→restore→re-advance
+   commit is then pure redundancy, so it is skipped; only partial-accept keeps
+   the snapshot→rewind→restore→re-advance path (intermediate per-position
+   recurrent states are not exposed by the GDN kernel, so partial accept still
+   rebuilds by re-running the accepted prefix). Result: **20.68 → 26.73 tok/s
+   (+29%)**; Primary graph replays 504 → 248 (64 redundant re-advance M=1
+   forwards/generation removed).
+
+2. **device-to-device recurrent snapshot (commit `2f48f801c`).** The remaining
+   snapshot copied every fixed-state (GDN SSM + conv1d) binding device→HOST every
+   spec step via `read_bytes_range` (~10ms/step PCIe round-trip), unused on the
+   full-accept steps. Replaced with a stream-ordered **device→device** snapshot
+   into a per-binding device scratch arena:
+   - `ExecutionProvider::copy_device_to_device` (default errors; CUDA overrides
+     with `cuMemcpyDtoDAsync` on the EP stream → snapshot is ordered ahead of the
+     verify forward that overwrites the state, no host sync).
+   - `DeviceIoBinding::snapshot_device_into` / `restore_device_from` + an
+     `allocator()` accessor; `DeviceBuffer` re-exported from onnx-runtime-session.
+   - `DecodeCudaState.fixed_state_snapshot_scratch` allocated lazily and sized
+     from `fixed_state_binding_range` (no hardcoded layer/dim numbers);
+     `snapshot_fixed_states_device` / `restore_fixed_states_device` replace the
+     host-bytes variants. `RecurrentStateSnapshot` carries a `device_scratch`
+     flag; the CPU host-past path is unchanged. Result: **26.73 → 34.34 tok/s**.
+
+Also fixed a pre-existing red on main (commit `c799a835b`): `InferenceSession`
+gained `verify_exec` in #1658 but the two `executor/tests.rs` initializers were
+never updated, breaking `cargo test -p onnx-runtime-session --features cuda`.
+
+**Per-step target-forward budget (128-tok run, 48 verify steps):**
+- Before approach-B: 48 base(M=1) + 80 re-advance(M=1) + 48 verify(M=2)  [Primary replays=504].
+- After approach-B:  48 base(M=1) + 16 re-advance(M=1) + 48 verify(M=2)  [Primary replays=248].
+- D2D snapshot does not change the forward count; it removes the per-step PCIe
+  D2H transfer (state staging) that was the dominant remaining wall-clock cost.
+
+**Validation (all green):**
+- Both slots replay, fallbacks=0: `cuda_graph: captures=4 replays=248 fallbacks=0 invalidations=3`;
+  `cuda_graph_verify: captures=4 replays=184 fallbacks=0 invalidations=0`.
+- MTP token-identical to baseline: `generated_token_ids` md5 `be7ed565` across
+  median-of-5, no NaN (finite-guard clean).
+- Engine lib `--features native-backend` **579 passed / 0 failed** (greedy inert
+  — the snapshot machinery only runs under spec-decode with recurrent state).
+- ep-cuda `graph::tests` **8/8**; session `--features cuda` green (191 + suites).
+- Median-of-5 A/B (H200, `CUDA_VISIBLE_DEVICES=5`, batch=1, tokens=128, warmups=3,
+  int4 block-32, ORT 1.28 cuda13, `--release --features bench-native,native-cuda,cuda-13000`,
+  `profile_native --steady`): **34.34 tok/s** (34.29–34.36, very tight).
+  Progression: 20.68 (#1658) → 26.73 (approach-B) → 34.34 (D2D snapshot). Total +66%.
+
+**Why:**
+On a launch-bound decode, MTP only beats greedy if per accepted step it costs
+fewer target forwards than tokens produced. approach-B removes lot (c) — the
+re-advance — on the 67% full-accept path (correct because the M=K verify
+post-state IS byte-equivalent to K sequential M=1 advances for the accepted
+prefix; proven by md5 equality). The D2D snapshot removes the last big per-step
+overhead (the PCIe state stage) with an exact device copy. Both are byte-identical
+and inert for greedy, so they carry no correctness risk.
+
+**Remaining gap to ~56 greedy (honest):** still below greedy. The structural
+lever left is **fusing out the separate M=1 base decode** so per-step cost = a
+single M=2 verify forward. It is NOT trivially fusible: circular dependency — the
+MTP head needs the hidden state AFTER the last committed token to draft, but that
+token/bonus is only processed by the base decode. EAGLE-style seeding the MTP
+head from the previous verify's frontier hidden would break the cycle and stay
+token-identical (the target verify corrects every draft, so drafts only affect
+acceptance RATE, not correctness), but the speed benefit hinges on preserving the
+83.3% acceptance through the changed seeding — needs MTP-head reseeding-semantics
+verification. Deferred as the next (riskier, bigger) turn; approach-B + D2D are
+landed and durable regardless.

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

@@ -683,6 +683,13 @@ pub(crate) struct DecodeCudaState {
     /// from the declared `init: zeros` state (see `rewind`). Empty for pure-KV
     /// decoders.
     pub(crate) fixed_state_binding_range: std::ops::Range<usize>,
+    /// Device-resident scratch buffers for the speculative recurrent-state
+    /// snapshot, keyed by fixed-state binding index. Allocated lazily on the
+    /// first device snapshot and reused thereafter (state shapes are fixed), so
+    /// snapshotting the destructive GDN/conv state costs one device→device copy
+    /// per binding instead of a PCIe round-trip through host memory. Empty until
+    /// the CUDA device-snapshot path first runs; inert for greedy decode.
+    fixed_state_snapshot_scratch: Vec<(usize, DeviceBuffer)>,
     pub(crate) auxiliary_binding_range: std::ops::Range<usize>,
     pub(crate) input_ids_binding: usize,
     pub(crate) position_ids_binding: Option<usize>,
@@ -4965,6 +4972,7 @@ impl DecodeCudaState {
             base_binding_count,
             kv_binding_range: kv_start..kv_end,
             fixed_state_binding_range: kv_end..fixed_state_end,
+            fixed_state_snapshot_scratch: Vec::new(),
             auxiliary_binding_range: auxiliary_start..auxiliary_end,
             input_ids_binding,
             position_ids_binding,
@@ -5491,57 +5499,94 @@ impl DecodeCudaState {
         self.set_logical_len(target_len)
     }
 
-    /// Copy every fixed-size recurrent/conv binding off the device into host
-    /// buffers keyed by binding index, capturing the state as of the last
-    /// committed token. Paired with [`Self::restore_fixed_states`] for
-    /// speculative recurrent-state commit; the bindings are the existing
-    /// `fixed_state_binding_range`, so no layer/dim geometry is hardcoded.
-    pub(crate) fn snapshot_fixed_states(&mut self) -> anyhow::Result<Vec<(usize, Vec<u8>)>> {
-        let mut snapshot = Vec::with_capacity(self.fixed_state_binding_range.len());
+    /// Ensure a device scratch buffer exists for every fixed-state binding,
+    /// sized to that binding's current byte length. Allocated once (state
+    /// shapes are fixed) via the binding's own execution provider, then reused.
+    /// Returns without allocating when there are no fixed-state bindings.
+    fn ensure_fixed_state_snapshot_scratch(&mut self) -> anyhow::Result<()> {
+        if self.fixed_state_snapshot_scratch.len() == self.fixed_state_binding_range.len() {
+            return Ok(());
+        }
+        // 256-byte alignment matches CUDA's native allocation granularity and is
+        // a power of two as `allocate` requires; the copy correctness does not
+        // depend on the base alignment, only on the byte length.
+        const SCRATCH_ALIGN: usize = 256;
+        let mut scratch = Vec::with_capacity(self.fixed_state_binding_range.len());
         for index in self.fixed_state_binding_range.clone() {
-            let binding = &mut self.bindings[index];
+            let binding = &self.bindings[index];
             let bytes = checked_shape_bytes(binding.physical_shape(), binding.dtype)
                 .with_context(|| {
                     format!(
-                        "fixed CUDA decoder state snapshot size overflow for binding {index} shape {:?}",
+                        "fixed CUDA decoder state scratch size overflow for binding {index} shape {:?}",
                         binding.physical_shape()
                     )
                 })?;
-            let host = binding
-                .read_bytes_range(0, bytes)
-                .with_context(|| format!("read device recurrent/conv state for binding {index}"))?;
-            snapshot.push((index, host));
+            let buffer = binding
+                .allocator()
+                .allocate(bytes.max(1), SCRATCH_ALIGN)
+                .with_context(|| format!("allocate device snapshot scratch for binding {index}"))?;
+            scratch.push((index, buffer));
         }
-        Ok(snapshot)
+        self.fixed_state_snapshot_scratch = scratch;
+        Ok(())
     }
 
-    /// Write a [`Self::snapshot_fixed_states`] capture back into the device
-    /// recurrent/conv bindings. Restoring bytes leaves every binding's shape
-    /// unchanged, so this never invalidates a captured decode graph.
-    pub(crate) fn restore_fixed_states(
-        &mut self,
-        snapshot: &[(usize, Vec<u8>)],
-    ) -> anyhow::Result<()> {
-        for (index, host) in snapshot {
-            let binding = self.bindings.get_mut(*index).with_context(|| {
-                format!("recurrent snapshot names out-of-range fixed-state binding {index}")
-            })?;
+    /// Device-resident analogue of [`Self::snapshot_fixed_states`]: stage every
+    /// fixed recurrent/conv binding into its device scratch buffer with a
+    /// stream-ordered device→device copy, avoiding the PCIe round-trip through
+    /// host memory. The snapshot is ordered on the EP stream ahead of the verify
+    /// forward that overwrites the state, so no host synchronization is needed.
+    /// Paired with [`Self::restore_fixed_states_device`].
+    pub(crate) fn snapshot_fixed_states_device(&mut self) -> anyhow::Result<()> {
+        self.ensure_fixed_state_snapshot_scratch()?;
+        for slot in 0..self.fixed_state_snapshot_scratch.len() {
+            let (index, ref mut scratch) = self.fixed_state_snapshot_scratch[slot];
+            let binding = &self.bindings[index];
             let bytes = checked_shape_bytes(binding.physical_shape(), binding.dtype)
                 .with_context(|| {
                     format!(
-                        "fixed CUDA decoder state restore size overflow for binding {index} shape {:?}",
+                        "fixed CUDA decoder state device snapshot size overflow for binding {index} shape {:?}",
                         binding.physical_shape()
                     )
                 })?;
-            if host.len() != bytes {
-                bail!(
-                    "recurrent snapshot for binding {index} has {} bytes but the binding needs {bytes}",
-                    host.len()
-                );
-            }
-            binding.write_bytes(0, host).with_context(|| {
-                format!("restore device recurrent/conv state for binding {index}")
+            binding
+                .snapshot_device_into(scratch, bytes)
+                .with_context(|| {
+                    format!("device snapshot of recurrent/conv state for binding {index}")
+                })?;
+        }
+        Ok(())
+    }
+
+    /// Restore a [`Self::snapshot_fixed_states_device`] capture back into the
+    /// device recurrent/conv bindings with a stream-ordered device→device copy.
+    /// Leaves every binding's shape unchanged, so it never invalidates a
+    /// captured decode graph.
+    pub(crate) fn restore_fixed_states_device(&mut self) -> anyhow::Result<()> {
+        if self.fixed_state_snapshot_scratch.len() != self.fixed_state_binding_range.len() {
+            bail!(
+                "device recurrent restore requested before a matching device snapshot (have {} scratch buffers, need {})",
+                self.fixed_state_snapshot_scratch.len(),
+                self.fixed_state_binding_range.len()
+            );
+        }
+        for slot in 0..self.fixed_state_snapshot_scratch.len() {
+            let (index, ref scratch) = self.fixed_state_snapshot_scratch[slot];
+            let binding = self.bindings.get_mut(index).with_context(|| {
+                format!("recurrent device snapshot names out-of-range fixed-state binding {index}")
             })?;
+            let bytes = checked_shape_bytes(binding.physical_shape(), binding.dtype)
+                .with_context(|| {
+                    format!(
+                        "fixed CUDA decoder state device restore size overflow for binding {index} shape {:?}",
+                        binding.physical_shape()
+                    )
+                })?;
+            binding
+                .restore_device_from(scratch, bytes)
+                .with_context(|| {
+                    format!("device restore of recurrent/conv state for binding {index}")
+                })?;
         }
         Ok(())
     }

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

@@ -11,7 +11,8 @@ use onnx_genai_ort::Tokenizer;
 use onnx_runtime_ir::{DataType, DeviceType, Dim, SymbolId};
 use onnx_runtime_session::{
     CaptureDeclineReport, DecodePrecision, DeviceAllocationCounts, DeviceBindingTransferStats,
-    DeviceGraphCaptureResult, DeviceIoBinding, DevicePreference, InferenceSession, Tensor,
+    DeviceBuffer, DeviceGraphCaptureResult, DeviceIoBinding, DevicePreference, InferenceSession,
+    Tensor,
 };
 use onnx_runtime_tracer::{Args, TraceContext, capture_rejected};
 use std::collections::{HashMap, HashSet};
@@ -181,14 +182,20 @@ impl NativePastSnapshot {
 /// Unlike [`NativePastSnapshot`] (which clones the whole host past for prefix
 /// caching), this captures *only* the recurrent/conv bindings and works on both
 /// the host past path and the CUDA fixed-state bindings. Exactly one of `host`
-/// (rank-carrying host tensors keyed by past-input name) or `device` (raw
-/// device bytes keyed by fixed-state binding index) is populated, depending on
-/// which decode backend produced it. `len` is the committed length the snapshot
-/// was taken at, asserted against the commit's `base_len`.
+/// (rank-carrying host tensors keyed by past-input name) or `device_scratch`
+/// (the CUDA fixed-state bindings staged into device scratch) indicates where
+/// the captured state lives, depending on which decode backend produced it.
+/// `len` is the committed length the snapshot was taken at, asserted against the
+/// commit's `base_len`.
 pub(crate) struct RecurrentStateSnapshot {
     len: usize,
     host: Option<HashMap<String, Tensor>>,
-    device: Option<Vec<(usize, Vec<u8>)>>,
+    /// True when the CUDA fixed-state bindings were staged into the session's
+    /// device scratch buffers (a stream-ordered device→device snapshot). The
+    /// bytes live in the CUDA decode state rather than in this handle, so a
+    /// restore copies them back from there. Only one such snapshot is live at a
+    /// time (per speculative step), matching the single scratch arena.
+    device_scratch: bool,
 }
 
 impl RecurrentStateSnapshot {
@@ -503,11 +510,11 @@ impl NativeDecodeSession {
         }
         let len = self.current_len;
         if let Some(cuda) = self.cuda.as_mut() {
-            let device = cuda.snapshot_fixed_states()?;
+            cuda.snapshot_fixed_states_device()?;
             return Ok(RecurrentStateSnapshot {
                 len,
                 host: None,
-                device: Some(device),
+                device_scratch: true,
             });
         }
         if self.cpu_kv.is_some() {
@@ -532,7 +539,7 @@ impl NativeDecodeSession {
         Ok(RecurrentStateSnapshot {
             len,
             host: Some(host),
-            device: None,
+            device_scratch: false,
         })
     }
 
@@ -544,11 +551,11 @@ impl NativeDecodeSession {
         &mut self,
         snapshot: &RecurrentStateSnapshot,
     ) -> anyhow::Result<()> {
-        if let Some(device) = &snapshot.device {
+        if snapshot.device_scratch {
             let cuda = self.cuda.as_mut().context(
                 "recurrent snapshot targets the CUDA fixed-state bindings but this session has no CUDA state",
             )?;
-            cuda.restore_fixed_states(device)?;
+            cuda.restore_fixed_states_device()?;
             return Ok(());
         }
         if let Some(host) = &snapshot.host {

crates/onnx-genai-engine/src/native_speculative.rs

@@ -428,13 +428,34 @@ impl<'a> NativeSpeculativeDriver<'a> {
             // destructive GDN/conv state stranded at `base + K`; commit it to
             // exactly the accepted prefix instead (snapshot restore + accepted-
             // token re-advance), which also performs the KV rewind. See #1598.
-            match recurrent_snapshot.as_ref() {
-                Some(snapshot) => self.session.commit_recurrent_state_to_accepted(
-                    snapshot,
-                    base,
-                    &draft[..accepted],
-                )?,
-                None => self.session.rewind(base + accepted)?,
+            //
+            // Approach-B fast path (full accept): when EVERY draft token is
+            // accepted, the eager verify forward already advanced BOTH the KV
+            // (`base + K`) and the destructive GDN/conv recurrent state by exactly
+            // those `K == accepted` tokens, and the committed length equals the
+            // current length — so there is nothing to roll back. Skip the KV
+            // rewind, the snapshot restore, AND the accepted-token re-advance
+            // entirely: the verify's own post-state IS the committed state. This
+            // removes the redundant re-advance target forwards on the majority of
+            // steps (every multi-token accept). The snapshot is still taken before
+            // the verify because acceptance is only known after it; it is simply
+            // unused on this path. Only a PARTIAL accept (`accepted < draft.len()`)
+            // needs the snapshot→restore→re-advance rebuild to a shorter prefix.
+            if accepted == draft.len() {
+                debug_assert_eq!(
+                    self.session.current_len(),
+                    base + accepted,
+                    "full-accept commit: verify must leave KV + state at base + accepted"
+                );
+            } else {
+                match recurrent_snapshot.as_ref() {
+                    Some(snapshot) => self.session.commit_recurrent_state_to_accepted(
+                        snapshot,
+                        base,
+                        &draft[..accepted],
+                    )?,
+                    None => self.session.rewind(base + accepted)?,
+                }
             }
 
             // Commit accepted draft tokens followed by the bonus, honoring the