From 71c4fcf653e347b4655023c155f1a8cc036fcb3f Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 19 Aug 2026 20:58:37 +0000
Subject: [PATCH 1/2] Report the full error chain and close four kernel dtype
 gaps

A native CUDA pipeline run that failed anywhere below the driver surfaced
as "native CUDA decoder forward pass failed" and nothing else. anyhow's
Display prints only the outermost context, so every layer of diagnosis
the lower crates had carefully attached was discarded at the boundary
where a human would read it. Switching the two driver formats to the
alternate `{err:#}` selector prints the whole chain.

With the chain visible, four real dtype gaps became legible instead of
being guessed at:

  * Reduce* on CPU accepted Int64 only for Sum. A graph reducing Int32,
    or taking Max/Min/Prod over integers, hit the dtype guard rather
    than the integer path that was already sitting there.
  * SkipLayerNormalization on CPU required its optional `mean` and
    `inv_std_var` outputs to carry the input dtype. ORT's schema types
    them `float`, but exporters routinely emit them in the activation
    dtype instead, and the write already narrows to whichever was
    declared -- so both should be accepted. The bf16 skip-RMSNorm path
    documents the same allowance.
  * SkipLayerNormalization on CUDA made the mirror-image assumption,
    deriving the stat pointers' dtype from the input rather than from
    the declared output. Threading the resolved stat dtype through to
    the kernel keeps a float32 stat pair on a bf16 input from being
    written as bf16.
  * Range on CPU had no Int32 branch.

Each gap gets tests at the boundary that was previously unreachable.

Note for reviewers: `onnx-runtime-ep-cuda --lib` has five failures on
this machine (two in graph::tests, three in matmul_nbits fp16 tests).
They reproduce identically on an unmodified main, so they are not from
this change.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 0190e2eb-abe4-451f-b36d-44a035a99b7e
---
 crates/onnx-genai-server/src/driver.rs        |   8 +-
 .../src/kernels/contrib_fused.rs              |  93 +++++++++++-
 .../src/kernels/reduce_ops.rs                 | 136 ++++++++++++++++--
 .../src/kernels/sequence.rs                   |  44 ++++++
 .../src/kernels/normalization.rs              |  40 +++++-
 5 files changed, 295 insertions(+), 26 deletions(-)

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::<Vec<_>>();
         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::<Vec<_>>();
+        let kept = numel(&kept_shape);
+        let bytes = match output.dtype {
+            DataType::Int64 => acc
+                .iter()
+                .take(kept)
+                .flat_map(|value| value.to_le_bytes())
+                .collect::<Vec<_>>(),
+            DataType::Int32 => acc
+                .iter()
+                .take(kept)
+                .flat_map(|value| (*value as i32).to_le_bytes())
+                .collect::<Vec<_>>(),
+            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<u8> = 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::<i32>(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<u8> = 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::<i64>() 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(

From 07d947f6e8fc8b79ee7867a0210cfe794ce4d143 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 19 Aug 2026 21:00:49 +0000
Subject: [PATCH 2/2] Free executor intermediates once their last consumer has
 run

The executor already computed `last_use` for every value, but the only
consumer was `inplace_dead_inputs`, whose conditions are narrow enough
(output must match an input in both dtype and shape) that in practice it
covers elementwise chains and nothing else. Any operator that changes
shape -- MatMul, Reshape, Attention -- allocated a buffer that then sat
resident until the run ended. A graph is therefore charged for its peak
node count rather than its peak live set.

The vision encoder made the difference concrete. Its 2545 node outputs
were all held simultaneously, costing roughly 23 GB at 448px and making
full-resolution images unservable: a 960x686 image OOMed an 80 GB card.
With dead values freed the same run peaks at 55.8 GB and completes.

Four details worth calling out, because each is easy to get wrong:

`DeviceBuffer` is a bare handle with no `Drop`. Dropping it from the
value map looks like a release and is actually a leak -- the first
version of this change did exactly that and reported 7075 releases while
device memory did not move at all. The buffer has to go back through
`ep.deallocate`.

`buffer_shapes` has to be cleared alongside it. `ensure_buffer` treats a
surviving shape entry as proof the allocation is still there and skips
sizing it, so leaving one behind hands the next run a value it believes
is backed and is not. Every other path that takes a buffer out of the
map does the same thing for the same reason.

A memoized loop-invariant `If` skips its branch on later runs and serves
its outputs straight from the buffers an earlier run left resident, and
that memo survives an eager run. Freeing one of those outputs turns the
next skip into a missing-buffer error on the *second* request rather
than the first. `try_move_host_output` already declines for exactly this
reason; this pass is the second buffer-stealing path and now carries the
same guard.

The release is per-session and off by default because it is not safe
under device graph capture. Capture bakes device pointers into the
recorded graph, so freeing a buffer lets a later allocation land at a
different address and replay then reads the wrong memory; the decoder
fails immediately with a capture validation violation. Restricting this
to the eager path is not sufficient either, since one session runs eager
before it captures. Components that never capture -- the pipeline
components, which run once per request -- opt in; the decoder does not.

Aliased values (`pinned`) are never released, even after every alias has
died. That is deliberate: `pinned` is a monotone record of "was ever an
alias source", and releasing on the exact liveness test instead would be
wrong for the decode memo, which restores a step's view table on the
next run and expects the source buffer to still be resident. An earlier
draft did try to reclaim view sources once their views died; it turned
out to be unreachable (every view source is pinned) and, when measured,
worth 36 MiB out of ~12 GB. Not worth a cross-run aliasing analysis.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 0190e2eb-abe4-451f-b36d-44a035a99b7e
---
 .../onnx-genai-engine/src/native_component.rs |  8 ++-
 .../src/executor/build.rs                     | 33 ++++++++++
 .../src/executor/control_flow.rs              | 12 ++++
 .../src/executor/dispatch.rs                  | 66 +++++++++++++++++++
 .../onnx-runtime-session/src/executor/mod.rs  | 10 +++
 .../src/executor/state.rs                     | 12 ++++
 crates/onnx-runtime-session/src/lib.rs        | 11 ++++
 7 files changed, 151 insertions(+), 1 deletion(-)

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-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<ValueId> = 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<WeightStore> {
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<CompiledChildPlan> {
         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<bool>,
     /// Lazy weight inputs this node may ask the EP to page at dispatch time.
     pub lazy_weight_inputs: Vec<ValueId>,
+    /// 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<ValueId>,
 }
 
 /// 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<Option<u64>> {
         Ok(self.exec.set_weight_residency_budget(budget_bytes)?)
     }