CUDA EP: Scan carry state buffer allocation fails with symbolic dim_params in inlined functions (shape mismatch: {1,16,1,1} != {1,16,128,128})

@justinchuby is already working on this.

Since Jun 1, 2026.

• #28735 by @copilot-swe-agent — open

## Bug: CUDA EP Scan carry state buffer allocation fails with symbolic dim_params in inlined functions

### Summary

When a `Scan` op is inside an **inlined ONNX local function**, the CUDA EP memory planner fails to resolve symbolic `dim_param` annotations in the Scan body's carry state (loop state variable) shapes. Unresolved dims fall back to `dim_value=0`, which the buffer allocator treats as `1`, producing undersized carry state buffers. The model then fails with a shape mismatch at runtime:

```
Shape mismatch attempting to re-use buffer. {1,16,1,1} != {1,16,128,128}
```

CPU EP handles this correctly (it resolves carry shapes from the actual input tensor at runtime).

### Environment

- ORT version: 1.24.x (and likely earlier)
- Platform: Linux, CUDA EP

### Minimal repro

The following model has a local function `RunScan` that contains a Scan op with symbolic dims in the body carry state:

```python
import numpy as np
import onnx
import onnx.helper as oh
import onnxruntime as ort

# Scan body with symbolic dims in carry state
body_carry_in = oh.make_tensor_value_info("bci", onnx.TensorProto.FLOAT, ["B", "H", "d_k"])
body_scan_in = oh.make_tensor_value_info("bsi", onnx.TensorProto.FLOAT, ["B", "H", "d_k"])
body_carry_out = oh.make_tensor_value_info("bco", onnx.TensorProto.FLOAT, ["B", "H", "d_k"])
body_scan_out = oh.make_tensor_value_info("bso", onnx.TensorProto.FLOAT, ["B", "H", "d_k"])
body_graph = oh.make_graph(
[oh.make_node("Add", ["bci","bsi"], ["bco"]),
oh.make_node("Identity", ["bco"], ["bso"])],
"body", [body_carry_in, body_scan_in], [body_carry_out, body_scan_out])

fn_scan = oh.make_node("Scan", ["fn_init","fn_seq"], ["fn_final","fn_seq_out"],
num_scan_inputs=1, body=body_graph, scan_input_axes=[1], scan_output_axes=[1])

fn_proto = oh.make_function(
domain="local", fname="RunScan",
inputs=["fn_init","fn_seq"], outputs=["fn_final","fn_seq_out"],
nodes=[fn_scan], opset_imports=[oh.make_opsetid("",18)])

outer_graph = oh.make_graph(
[oh.make_node("RunScan", ["init","seq_in"], ["final","seq_out"], domain="local")],
"outer",
[oh.make_tensor_value_info("init", onnx.TensorProto.FLOAT, ["B","H","d_k"]),
oh.make_tensor_value_info("seq_in", onnx.TensorProto.FLOAT, ["B","seq_len","H","d_k"])],
[oh.make_tensor_value_info("final", onnx.TensorProto.FLOAT, ["B","H","d_k"]),
oh.make_tensor_value_info("seq_out", onnx.TensorProto.FLOAT, ["B","seq_len","H","d_k"])],
)
model = oh.make_model(outer_graph,
opset_imports=[oh.make_opsetid("",18), oh.make_opsetid("local",1)],
functions=[fn_proto])
model.ir_version = 10

B, seq_len, H, d_k = 1, 4, 16, 128
init_v = np.zeros((B, H, d_k), dtype=np.float32)
seq_v = np.ones((B, seq_len, H, d_k), dtype=np.float32)

# CPU: works
sess_cpu = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
print("CPU PASS:", sess_cpu.run(None, {"init": init_v, "seq_in": seq_v})[0].shape)

# CUDA: shape mismatch error
sess_cuda = ort.InferenceSession(model.SerializeToString(),
providers=["CUDAExecutionProvider","CPUExecutionProvider"])
print("CUDA:", sess_cuda.run(None, {"init": init_v, "seq_in": seq_v})[0].shape)
```

Expected: both pass, `final.shape == (1, 16, 128)`.
Actual on CUDA EP: shape mismatch error mentioning a buffer with `1` where `128` should be.

### Root cause analysis

The failure chain:

1. ORT inlines the `local.RunScan` function at load time (the inlined Scan node gets prefix `_inlfunc_RunScan_...`).
2. After inlining, the Scan body's carry state (loop state variable) input retains `dim_param` annotations (`"H"`, `"d_k"`) from the original function body.
3. The **CUDA EP memory planner** needs to pre-allocate carry state buffers before executing the Scan. It reads the static shape annotations from the Scan body to determine buffer size.
4. Since `dim_param` values cannot be resolved statically after inlining (no concrete outer-scope binding for `H` or `d_k`), `GetTensorShapeFromTensorShapeProto` returns `-1` for those dims.
5. The memory planner appears to treat unresolvable symbolic dims as `dim_value=0`, then as `1` during allocation.
6. The resulting buffer (e.g., `{1, 16, 1, 1}`) is too small for the actual carry state shape (`{1, 16, 128, 128}`), causing a shape mismatch on the first iteration.

**CPU EP** does not pre-allocate carry buffers statically — `OutputIterator::Initialize()` (in `scan_utils.cc`) resolves the carry state shape from the actual input tensor at runtime via `MakeShapeConcrete`, which is why CPU EP works correctly.

### Relevant code paths

- `onnxruntime/core/providers/cpu/controlflow/scan_utils.cc`, `OutputIterator::Initialize()` (lines ~448-456): CPU lazy resolution from input tensor shape — this is what's missing on CUDA.
- `onnxruntime/core/framework/execution_frame.cc`, `AllocateTensorWithPreAllocateBufferHelper`: The "Shape mismatch attempting to re-use buffer" error is thrown here.
- `onnxruntime/core/graph/function_utils.cc`, `Specialize()`: Function inlining — correctly renames values but does not substitute concrete dims into body `dim_param` annotations.

### Suggested fix direction

In the CUDA EP memory planner, when planning carry state buffer sizes for an inlined Scan node, fall back to lazy allocation (as CPU EP does) if the body's carry state shape contains unresolvable symbolic dims after inlining.

Alternatively, during function inlining in `Graph::InlineFunction` / `Specialize`, propagate the concrete outer-graph input shapes into the Scan body's carry state `dim_param` annotations so the memory planner can see concrete values.

### Workaround (applied in consuming library)

Use concrete integer dimensions instead of `dim_param` strings for all model-constant dimensions (num_heads, head_dim, etc.) in Scan body shapes. Only truly dynamic dimensions (batch size, sequence length) should remain symbolic. See [mobius PR #57](https://github.com/onnxruntime/mobius/pull/57) for an example.