diff --git a/crates/onnx-genai-engine/src/native_component.rs b/crates/onnx-genai-engine/src/native_component.rs index 9e9cbf8a12..86693751e5 100644 --- a/crates/onnx-genai-engine/src/native_component.rs +++ b/crates/onnx-genai-engine/src/native_component.rs @@ -126,7 +126,7 @@ impl NativeComponentSession { NativeDecodeDevice::Cuda { index } => DevicePreference::Gpu { index }, NativeDecodeDevice::Plugin { .. } => DevicePreference::Cpu, }; - let session = InferenceSession::builder() + let mut session = InferenceSession::builder() .model(path) .device(preference) .build() @@ -136,6 +136,12 @@ impl NativeComponentSession { path.display() ) })?; + // A component graph is run once per request and never records a device + // graph, so its intermediates can be freed as they die. The vision + // encoder is why this matters: holding every one of its 2545 node + // outputs at once cost ~23 GB at 448px and put full-resolution images + // out of reach entirely. + session.set_release_dead_values(true); if requested_cuda && let Some(report) = session.execution_provider_fallback_report() { tracing::warn!( model = %path.display(), diff --git a/crates/onnx-genai-server/src/driver.rs b/crates/onnx-genai-server/src/driver.rs index 33b21cd652..4b721e1209 100644 --- a/crates/onnx-genai-server/src/driver.rs +++ b/crates/onnx-genai-server/src/driver.rs @@ -210,7 +210,11 @@ impl DriverFailure { ) }); Self { - message: error.to_string(), + // Anyhow's Display shows only the outermost context, which for a + // decode failure is the generic "forward pass failed" wrapper. The + // alternate form keeps the whole chain, and this message is the + // only thing the client ever sees. + message: format!("{error:#}"), kind: if memory_overload { DriverFailureKind::MemoryOverload } else { @@ -1123,7 +1127,7 @@ fn run_static_batch_until_idle( route_continuous_admissions(manager.poll_admissions(), &mut routes); if let Err(err) = manager.step() { let mut failure = DriverFailure::from_engine_error(&err); - failure.message = format!("continuous batch generation failed: {err}"); + failure.message = format!("continuous batch generation failed: {err:#}"); for (_, mut route) in routes.drain() { if let Some(sender) = route.admission.take() { let _ = sender.send(Err(failure.clone())); diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/contrib_fused.rs b/crates/onnx-runtime-ep-cpu/src/kernels/contrib_fused.rs index 355492a580..f4e8e36ec6 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/contrib_fused.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/contrib_fused.rs @@ -56,14 +56,23 @@ fn require_skip_layer_norm_dtype(inputs: &[TensorView], outputs: &[TensorMut]) - if output.is_absent() { continue; } - let expected = if matches!(index, 1 | 2) { - DataType::Float32 - } else { - dtype - }; - if output.dtype != expected { + if matches!(index, 1 | 2) { + // `mean`/`inv_std_var` are training-only stats that inference + // graphs leave dangling. ORT's schema types them `float`, but + // exporters routinely emit them in the activation dtype instead + // (the bf16 skip-RMSNorm path documents the same allowance), and + // the write below narrows to whichever was declared. + if output.dtype != DataType::Float32 && output.dtype != dtype { + return Err(EpError::KernelFailed(format!( + "SkipLayerNormalization: output {index} must have dtype Float32 or {dtype:?}, got {:?}", + output.dtype + ))); + } + continue; + } + if output.dtype != dtype { return Err(EpError::KernelFailed(format!( - "SkipLayerNormalization: output {index} must have dtype {expected:?}, got {:?}", + "SkipLayerNormalization: output {index} must have dtype {dtype:?}, got {:?}", output.dtype ))); } @@ -811,6 +820,76 @@ mod tests { assert_close(&y.to_f32(), &want); } + /// A bf16 graph may declare the dangling `mean`/`inv_std_var` stats in the + /// activation dtype rather than the schema's float. The Muse Glimmer vision + /// tower does exactly this, so the kernel must accept it and narrow the + /// stat writes instead of rejecting the node. + #[test] + fn skip_layer_norm_accepts_bf16_stat_outputs() { + let x = Owned::bf16(&[2, 4], &[1., 2., 3., 4., 2., 3., 4., 5.]); + let skip = Owned::bf16(&[2, 4], &[0.5, -1., 1., 0., 1., 0., -1., 2.]); + let gamma = Owned::bf16(&[4], &[1., 2., 0.5, 1.5]); + let mut y = Owned::zeros(DataType::BFloat16, &[2, 4]); + let mut mean = Owned::zeros(DataType::BFloat16, &[2, 1]); + let mut inv_std = Owned::zeros(DataType::BFloat16, &[2, 1]); + SkipLayerNormKernel { epsilon: 1e-5 } + .execute( + &[x.view(), skip.view(), gamma.view()], + &mut [y.view_mut(), mean.view_mut(), inv_std.view_mut()], + ) + .unwrap(); + // sum = X + skip; row means are 2.625 and 4.0. + let means = mean.to_bf16_as_f32(); + assert!((means[0] - 2.625).abs() < 0.05, "got {means:?}"); + assert!((means[1] - 4.0).abs() < 0.05, "got {means:?}"); + assert!( + inv_std + .to_bf16_as_f32() + .iter() + .all(|v| v.is_finite() && *v > 0.) + ); + } + + /// The same node with schema-conformant float stats stays accepted. + #[test] + fn skip_layer_norm_accepts_f32_stat_outputs_in_a_bf16_graph() { + let x = Owned::bf16(&[2, 4], &[1., 2., 3., 4., 2., 3., 4., 5.]); + let skip = Owned::bf16(&[2, 4], &[0.5, -1., 1., 0., 1., 0., -1., 2.]); + let gamma = Owned::bf16(&[4], &[1., 2., 0.5, 1.5]); + let mut y = Owned::zeros(DataType::BFloat16, &[2, 4]); + let mut mean = Owned::zeros_f32(&[2, 1]); + let mut inv_std = Owned::zeros_f32(&[2, 1]); + SkipLayerNormKernel { epsilon: 1e-5 } + .execute( + &[x.view(), skip.view(), gamma.view()], + &mut [y.view_mut(), mean.view_mut(), inv_std.view_mut()], + ) + .unwrap(); + let means = mean.to_f32(); + assert!((means[0] - 2.625).abs() < 0.05, "got {means:?}"); + assert!((means[1] - 4.0).abs() < 0.05, "got {means:?}"); + } + + /// A stat output in an unrelated dtype is still a hard error. + #[test] + fn skip_layer_norm_rejects_unrelated_stat_dtype() { + let x = Owned::bf16(&[2, 4], &[1., 2., 3., 4., 2., 3., 4., 5.]); + let skip = Owned::bf16(&[2, 4], &[0.5, -1., 1., 0., 1., 0., -1., 2.]); + let gamma = Owned::bf16(&[4], &[1., 2., 0.5, 1.5]); + let mut y = Owned::zeros(DataType::BFloat16, &[2, 4]); + let mut mean = Owned::zeros(DataType::Float16, &[2, 1]); + let error = SkipLayerNormKernel { epsilon: 1e-5 } + .execute( + &[x.view(), skip.view(), gamma.view()], + &mut [y.view_mut(), mean.view_mut()], + ) + .unwrap_err(); + assert!( + error.to_string().contains("Float32 or BFloat16"), + "unexpected error: {error}" + ); + } + /// A 4-output SkipLayerNorm node writes `output`, `mean`, `inv_std_var`, and /// `input_skip_bias_sum` (= X + skip + bias), all numerically correct. #[test] diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/reduce_ops.rs b/crates/onnx-runtime-ep-cpu/src/kernels/reduce_ops.rs index 7e90a9e01b..2d75c5b034 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/reduce_ops.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/reduce_ops.rs @@ -158,8 +158,13 @@ impl Kernel for ReduceKernel { if matches!(self.op, ReduceOp::Max | ReduceOp::Min) && inputs[0].dtype == DataType::Bool { return self.execute_bool(inputs, &mut outputs[0]); } - if matches!(self.op, ReduceOp::Sum) && inputs[0].dtype == onnx_runtime_ir::DataType::Int64 { - return self.execute_i64_sum(inputs, &mut outputs[0]); + if matches!(inputs[0].dtype, DataType::Int64 | DataType::Int32) + && matches!( + self.op, + ReduceOp::Sum | ReduceOp::Max | ReduceOp::Min | ReduceOp::Prod + ) + { + return self.execute_integer(inputs, &mut outputs[0]); } let x = to_dense_f32_widen(self.op.name(), &inputs[0])?; let in_shape = inputs[0].shape; @@ -277,7 +282,17 @@ impl ReduceKernel { Ok(reduce) } - fn execute_i64_sum(&self, inputs: &[TensorView], output: &mut TensorMut) -> Result<()> { + /// Reduce an integer tensor without leaving the integer domain. + /// + /// The general path widens to `f32`, which silently rounds any magnitude + /// past 2^24 and rejects integer dtypes outright. ONNX defines + /// `ReduceSum`/`ReduceMax`/`ReduceMin`/`ReduceProd` over the integer types, + /// and callers use them on things like token counts and patch grid extents + /// where an off-by-rounding answer is worse than no answer. + /// + /// Accumulation happens in `i64` regardless of the input width; the result + /// is narrowed to whatever the output view declares. + fn execute_integer(&self, inputs: &[TensorView], output: &mut TensorMut) -> Result<()> { let x = to_dense_i64(&inputs[0])?; let in_shape = inputs[0].shape; let rank = in_shape.len(); @@ -288,7 +303,14 @@ impl ReduceKernel { .collect::>(); let in_strides = compute_contiguous_strides(in_shape); let kept_strides = compute_contiguous_strides(&kept_shape); - let mut sums = vec![0_i64; numel(&kept_shape).max(1)]; + let seed = match self.op { + ReduceOp::Sum => 0, + ReduceOp::Prod => 1, + ReduceOp::Max => i64::MIN, + ReduceOp::Min => i64::MAX, + _ => unreachable!("integer reduction is only valid for Sum/Max/Min/Prod"), + }; + let mut acc = vec![seed; numel(&kept_shape).max(1)]; if numel(in_shape) > 0 { let mut index = vec![0; rank]; loop { @@ -302,18 +324,40 @@ impl ReduceKernel { kept_axis += 1; } } - sums[output_offset] = sums[output_offset].wrapping_add(x[input_offset]); + let value = x[input_offset]; + let slot = &mut acc[output_offset]; + *slot = match self.op { + ReduceOp::Sum => slot.wrapping_add(value), + ReduceOp::Prod => slot.wrapping_mul(value), + ReduceOp::Max => (*slot).max(value), + ReduceOp::Min => (*slot).min(value), + _ => unreachable!("integer reduction is only valid for Sum/Max/Min/Prod"), + }; if !next_index(in_shape, &mut index) { break; } } } let _ = self.keepdims; - let bytes = sums - .iter() - .take(numel(&kept_shape)) - .flat_map(|value| value.to_le_bytes()) - .collect::>(); + let kept = numel(&kept_shape); + let bytes = match output.dtype { + DataType::Int64 => acc + .iter() + .take(kept) + .flat_map(|value| value.to_le_bytes()) + .collect::>(), + DataType::Int32 => acc + .iter() + .take(kept) + .flat_map(|value| (*value as i32).to_le_bytes()) + .collect::>(), + other => { + return Err(EpError::KernelFailed(format!( + "{}: integer reduction output must be Int64 or Int32, got {other:?}", + self.op.name() + ))); + } + }; write_dense_bytes(output, &bytes) } @@ -420,6 +464,78 @@ mod tests { assert_eq!(out.to_i64(), vec![6, 15]); } + #[test] + fn max_min_prod_int64_axis1() { + let x = Owned::i64(&[2, 3], &[1, 9, 3, 4, 5, 0]); + for (op, expected) in [ + (ReduceOp::Max, vec![9, 5]), + (ReduceOp::Min, vec![1, 0]), + (ReduceOp::Prod, vec![27, 0]), + ] { + let mut out = Owned::zeros(DataType::Int64, &[2, 1]); + run_attr(op, Some(vec![1]), &x, &mut out); + assert_eq!(out.to_i64(), expected); + } + } + + #[test] + fn max_int64_reduce_all_to_scalar() { + // The shape the Muse Glimmer vision tower produces: reduce a 1-D int64 + // patch-extent vector down to a scalar with no axes input. + let x = Owned::i64(&[3], &[7, 42, 13]); + let mut out = Owned::zeros(DataType::Int64, &[1]); + run_omitted_axes(ReduceOp::Max, &x, &mut out); + assert_eq!(out.to_i64(), vec![42]); + } + + #[test] + fn max_int64_preserves_values_beyond_f32_precision() { + // Widening through f32 rounds anything past 2^24, so these two distinct + // values would collapse onto the same float and pick the wrong maximum. + let low = (1_i64 << 40) + 1; + let high = (1_i64 << 40) + 2; + let x = Owned::i64(&[2], &[high, low]); + let mut out = Owned::zeros(DataType::Int64, &[1]); + run_omitted_axes(ReduceOp::Max, &x, &mut out); + assert_eq!(out.to_i64(), vec![high]); + + let mut out = Owned::zeros(DataType::Int64, &[1]); + run_omitted_axes(ReduceOp::Min, &x, &mut out); + assert_eq!(out.to_i64(), vec![low]); + } + + #[test] + fn min_int64_negative_values() { + let x = Owned::i64(&[2, 2], &[-5, 3, -9, -1]); + let mut out = Owned::zeros(DataType::Int64, &[2, 1]); + run_attr(ReduceOp::Min, Some(vec![1]), &x, &mut out); + assert_eq!(out.to_i64(), vec![-5, -9]); + } + + #[test] + fn max_int32_reduces_to_int32_output() { + // The shape the Muse Glimmer vision tower attention produces: an int32 + // window-extent vector reduced to a single int32. + let x = Owned::i32(&[6], &[3, 11, 7, 11, 2, 9]); + let mut out = Owned::zeros(DataType::Int32, &[1]); + run_omitted_axes(ReduceOp::Max, &x, &mut out); + assert_eq!(out.to_i32(), vec![11]); + } + + #[test] + fn int32_sum_min_prod_reduce_to_int32_output() { + let x = Owned::i32(&[2, 3], &[1, 9, 3, 4, 5, 0]); + for (op, expected) in [ + (ReduceOp::Sum, vec![13, 9]), + (ReduceOp::Min, vec![1, 0]), + (ReduceOp::Prod, vec![27, 0]), + ] { + let mut out = Owned::zeros(DataType::Int32, &[2, 1]); + run_attr(op, Some(vec![1]), &x, &mut out); + assert_eq!(out.to_i32(), expected); + } + } + #[test] fn sum_axes_from_input() { // Modern opset-13 signature: axes come as input 1. diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/sequence.rs b/crates/onnx-runtime-ep-cpu/src/kernels/sequence.rs index 8184057a06..a707ac2502 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/sequence.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/sequence.rs @@ -139,6 +139,27 @@ impl Kernel for RangeKernel { let bytes: Vec = out.iter().flat_map(|v| v.to_le_bytes()).collect(); write_dense_bytes(&mut outputs[0], &bytes) } + DataType::Int32 => { + let (start, limit, delta) = ( + to_dense_i64(&inputs[0])?[0], + to_dense_i64(&inputs[1])?[0], + to_dense_i64(&inputs[2])?[0], + ); + if delta == 0 { + return Err(EpError::KernelFailed( + "Range: delta must not be zero".into(), + )); + } + let count = int_range_count(start, limit, delta)?; + let mut out = alloc_range_output::(count)?; + for i in 0..count { + let value = i32::try_from(start as i128 + i as i128 * delta as i128) + .map_err(|_| EpError::KernelFailed("Range: value overflow".into()))?; + out.push(value); + } + let bytes: Vec = out.iter().flat_map(|v| v.to_le_bytes()).collect(); + write_dense_bytes(&mut outputs[0], &bytes) + } dtype => Err(EpError::KernelFailed(format!( "Range: unsupported dtype {dtype:?}" ))), @@ -454,6 +475,29 @@ mod tests { assert_eq!(y.to_f32().last(), Some(&16_777_216.)); } #[test] + fn range_int32_produces_int32_output() { + // The Muse Glimmer vision tower builds int32 window offsets this way. + let a = Owned::i32(&[], &[2]); + let b = Owned::i32(&[], &[8]); + let d = Owned::i32(&[], &[2]); + let mut y = Owned::zeros(DataType::Int32, &[3]); + RangeKernel + .execute(&[a.view(), b.view(), d.view()], &mut [y.view_mut()]) + .unwrap(); + assert_eq!(y.to_i32(), vec![2, 4, 6]); + } + #[test] + fn range_int32_negative_delta_counts_down() { + let a = Owned::i32(&[], &[5]); + let b = Owned::i32(&[], &[1]); + let d = Owned::i32(&[], &[-2]); + let mut y = Owned::zeros(DataType::Int32, &[2]); + RangeKernel + .execute(&[a.view(), b.view(), d.view()], &mut [y.view_mut()]) + .unwrap(); + assert_eq!(y.to_i32(), vec![5, 3]); + } + #[test] fn range_int64_overflow_returns_error() { // count ~= i64::MAX elements; count * size_of::() overflows the // addressable-byte limit, so the kernel must return an Err before any diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/normalization.rs b/crates/onnx-runtime-ep-cuda/src/kernels/normalization.rs index f36b6dd96d..07a6032bac 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/normalization.rs +++ b/crates/onnx-runtime-ep-cuda/src/kernels/normalization.rs @@ -1240,6 +1240,7 @@ extern "C" __global__ void skip_layernorm( const int gamma_dtype, const int beta_dtype, const int bias_dtype, + const int stat_dtype, const int has_beta, const int has_bias, const float epsilon) @@ -1290,8 +1291,8 @@ extern "C" __global__ void skip_layernorm( const float var = red[0] / (float)norm_size; const float inv_std = 1.0f / sqrtf(var + epsilon); if (tid == 0) { - if (mean_out) skip_ln_store(mean_out, g, mean, dtype); - if (invstd_out) skip_ln_store(invstd_out, g, inv_std, dtype); + if (mean_out) skip_ln_store(mean_out, g, mean, stat_dtype); + if (invstd_out) skip_ln_store(invstd_out, g, inv_std, stat_dtype); } __syncthreads(); @@ -2880,7 +2881,7 @@ impl SkipLayerNormKernel { // Optional outputs: mean (slot 1), inv_std_var (slot 2) — length // num_groups; input_skip_bias_sum (slot 3) — length input.numel(). - let (mean_ptr, invstd_ptr) = + let (mean_ptr, invstd_ptr, stat_dtype) = optional_stat_ptrs_typed(op, outputs, num_groups, input.dtype)?; let sum_ptr = match outputs.get_mut(3) { None => 0u64, @@ -2913,6 +2914,7 @@ impl SkipLayerNormKernel { let gamma_dtype = storage_kind(gamma.dtype); let beta_dtype = beta.map_or(dtype, |tensor| storage_kind(tensor.dtype)); let bias_dtype = bias.map_or(dtype, |tensor| storage_kind(tensor.dtype)); + let stat_dtype_kind = storage_kind(stat_dtype); let func = self.runtime.nvrtc_function( SKIP_LAYERNORM_MODULE, @@ -2944,6 +2946,7 @@ impl SkipLayerNormKernel { .arg(&gamma_dtype) .arg(&beta_dtype) .arg(&bias_dtype) + .arg(&stat_dtype_kind) .arg(&has_beta) .arg(&has_bias) .arg(&eps); @@ -3037,10 +3040,33 @@ fn optional_stat_ptrs_typed( outputs: &mut [TensorMut], num_groups: usize, dtype: DataType, -) -> Result<(CUdeviceptr, CUdeviceptr)> { - let mean = optional_out_ptr_typed(op, "Mean", outputs, 1, num_groups, dtype)?; - let invstd = optional_out_ptr_typed(op, "InvStdDev", outputs, 2, num_groups, dtype)?; - Ok((mean, invstd)) +) -> Result<(CUdeviceptr, CUdeviceptr, DataType)> { + // ORT's schema types `mean`/`inv_std_var` as float, but exporters commonly + // emit these dangling training-only stats in the activation dtype. Accept + // either and tell the kernel which storage to write. + let mut stat_dtype = DataType::Float32; + for idx in [1, 2] { + if let Some(t) = outputs.get(idx) { + if t.dtype != DataType::Float32 && t.dtype != dtype { + return Err(EpError::KernelFailed(format!( + "cuda_ep {op}: stat output {idx} dtype {:?} must be Float32 or the input dtype {dtype:?}", + t.dtype + ))); + } + stat_dtype = t.dtype; + } + } + if let (Some(mean), Some(invstd)) = (outputs.get(1), outputs.get(2)) + && mean.dtype != invstd.dtype + { + return Err(EpError::KernelFailed(format!( + "cuda_ep {op}: Mean dtype {:?} and InvStdDev dtype {:?} must match", + mean.dtype, invstd.dtype + ))); + } + let mean = optional_out_ptr_typed(op, "Mean", outputs, 1, num_groups, stat_dtype)?; + let invstd = optional_out_ptr_typed(op, "InvStdDev", outputs, 2, num_groups, stat_dtype)?; + Ok((mean, invstd, stat_dtype)) } fn optional_out_ptr_typed( diff --git a/crates/onnx-runtime-session/src/executor/build.rs b/crates/onnx-runtime-session/src/executor/build.rs index cb10afa9ed..3cc7540603 100644 --- a/crates/onnx-runtime-session/src/executor/build.rs +++ b/crates/onnx-runtime-session/src/executor/build.rs @@ -535,6 +535,7 @@ impl Executor { decode_view_plan_sig_mismatch_streak: 0, decode_view_plan_disabled: false, compute_in_place_enabled: compute_in_place_env_enabled(), + release_dead_values_enabled: false, compute_in_place_alias_count: 0, scan_inline_single_trip_enabled: scan_inline_single_trip_env_enabled(), scan_inline_single_trip_count: 0, @@ -1058,6 +1059,7 @@ impl Executor { input_dtypes, output_dtypes, inplace_dead_inputs: Vec::new(), + dead_after: Vec::new(), lazy_weight_inputs, }); } @@ -1091,6 +1093,27 @@ impl Executor { }) }) .collect(); + // Values this node consumes for the last time. Weights are excluded + // here rather than at runtime because an initializer's buffer is + // built once and reused by every later run -- releasing it after its + // final *use in this run* would leave the next run with no weights. + // Graph inputs are excluded for the same reason at a different + // lifetime: their storage belongs to the caller's binding. + let mut seen = std::collections::HashSet::new(); + let dead_after: Vec = node + .inputs + .iter() + .flatten() + .copied() + .filter(|vid| { + last_use.get(vid) == Some(&pi) + && !graph_outputs.contains(vid) + && !graph.initializers.contains_key(vid) + && !graph.inputs.contains(vid) + && seen.insert(*vid) + }) + .collect(); + node.dead_after = dead_after; } if let Some(span) = plan_span.as_mut() { span.set_args( @@ -1942,6 +1965,16 @@ impl Executor { self.trace = trace; } + /// Enable/disable dead-value buffer release, propagating to control-flow + /// child executors so a Scan/Loop body frees its intermediates too -- the + /// vision encoder's per-image Scan bodies are a large part of its live set. + pub(crate) fn set_release_dead_values(&mut self, enabled: bool) { + for child in self.subgraph_execs.values_mut() { + child.set_release_dead_values(enabled); + } + self.release_dead_values_enabled = enabled; + } + /// Live weight bytes backing the graph, needed alongside [`Self::graph`] so /// the EPContext dump can encode initializers into the context model. pub(crate) fn weights(&self) -> &Arc { diff --git a/crates/onnx-runtime-session/src/executor/control_flow.rs b/crates/onnx-runtime-session/src/executor/control_flow.rs index 57a96027a0..477ce5566c 100644 --- a/crates/onnx-runtime-session/src/executor/control_flow.rs +++ b/crates/onnx-runtime-session/src/executor/control_flow.rs @@ -58,6 +58,7 @@ impl ChildExecutor { builds: 0, runs: 0, trace: TraceContext::noop(), + release_dead_values: false, }) } @@ -77,6 +78,15 @@ impl ChildExecutor { self.trace = trace; } + /// Mirror of [`Self::set_trace_context`] for dead-value release: applied to + /// already-compiled child plans and remembered for plans compiled later. + pub(crate) fn set_release_dead_values(&mut self, enabled: bool) { + for plan in &mut self.compiled { + plan.exec.set_release_dead_values(enabled); + } + self.release_dead_values = enabled; + } + pub(super) fn compile(&self, externals: &[&Tensor]) -> Result { let mut graph = self.template.clone(); // GraphProto has no opset table: nested graphs inherit the model-level @@ -124,6 +134,7 @@ impl ChildExecutor { exec: { let mut exec = Executor::build(graph, self.weights.clone(), self.ep.clone())?; exec.set_trace_context(self.trace.clone()); + exec.set_release_dead_values(self.release_dead_values); exec }, signature: externals @@ -623,6 +634,7 @@ impl Executor { self.ep.clone(), )?; child.set_trace_context(self.trace.clone()); + child.set_release_dead_values(self.release_dead_values_enabled); self.subgraph_execs.insert(prepared.key.clone(), child); } diff --git a/crates/onnx-runtime-session/src/executor/dispatch.rs b/crates/onnx-runtime-session/src/executor/dispatch.rs index 39d63a6492..89913a22d9 100644 --- a/crates/onnx-runtime-session/src/executor/dispatch.rs +++ b/crates/onnx-runtime-session/src/executor/dispatch.rs @@ -171,6 +171,7 @@ impl Executor { entry.0 += elapsed; entry.1 += 1; result?; + self.release_dead_values(pi, external); } print_op_profile(run_start.elapsed(), timings); } else { @@ -180,11 +181,76 @@ impl Executor { } self.prefetch_lazy_weights_after(pi)?; self.exec_plan_node(pi, resolved, outer_scope, external, OpCaptureTrace::Eager)?; + self.release_dead_values(pi, external); } } Ok(()) } + /// Release the buffers of values whose last consumer was plan node `pi`. + /// + /// The guards mirror the in-place-overwrite path: anything the caller owns, + /// anything aliased by a view, and anything a sequence or shared buffer still + /// refers to is left alone. A value that fails a guard simply stays resident, + /// which is the pre-existing behaviour for every value. + /// Free the buffers of values whose last consumer was node `pi`. + /// + /// The guard list is deliberately the same one `inplace_input` applies, + /// because both are asking the same question: may this buffer be taken away + /// from the value that currently holds it? Any guard added there belongs + /// here too. + /// + /// A value that is aliased -- by a view, by a sequence's shared storage -- + /// is never released, even after every alias has died. `pinned` records + /// that a value was ever an alias source and is only cleared between runs, + /// so it is a deliberately conservative over-approximation. Releasing on + /// the exact liveness test (`views` no longer names it as a source) would + /// be wrong for the decode memo, which restores a step's view table from + /// `retained_views` on the *next* run and expects the source buffer to + /// still be resident. Measured cost of the conservative choice on the + /// vision encoder is 36 MiB out of ~12 GB reclaimed, which does not buy an + /// aliasing analysis that has to be right across runs. + fn release_dead_values(&mut self, pi: usize, external: &ExternalBindings) { + if !self.release_dead_values_enabled { + return; + } + let dead = std::mem::take(&mut self.plan[pi].dead_after); + for &vid in &dead { + if external.inputs.contains_key(&vid) + || external.outputs.contains_key(&vid) + || self.pinned.contains(&vid) + || self.shared_buffers.contains_key(&vid) + || self.seq_elem_values.contains_key(&vid) + || self.views.contains_key(&vid) + { + continue; + } + // A memoized loop-invariant `If` skips its branch on later runs and + // serves the outputs straight from the buffers the first run left + // resident (see `exec_if`), and the memo outlives an eager run. + // Freeing one of those outputs turns the next skip into a missing + // buffer. `try_move_host_output` declines for the same reason. + if let Some(producer) = self.graph.try_value(vid).and_then(|value| value.producer) + && self.if_last_predicate.contains_key(&producer) + { + continue; + } + // A borrowed buffer aliases memory this session does not own. + if self.buffers.get(&vid).is_some_and(|b| b.is_borrowed()) { + continue; + } + // `DeviceBuffer` is a bare handle with no `Drop`: dropping it + // strands the allocation, so the owning EP has to free it. + if let Some(buffer) = self.buffers.remove(&vid) { + // The reuse fast path treats a surviving shape entry as proof + // the allocation is still there, so it has to go with it. + self.buffer_shapes.remove(&vid); + let _ = self.ep.deallocate(buffer); + } + } + self.plan[pi].dead_after = dead; + } + fn prefetch_lazy_weights_after(&self, pi: usize) -> Result<()> { if self.prefetch_lookahead_nodes == 0 { return Ok(()); diff --git a/crates/onnx-runtime-session/src/executor/mod.rs b/crates/onnx-runtime-session/src/executor/mod.rs index 907844c569..6bf47cb446 100644 --- a/crates/onnx-runtime-session/src/executor/mod.rs +++ b/crates/onnx-runtime-session/src/executor/mod.rs @@ -518,6 +518,16 @@ pub(crate) struct NodePlan { pub inplace_dead_inputs: Vec, /// Lazy weight inputs this node may ask the EP to page at dispatch time. pub lazy_weight_inputs: Vec, + /// Intermediate values whose last consumer is this node, and whose buffers + /// may therefore be released once it has run. + /// + /// Distinct from [`NodePlan::inplace_dead_inputs`], which only ever lets a + /// kernel overwrite an input whose shape and dtype already match its output. + /// That covers elementwise chains and nothing else, so before this list + /// existed every intermediate buffer a graph produced stayed resident for the + /// whole run: a 2545-node vision encoder held all 2545 of them at once and + /// spent ~20x the memory its live set needs. + pub dead_after: Vec, } /// Map a [`crate::sequence::SequenceError`] into an actionable `SessionError`. diff --git a/crates/onnx-runtime-session/src/executor/state.rs b/crates/onnx-runtime-session/src/executor/state.rs index fb2c1a39ab..5e518f162a 100644 --- a/crates/onnx-runtime-session/src/executor/state.rs +++ b/crates/onnx-runtime-session/src/executor/state.rs @@ -267,6 +267,15 @@ pub(crate) struct Executor { pub(super) decode_view_plan_disabled: bool, /// Master switch for graph-level compute-in-place aliasing. pub(super) compute_in_place_enabled: bool, + /// Release an intermediate value's buffer once its last consumer has run. + /// + /// Opt-in per session, and default OFF, because freeing lets a later run + /// place the same value at a different address. A session that records a + /// device graph needs its buffer addresses to stay put between the recording + /// run and every replay, so capture-eligible sessions (the decoder) leave + /// this alone. Prompt-phase component graphs never capture and are where the + /// resident set actually hurts, so they turn it on. + pub(super) release_dead_values_enabled: bool, /// Successful dead-input buffer aliases, retained for parity/safety tests. pub(super) compute_in_place_alias_count: u64, /// Opt-in (default OFF) master switch for the single-trip `Scan` inline @@ -781,6 +790,9 @@ pub(crate) struct ChildExecutor { pub(super) runs: u64, /// Shared trace context, propagated to every compiled child plan's executor. pub(super) trace: TraceContext, + /// Dead-value release setting, propagated to every compiled child plan's + /// executor so a Scan/Loop body inherits the parent session's choice. + pub(super) release_dead_values: bool, } #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] diff --git a/crates/onnx-runtime-session/src/lib.rs b/crates/onnx-runtime-session/src/lib.rs index 6765d15926..858f1cfa2e 100644 --- a/crates/onnx-runtime-session/src/lib.rs +++ b/crates/onnx-runtime-session/src/lib.rs @@ -1486,6 +1486,17 @@ impl InferenceSession { Ok(self.exec.adopt_memory_governor(governor, tier, holder)?) } + /// Release each intermediate value's buffer once its last consumer has run. + /// + /// Off by default: a session that records a device graph needs stable buffer + /// addresses across the recording run and its replays, and freeing lets the + /// allocator hand out a different address next time. Enable it on graphs that + /// never capture -- a prompt-phase vision encoder is the motivating case, + /// where holding all 2545 intermediates at once cost ~20x the live set. + pub fn set_release_dead_values(&mut self, enabled: bool) { + self.exec.set_release_dead_values(enabled); + } + pub fn set_weight_residency_budget(&self, budget_bytes: u64) -> Result> { Ok(self.exec.set_weight_residency_budget(budget_bytes)?) }