fix(session): expose bound inputs to control flow

.squad/decisions/inbox/mary-384-value1414.md

@@ -0,0 +1,4 @@
+### 2026-07-30: Carry persistent device state through control flow
+**By:** Mary
+**What:** `value#1414` is the graph input `past_key_values.0.recurrent_state`, a persistent CUDA binding captured by the `Scan` inside the inlined `com.microsoft::LinearAttention` function. Control-flow capture preparation originally ignored external bindings, so the executor could not read this state. After making bound inputs readable, a second correctness bug remained: `If`/`Loop`/`Scan` outputs were always stored in executor-owned buffers, bypassing persistent output bindings. The `present.*.recurrent_state` output therefore never replaced the paired `past.*` state, causing native decode to repeat the prompt. Control-flow output storage now reuses the general external-output-aware tensor storage path. The regression binds an `If` input and output to one device allocation, runs twice, and verifies `[5, 9] -> [6, 10] -> [7, 11]`.
+**Why:** Persistent device I/O must be honored at every control-flow materialization and storage boundary. The fix is general across externally bound `If`, `Loop`, and `Scan` values, independent of tensor name, rank, or model. Qwen3.6-27B native CUDA now exactly matches the trusted ORT-CPU reference for the first 16 greedy tokens on `The capital of France is`: both produce `[11751, 13, 271, 248068, 271, 248069, 271, 4639, 369, 4252, 13, 11751, 369, 279, 6511, 321]`, decoding to ` Paris.\n\n<think>\n\n</think>\n\nThat is correct. Paris is the capital and`. Native measured 5.26 tok/s (15-token steady window after one skipped token); ORT-CPU measured 1.67 tok/s. Verdict: **FIXED — persistent recurrent state now carries between decode steps and token parity is proven.**

crates/onnx-runtime-session/src/executor/control_flow.rs

@@ -234,6 +234,7 @@ impl Executor {
         &self,
         vid: ValueId,
         resolved: &HashMap<ValueId, Vec<usize>>,
+        external: &ExternalBindings,
     ) -> Result<Tensor> {
         let dtype = self.value_dtypes[&vid];
         let shape = resolved.get(&vid).cloned().ok_or_else(|| {
@@ -247,6 +248,28 @@ impl Executor {
                 op: "<control-flow input>".to_string(),
             }
         })?;
+        if let Some(value) = external
+            .inputs
+            .get(&vid)
+            .or_else(|| external.outputs.get(&vid))
+        {
+            let n = checked_storage_bytes(
+                dtype,
+                checked_numel(&shape, || format!("value#{}", vid.0))?,
+                || format!("value#{}", vid.0),
+                &shape,
+            )?;
+            if value.dtype != dtype || value.len < n {
+                return Err(SessionError::Internal(format!(
+                    "external binding for value#{} cannot materialize {:?} {:?} (binding: {:?}, {} bytes)",
+                    vid.0, dtype, shape, value.dtype, value.len
+                )));
+            }
+            let buffer = value.readable_buffer()?;
+            let mut bytes = vec![0_u8; n];
+            self.ep.copy_to_host(&buffer, &mut bytes)?;
+            return Tensor::from_raw(dtype, shape, &bytes);
+        }
         // A view value owns no buffer; materialize its strided bytes contiguous.
         let bytes = self.contiguous_bytes(vid, &shape, dtype)?;
         Tensor::from_raw(dtype, shape, &bytes)
@@ -394,7 +417,7 @@ impl Executor {
             let buf = self
                 .buffers
                 .get(&vid)
-                .ok_or_else(|| SessionError::Internal(format!("value#{} not produced", vid.0)))?;
+                .ok_or_else(|| SessionError::Internal(format!("{} not produced", value_name())))?;
             let mut host = vec![0u8; n];
             self.ep.copy_to_host(buf, &mut host)?;
             Ok(host)
@@ -412,13 +435,15 @@ impl Executor {
         vid: ValueId,
         tensor: &Tensor,
         resolved: &mut HashMap<ValueId, Vec<usize>>,
+        external: &ExternalBindings,
     ) -> Result<()> {
         self.store_output_bytes(
             vid,
             tensor.dtype,
             tensor.shape.clone(),
             tensor.as_bytes(),
             resolved,
+            external,
         )
     }
 
@@ -429,31 +454,9 @@ impl Executor {
         dims: Vec<usize>,
         bytes: &[u8],
         resolved: &mut HashMap<ValueId, Vec<usize>>,
+        external: &ExternalBindings,
     ) -> Result<()> {
-        let numel = checked_numel(&dims, || format!("value#{}", vid.0))?;
-        let need =
-            checked_storage_bytes(dtype, numel, || format!("value#{}", vid.0), &dims)?.max(1);
-        let fits = self
-            .buffers
-            .get(&vid)
-            .map(|b| b.len() == need)
-            .unwrap_or(false);
-        if !fits {
-            if let Some(old) = self.buffers.remove(&vid) {
-                self.ep.deallocate(old)?;
-            }
-            self.shared_buffers.remove(&vid);
-            let buf = self
-                .ep
-                .allocate(need, TensorLayout::contiguous().alignment)?;
-            self.buffers.insert(vid, buf);
-        }
-        let buf = self.buffers.get_mut(&vid).expect("just ensured");
-        self.ep.copy_from_host(bytes, buf)?;
-        self.value_dtypes.insert(vid, dtype);
-        self.buffer_shapes.insert(vid, dims.clone());
-        resolved.insert(vid, dims);
-        Ok(())
+        self.store_raw_tensor_output(vid, dtype, dims, bytes, resolved, external)
     }
 
     /// Prepare one selected control-flow subgraph and materialize only the free
@@ -466,6 +469,7 @@ impl Executor {
         attr_key: &str,
         resolved: &HashMap<ValueId, Vec<usize>>,
         outer_scope: &HashMap<String, Tensor>,
+        external: &ExternalBindings,
     ) -> Result<PreparedSubgraph> {
         let key = (node_id, attr_key.to_string());
         let body = self.graph.subgraphs.get(&key).ok_or_else(|| {
@@ -482,9 +486,11 @@ impl Executor {
             let tensor = if let Some(&vid) = self.name_index.get(&name) {
                 let materialized = self.buffers.contains_key(&vid)
                     || self.views.contains_key(&vid)
-                    || self.seq_elem_values.contains_key(&vid);
+                    || self.seq_elem_values.contains_key(&vid)
+                    || external.inputs.contains_key(&vid)
+                    || external.outputs.contains_key(&vid);
                 if resolved.contains_key(&vid) && materialized {
-                    self.value_tensor(vid, resolved)?
+                    self.value_tensor(vid, resolved, external)?
                 } else {
                     outer_scope
                         .get(&name)
@@ -552,12 +558,13 @@ impl Executor {
         pi: usize,
         resolved: &mut HashMap<ValueId, Vec<usize>>,
         outer_scope: &HashMap<String, Tensor>,
+        external: &ExternalBindings,
     ) -> Result<()> {
         let node = self.graph.node(self.plan[pi].node_id).clone();
         match node.op_type.as_str() {
-            "If" => self.exec_if(&node, resolved, outer_scope),
-            "Loop" => self.exec_loop(&node, resolved, outer_scope),
-            "Scan" => self.exec_scan(&node, resolved, outer_scope),
+            "If" => self.exec_if(&node, resolved, outer_scope, external),
+            "Loop" => self.exec_loop(&node, resolved, outer_scope, external),
+            "Scan" => self.exec_scan(&node, resolved, outer_scope, external),
             other => Err(SessionError::Internal(format!(
                 "exec_control_flow reached non-control-flow op {other:?}"
             ))),
@@ -571,6 +578,7 @@ impl Executor {
         node: &Node,
         resolved: &mut HashMap<ValueId, Vec<usize>>,
         outer_scope: &HashMap<String, Tensor>,
+        external: &ExternalBindings,
     ) -> Result<()> {
         {
             let then_branch = self
@@ -612,7 +620,7 @@ impl Executor {
                     op: "If".to_string(),
                     reason: "missing required 'cond' input".to_string(),
                 })?;
-        let cond_t = self.value_tensor(cond_vid, resolved)?;
+        let cond_t = self.value_tensor(cond_vid, resolved, external)?;
         if cond_t.dtype != DataType::Bool {
             return Err(SessionError::DtypeMismatch {
                 name: "If cond".to_string(),
@@ -655,7 +663,7 @@ impl Executor {
             .unwrap_or(false);
         let prepared = {
             let _s = phase_span!("execif.prepare_subgraph");
-            self.prepare_subgraph(node.id, attr_key, resolved, outer_scope)?
+            self.prepare_subgraph(node.id, attr_key, resolved, outer_scope, external)?
         };
         let outs = {
             let _s = phase_span!("execif.run_subgraph");
@@ -672,7 +680,7 @@ impl Executor {
         {
             let _s = phase_span!("execif.store_output");
             for (vid, t) in node.outputs.iter().zip(outs.iter()) {
-                self.store_output_tensor(*vid, t, resolved)?;
+                self.store_output_tensor(*vid, t, resolved, external)?;
             }
         }
         // Only enable future skips when the taken branch is loop-invariant.
@@ -814,12 +822,13 @@ impl Executor {
         node: &Node,
         resolved: &mut HashMap<ValueId, Vec<usize>>,
         outer_scope: &HashMap<String, Tensor>,
+        external: &ExternalBindings,
     ) -> Result<()> {
         // Inputs: [M, cond, v_initial...]. M and cond may be omitted (None slot)
         // or an empty-name optional; absence means "unbounded" / "true".
         let m: Option<i64> = match node.inputs.first().and_then(|s| *s) {
             Some(vid) => {
-                let t = self.value_tensor(vid, resolved)?;
+                let t = self.value_tensor(vid, resolved, external)?;
                 if t.dtype != DataType::Int64 {
                     return Err(SessionError::DtypeMismatch {
                         name: "Loop M".to_string(),
@@ -841,7 +850,7 @@ impl Executor {
         let mut cond: Option<bool> =
             match node.inputs.get(1).and_then(|s| *s) {
                 Some(vid) => {
-                    let t = self.value_tensor(vid, resolved)?;
+                    let t = self.value_tensor(vid, resolved, external)?;
                     if t.dtype != DataType::Bool {
                         return Err(SessionError::DtypeMismatch {
                             name: "Loop cond".to_string(),
@@ -870,7 +879,7 @@ impl Executor {
                         .to_string(),
                 )
             })?;
-            carried.push(self.value_tensor(vid, resolved)?);
+            carried.push(self.value_tensor(vid, resolved, external)?);
         }
         let num_carried = carried.len();
         let carried_invariants: Vec<(DataType, Vec<usize>)> = carried
@@ -891,7 +900,7 @@ impl Executor {
         let mut scan_acc: Vec<TensorStackAccumulator> = (0..num_scan)
             .map(|_| TensorStackAccumulator::new())
             .collect();
-        let prepared = self.prepare_subgraph(node.id, "body", resolved, outer_scope)?;
+        let prepared = self.prepare_subgraph(node.id, "body", resolved, outer_scope, external)?;
         let mut iter_tensor = scalar_i64_tensor(0)?;
         let mut cond_tensor = scalar_bool_tensor(cond.unwrap_or(true))?;
 
@@ -972,7 +981,7 @@ impl Executor {
 
         // Emit outputs: carried finals, then stacked scan outputs.
         for (i, t) in carried.iter().enumerate() {
-            self.store_output_tensor(node.outputs[i], t, resolved)?;
+            self.store_output_tensor(node.outputs[i], t, resolved, external)?;
         }
         for (s, (acc, empty_spec)) in scan_acc.into_iter().zip(empty_scan_specs).enumerate() {
             let (dtype, shape, bytes) = acc.finish_with_empty(empty_spec, s)?;
@@ -982,6 +991,7 @@ impl Executor {
                 shape,
                 &bytes,
                 resolved,
+                external,
             )?;
         }
         Ok(())
@@ -1142,6 +1152,7 @@ impl Executor {
         node: &Node,
         resolved: &mut HashMap<ValueId, Vec<usize>>,
         outer_scope: &HashMap<String, Tensor>,
+        external: &ExternalBindings,
     ) -> Result<()> {
         let raw_num_scan_inputs = node
             .attr("num_scan_inputs")
@@ -1208,15 +1219,15 @@ impl Executor {
                 reason: "an initial-state input is omitted (empty), which ONNX does not allow"
                     .to_string(),
             })?;
-            state.push(self.value_tensor(vid, resolved)?);
+            state.push(self.value_tensor(vid, resolved, external)?);
         }
         let mut scan_inputs: Vec<Tensor> = Vec::with_capacity(num_scan_inputs);
         for slot in node.inputs.iter().skip(num_state) {
             let vid = slot.ok_or_else(|| SessionError::ControlFlow {
                 op: "Scan".to_string(),
                 reason: "a scan input is omitted (empty), which ONNX does not allow".to_string(),
             })?;
-            scan_inputs.push(self.value_tensor(vid, resolved)?);
+            scan_inputs.push(self.value_tensor(vid, resolved, external)?);
         }
 
         let mut input_axes = Vec::with_capacity(num_scan_inputs);
@@ -1262,7 +1273,7 @@ impl Executor {
         let mut scan_acc: Vec<TensorStackAccumulator> = (0..num_scan_outputs)
             .map(|_| TensorStackAccumulator::new())
             .collect();
-        let prepared = self.prepare_subgraph(node.id, "body", resolved, outer_scope)?;
+        let prepared = self.prepare_subgraph(node.id, "body", resolved, outer_scope, external)?;
         let mut scan_slices = Vec::with_capacity(num_scan_inputs);
         if trip_count != 0 {
             for (index, ((input, &axis), &direction)) in scan_inputs
@@ -1339,7 +1350,7 @@ impl Executor {
         }
 
         for (i, t) in state.iter().enumerate() {
-            self.store_output_tensor(node.outputs[i], t, resolved)?;
+            self.store_output_tensor(node.outputs[i], t, resolved, external)?;
         }
         for (s, ((acc, empty_spec), (&axis, &direction))) in scan_acc
             .into_iter()
@@ -1348,7 +1359,14 @@ impl Executor {
             .enumerate()
         {
             let (dtype, shape, bytes) = acc.finish_scan(axis, direction, empty_spec, s)?;
-            self.store_output_bytes(node.outputs[num_state + s], dtype, shape, &bytes, resolved)?;
+            self.store_output_bytes(
+                node.outputs[num_state + s],
+                dtype,
+                shape,
+                &bytes,
+                resolved,
+                external,
+            )?;
         }
         Ok(())
     }

crates/onnx-runtime-session/src/executor/dispatch.rs

@@ -86,7 +86,7 @@ impl Executor {
             (is_control_flow, is_sequence, span)
         };
         if is_control_flow {
-            self.exec_control_flow(pi, resolved, outer_scope)
+            self.exec_control_flow(pi, resolved, outer_scope, external)
         } else if is_sequence {
             self.exec_sequence_node(pi, resolved, external)
         } else {

crates/onnx-runtime-session/src/executor/state.rs

@@ -580,6 +580,20 @@ impl ExternalValue {
         }
         .ok_or_else(|| SessionError::Internal("external output binding has a null pointer".into()))
     }
+
+    pub(super) fn readable_buffer(&self) -> Result<DeviceBuffer> {
+        if self.ptr.is_null() {
+            return Err(SessionError::Internal(
+                "external input binding has a null pointer".into(),
+            ));
+        }
+        // SAFETY: `prepare_external_bindings` obtains this pointer from a live
+        // `DeviceIoBinding` borrowed for the complete run. This read-only alias
+        // neither owns nor mutates the binding's allocation.
+        Ok(unsafe {
+            DeviceBuffer::from_borrowed_parts(self.ptr, self.device, self.len, self.alignment)
+        })
+    }
 }
 
 #[derive(Default)]

crates/onnx-runtime-session/tests/control_flow.rs

@@ -324,6 +324,40 @@ fn if_never_memoizes_branch_that_reads_changing_captures() {
     }
 }
 
+#[test]
+fn if_materializes_outer_capture_from_persistent_device_input() {
+    let mut g = new_parent();
+    let cond = input(&mut g, "cond", DataType::Bool, &[1]);
+    let _state = input(&mut g, "X", DataType::Float32, &[2]);
+    let y = g.create_named_value("Y", DataType::Float32, static_shape([2]));
+    let node = control_flow_node(&mut g, "If", vec![Some(cond)], vec![y], &[]);
+    register(&mut g, node, "then_branch", if_branch("Add"));
+    register(&mut g, node, "else_branch", if_branch("Sub"));
+    g.add_output(y);
+
+    let mut session = InferenceSession::from_graph(g).expect("build session");
+    let mut state = session
+        .allocate_device_binding("X", Some("Y"), DataType::Float32, vec![2], vec![2])
+        .expect("state binding");
+    state
+        .write_bytes(0, &f32_bytes(&[5.0, 9.0]))
+        .expect("seed state");
+    let cond_t = Tensor::from_raw(DataType::Bool, vec![1], &[1]).unwrap();
+    for expected in [[6.0, 10.0], [7.0, 11.0]] {
+        let outputs = session
+            .run_with_device_bindings(&[("cond", &cond_t)], std::slice::from_mut(&mut state))
+            .expect("run with captured persistent state");
+        assert!(outputs[0].is_none(), "bound output must stay on device");
+        let values = state
+            .read_bytes()
+            .unwrap()
+            .chunks_exact(4)
+            .map(|bytes| f32::from_le_bytes(bytes.try_into().unwrap()))
+            .collect::<Vec<_>>();
+        assert_eq!(values, expected);
+    }
+}
+
 #[test]
 fn if_rejects_mismatched_branch_output_counts_before_running_selected_branch() {
     let mut g = new_parent();