Skip to content

ONNX: support tensors in CUDA device memory, IO binding and CUDA provider options - #281

Open
larissakl wants to merge 4 commits into
masterfrom
onnx-device-memory
Open

larissakl wants to merge 4 commits into
masterfrom
onnx-device-memory

Conversation

@larissakl

Copy link
Copy Markdown
Contributor

This PR prepares the ONNX module for a transformer state manager that keeps its KV cache states on the GPU. As long as all inputs and outputs are in host memory, Session::run uses the previous code path, so existing setups are not effected.

Onnx::Value can now hold tensors in CUDA device memory. New helpers check where a tensor is, copy it to the host or wrap existing device memory. Host accessors abort with a clear message when called on a device tensor.
Session::run gets an overload with a memory location per output. If any input or output is on the device, the run uses an Ort::IoBinding. By default, a state output stays on the device if its state input came from the device and everything else goes to the host. This way label scorers and LMs need no changes once a state manager provides device inputs.
All CUDA sessions of a process now share one CUDA stream, so other code can write device inputs on the same stream without extra synchronization. Note that each run synchronizes this stream, so sessions used from several threads wait for each other's GPU work. The new parameter cuda-use-tf32 allows disabling TF32. Requires onnxruntime >= 1.16.

Depends on #279.

Base automatically changed from onnx-value-noexcept to master October 9, 2026 12:17
@larissakl
larissakl marked this pull request as ready for review October 16, 2026 12:13
@larissakl
larissakl requested review from SimBe195 and curufinwe and removed request for curufinwe November 2, 2026 12:49

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant