Free executor intermediates once their last consumer has run

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(),

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()));

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]

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.

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