diff --git a/.squad/decisions/inbox/gaff-mtp-approach-b-snapshot-d2d.md b/.squad/decisions/inbox/gaff-mtp-approach-b-snapshot-d2d.md new file mode 100644 index 0000000000..63d13b98b5 --- /dev/null +++ b/.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. diff --git a/crates/onnx-genai-engine/src/native_decode/cuda.rs b/crates/onnx-genai-engine/src/native_decode/cuda.rs index 63d317b8b0..95cbeba895 100644 --- a/crates/onnx-genai-engine/src/native_decode/cuda.rs +++ b/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, + /// 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, pub(crate) input_ids_binding: usize, pub(crate) position_ids_binding: Option, @@ -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)>> { - 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)], - ) -> 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(()) } diff --git a/crates/onnx-genai-engine/src/native_decode/mod.rs b/crates/onnx-genai-engine/src/native_decode/mod.rs index b43084bcfa..434cb05a20 100644 --- a/crates/onnx-genai-engine/src/native_decode/mod.rs +++ b/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>, - device: Option)>>, + /// 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 { diff --git a/crates/onnx-genai-engine/src/native_speculative.rs b/crates/onnx-genai-engine/src/native_speculative.rs index 8485741d6c..aad5e055df 100644 --- a/crates/onnx-genai-engine/src/native_speculative.rs +++ b/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 diff --git a/crates/onnx-runtime-ep-api/src/provider.rs b/crates/onnx-runtime-ep-api/src/provider.rs index 5ce57720ee..644e05a4cc 100644 --- a/crates/onnx-runtime-ep-api/src/provider.rs +++ b/crates/onnx-runtime-ep-api/src/provider.rs @@ -1157,6 +1157,30 @@ pub trait ExecutionProvider: Send + Sync { /// Block until all pending work on this EP completes. fn sync(&self) -> Result<()>; + /// Copy `bytes` device→device, from `src[src_offset..]` into + /// `dst[dst_offset..]`, both allocations owned by this EP. + /// + /// The default errors: only device EPs with a native device-to-device copy + /// (CUDA) implement it. Used by the speculative recurrent-state snapshot to + /// stage the destructive GDN/conv state into device scratch WITHOUT a PCIe + /// round-trip through host memory. The copy is stream-ordered on the EP's + /// compute stream, so it composes with the surrounding forward passes + /// (snapshot before the verify overwrites the state; restore before the + /// accepted-token re-advance) with no host synchronization. + fn copy_device_to_device( + &self, + _src: &DeviceBuffer, + _src_offset: usize, + _dst: &mut DeviceBuffer, + _dst_offset: usize, + _bytes: usize, + ) -> Result<()> { + Err(EpError::KernelFailed(format!( + "{}: device-to-device copy is not implemented", + self.name() + ))) + } + /// EP-specific optimization passes, run after the generic optimizer. fn custom_passes(&self) -> Vec> { Vec::new() diff --git a/crates/onnx-runtime-ep-cuda/src/provider.rs b/crates/onnx-runtime-ep-cuda/src/provider.rs index 7ff553baaf..33419c1d2e 100644 --- a/crates/onnx-runtime-ep-cuda/src/provider.rs +++ b/crates/onnx-runtime-ep-cuda/src/provider.rs @@ -1508,6 +1508,65 @@ impl ExecutionProvider for CudaExecutionProvider { unsafe { self.runtime.dtoh(dst, cuptr(src.as_ptr())) } } + fn copy_device_to_device( + &self, + src: &DeviceBuffer, + src_offset: usize, + dst: &mut DeviceBuffer, + dst_offset: usize, + bytes: usize, + ) -> Result<()> { + assert_eq!( + src.device(), + self.device, + "cuda_ep::copy_device_to_device: foreign src buffer" + ); + assert_eq!( + dst.device(), + self.device, + "cuda_ep::copy_device_to_device: foreign dst buffer" + ); + if bytes == 0 { + return Ok(()); + } + let src_end = src_offset.checked_add(bytes).ok_or_else(|| { + EpError::KernelFailed("cuda_ep::copy_device_to_device: src range overflows".into()) + })?; + if src_end > src.len() { + return Err(EpError::KernelFailed(format!( + "cuda_ep::copy_device_to_device: src range {src_offset}..{src_end} exceeds {}", + src.len() + ))); + } + let dst_end = dst_offset.checked_add(bytes).ok_or_else(|| { + EpError::KernelFailed("cuda_ep::copy_device_to_device: dst range overflows".into()) + })?; + if dst_end > dst.len() { + return Err(EpError::KernelFailed(format!( + "cuda_ep::copy_device_to_device: dst range {dst_offset}..{dst_end} exceeds {}", + dst.len() + ))); + } + let src_ptr = cuptr(src.as_ptr()) + .checked_add(src_offset as u64) + .ok_or_else(|| { + EpError::KernelFailed( + "cuda_ep::copy_device_to_device: src pointer offset overflows".into(), + ) + })?; + let dst_ptr = cuptr(dst.as_mut_ptr()) + .checked_add(dst_offset as u64) + .ok_or_else(|| { + EpError::KernelFailed( + "cuda_ep::copy_device_to_device: dst pointer offset overflows".into(), + ) + })?; + // SAFETY: both pointers name checked byte ranges within live device + // allocations that outlive the enqueued copy; the copy is ordered on the + // EP stream against the surrounding forward passes that read/write them. + unsafe { self.runtime.dtod_async(src_ptr, dst_ptr, bytes) } + } + fn begin_device_graph_capture(&self, kernels: &[&dyn Kernel]) -> Result<()> { self.runtime.begin_graph_capture(kernels) } diff --git a/crates/onnx-runtime-session/src/executor/tests.rs b/crates/onnx-runtime-session/src/executor/tests.rs index b3c0799d95..837ab1e9a0 100644 --- a/crates/onnx-runtime-session/src/executor/tests.rs +++ b/crates/onnx-runtime-session/src/executor/tests.rs @@ -2988,6 +2988,7 @@ fn inference_session_fallback_workspace_grows_retries_and_reuses() { model_metadata: crate::ModelMetadata::default(), exec, decode_inline_exec: None, + verify_exec: None, ep_context_config: crate::EpContextDumpConfig::default(), }; @@ -3091,6 +3092,7 @@ fn prepared_session_reprepares_workspace_when_execution_rebuckets() { model_metadata: crate::ModelMetadata::default(), exec, decode_inline_exec: None, + verify_exec: None, ep_context_config: crate::EpContextDumpConfig::default(), }; diff --git a/crates/onnx-runtime-session/src/lib.rs b/crates/onnx-runtime-session/src/lib.rs index d80364364b..2d846eaae6 100644 --- a/crates/onnx-runtime-session/src/lib.rs +++ b/crates/onnx-runtime-session/src/lib.rs @@ -38,6 +38,7 @@ pub use executor::{ exec_phase_stats, plan_double_buffer, print_exec_phase_profile, reset_dense_prefetch_gap_stats, reset_exec_phase_profile, }; +pub use onnx_runtime_ep_api::DeviceBuffer; pub use onnx_runtime_ep_api::DeviceGraphSlot; pub use onnx_runtime_ep_api::WorkspaceRequirement; pub use onnx_runtime_loader::{ diff --git a/crates/onnx-runtime-session/src/tensor.rs b/crates/onnx-runtime-session/src/tensor.rs index 6281c339e0..efd3324ead 100644 --- a/crates/onnx-runtime-session/src/tensor.rs +++ b/crates/onnx-runtime-session/src/tensor.rs @@ -657,6 +657,56 @@ impl DeviceIoBinding { self.buffer().as_ptr() } + /// The execution provider that owns this binding's device allocation. + /// Callers use it to allocate device scratch compatible with + /// [`Self::snapshot_device_into`] / [`Self::restore_device_from`]. + pub fn allocator(&self) -> &Arc { + &self.allocator + } + + /// Copy the first `bytes` of this binding's device allocation into + /// `scratch` (device→device, no host round-trip). Used to snapshot the + /// destructive recurrent/conv state before a speculative verify overwrites + /// it. `scratch` must be owned by the same EP and hold at least `bytes`. + pub fn snapshot_device_into(&self, scratch: &mut DeviceBuffer, bytes: usize) -> Result<()> { + let buffer = self.buffer(); + if bytes > buffer.len() { + return Err(SessionError::ExternalBuffer { + binding: self.input_name.clone(), + reason: format!( + "device snapshot of {bytes} bytes exceeds allocation of {}", + buffer.len() + ), + }); + } + self.allocator + .copy_device_to_device(buffer, 0, scratch, 0, bytes)?; + Ok(()) + } + + /// Copy the first `bytes` of `scratch` back into this binding's device + /// allocation (device→device). Inverse of [`Self::snapshot_device_into`]; + /// leaves the binding shape unchanged so it never invalidates a captured + /// decode graph. + pub fn restore_device_from(&mut self, scratch: &DeviceBuffer, bytes: usize) -> Result<()> { + let buffer = self + .buffer + .as_mut() + .expect("DeviceIoBinding buffer taken only in Drop"); + if bytes > buffer.len() { + return Err(SessionError::ExternalBuffer { + binding: self.input_name.clone(), + reason: format!( + "device restore of {bytes} bytes exceeds allocation of {}", + buffer.len() + ), + }); + } + self.allocator + .copy_device_to_device(scratch, 0, buffer, 0, bytes)?; + Ok(()) + } + pub fn transfer_stats(&self) -> DeviceBindingTransferStats { self.transfer_stats }