Skip to content

Release captured graph resources when generator is destroyed - #2106

Merged
kunal-vaishnavi merged 15 commits into
mainfrom
per-graph-buffer-manager-webgpu
Jun 30, 2026
Merged

Release captured graph resources when generator is destroyed#2106
kunal-vaishnavi merged 15 commits into
mainfrom
per-graph-buffer-manager-webgpu

Conversation

@qjia7

@qjia7 qjia7 commented Apr 29, 2026

Copy link
Copy Markdown
Contributor

Call ReleaseGraph on the ORT session in State::~State() to free per-graph buffer resources in GPU. This pairs with the per-graph buffer manager changes in ORT microsoft/onnxruntime#28260.

Changes:

  • Add OrtSession::ReleaseGraph wrapper calling the new C API
  • Store graph_capture_session_ pointer in State during graph capture runs
  • Call ReleaseGraph in State destructor with the graph annotation ID

qjia7 added a commit to microsoft/onnxruntime that referenced this pull request May 28, 2026
)

### Summary
Add per-annotation-ID buffer managers and captured command storage so
multiple generators can each capture and replay their own graph
independently without cross-contamination
Add ReleaseGraph API through the full ORT stack (EP base → C API →
InferenceSession → plugin EP) to release captured commands and GPU
buffers when a generator is destroyed
Replace the single graph_buffer_mgr_ / is_graph_captured_ bool with
per_graph_buffer_mgrs_ map and captured_graph_ids_ set keyed by
annotation ID
Use a std::function getter with cached pointer pattern in
GpuBufferAllocator to dynamically route allocations to the active
per-graph buffer manager during runs, while keeping Alloc/Free as simple
pointer dereferences
### Motivation
Edge's Prompt API speed benchmark creates multiple sessions/generators
sequentially with graph capture enabled. With the existing single-graph
design, the second generator replays the first generator's captured
commands with wrong buffers, producing incorrect output and ultimately a
QuotaExceededError in the browser. This PR isolates each generator's
graph capture state so they don't interfere with each other.

### Related PR
The GenAI side change is in
microsoft/onnxruntime-genai#2106, which calls
SessionReleaseGraph when a generator is destroyed to release the
captured graph's GPU buffers.

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: qjia7 <4221210+qjia7@users.noreply.github.com>
qjia7 added 2 commits June 16, 2026 15:09
Call ReleaseGraph on the ORT session in State::~State() to free per-graph
buffer resources in GPU. This pairs with the per-graph buffer manager changes in ORT.

Changes:
- Add OrtSession::ReleaseGraph wrapper calling the new C API
- Store graph_capture_session_ pointer in State during graph capture runs
- Call ReleaseGraph in State destructor with the graph annotation ID
Update ORT API bindings to match the renamed C API function
in onnxruntime PR #28260.
@qjia7
qjia7 force-pushed the per-graph-buffer-manager-webgpu branch from 88b5202 to 3246c53 Compare June 16, 2026 09:12
@qjia7
qjia7 requested a review from Copilot June 16, 2026 09:16
The genai WebGPU branch needs the SessionReleaseCapturedGraph C API
that landed in onnxruntime PR #28260. The 1.27.0 stable release is
built from exactly that commit and predates the MatMulNBitsMlpFusion
transformer (PR #28280) that emits a contrib op not yet registered
by the v0.1.x WebGPU plugin EP installed on most developer hosts.
@qjia7
qjia7 force-pushed the per-graph-buffer-manager-webgpu branch from 3246c53 to bbfb909 Compare June 16, 2026 09:19

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR adds lifecycle cleanup for ONNX Runtime graph-capture resources by releasing captured graph buffers when a Generators::State is destroyed, aligning GenAI teardown with new ORT per-graph buffer manager behavior.

Changes:

  • Added an OrtSession::ReleaseCapturedGraph(int) wrapper around the new ORT C API.
  • Stored the session used during graph capture in State and invoked captured-graph release in State::~State().
  • Updated the auto-downloaded ORT package version to 1.27.0.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
src/models/onnxruntime_inline.h Adds the inline implementation of OrtSession::ReleaseCapturedGraph.
src/models/onnxruntime_api.h Declares OrtSession::ReleaseCapturedGraph in the local C++ wrapper.
src/models/model.h Stores the session pointer used for graph capture in State.
src/models/model.cpp Captures the session pointer during runs and releases captured-graph resources in State destructor.
cmake/ortlib.cmake Bumps default ORT version for auto-download to 1.27.0.

Comment thread src/models/model.cpp
Comment thread src/models/onnxruntime_inline.h
qjia7 and others added 3 commits June 16, 2026 18:06
Address review feedback and resolve the CI runtime-ORT version mismatch
that produced "The requested API version [27] is not available" on the
Windows CUDA and Windows CPU arm64 lanes.

Source changes:
- Wrap the ReleaseCapturedGraph call in ~State() in try/catch so a
  non-OK status from the EP cannot propagate as an exception out of a
  destructor and trigger std::terminate during unwinding.
- Guard the OrtSession::ReleaseCapturedGraph declaration, inline
  definition, and call site with #if ORT_API_VERSION >= 27 so builds
  using ORT_HOME pointing at an older ORT continue to compile.
- Store the graph annotation id as an int alongside the existing string
  form to avoid the std::atoi reparse in the destructor.
- Document the lifetime invariant of the non-owning graph_capture_session_
  pointer and the EP no-op contract for unknown graph ids.

Build/CI changes:
- Bump the win-cuda-x64 workflow's nuget-installed runtime ORT from
  1.26.0 to 1.27.0 to match the build-time FetchContent pin.
- Bump the pip-installed onnxruntime / onnxruntime-gpu in
  test/python/{cpu,cuda}/ort/requirements.txt and in the ADO
  integration-pytest-step.yml from 1.26.0 to 1.27.0 for the same reason.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
onnxruntime-gpu==1.27.0 is not yet published on PyPI, so the pip install
fails on every CI lane that installs the GPU wheel from PyPI. Roll the
three GPU pins back to 1.26.0 (the CPU package onnxruntime==1.27.0 is
already on PyPI and stays at 1.27.0).
ORT 1.27.0 NuGet for the GPU/Windows variant and the CPU PyPI wheel
have not propagated to every CI channel yet, so the runtime ORT in
CUDA and macOS lanes is still 1.26. Pin the build-time ORT to 1.26.0
as well so the genai DLL only binds API 26.

The SessionReleaseCapturedGraph code paths in model.cpp, model.h,
onnxruntime_api.h, and onnxruntime_inline.h remain in place; they are
gated by #if ORT_API_VERSION >= 27 and are elided when built against
1.26 headers, so no source revert is needed.
@qjia7
qjia7 marked this pull request as ready for review June 16, 2026 14:27
@qjia7
qjia7 requested a review from a team as a code owner June 16, 2026 14:27
Comment thread src/models/model.cpp
Comment thread src/models/model.cpp
The dtor swallowed exceptions from ReleaseCapturedGraph silently, making
EP-side cleanup failures invisible during debugging. Route the failure
through the existing ort_lib log channel so it surfaces when enabled,
while keeping the destructor non-throwing.
@qjia7

qjia7 commented Jun 18, 2026

Copy link
Copy Markdown
Contributor Author

It seems that the webgpu CI is not stable. The failure should not be related with this PR's changes.

@qjia7
qjia7 requested a review from baijumeswani June 18, 2026 05:44
qjia7 added 5 commits June 22, 2026 10:08
Now that ORT 1.27 is released, point the FetchContent default,
the Windows CUDA nightly install, the Win/Linux/macOS integration
pytest installs, and the CPU/CUDA python test requirements at
1.27.0. Plugin EP onnxruntime-ep-webgpu stays at 0.1.0 since
0.2.0 has not been published yet.
onnxruntime-gpu 1.27.0 links libcudart.so.13 (CUDA 13) but CI runners
ship CUDA 12.x, and onnxruntime-ep-webgpu 0.2.0 has not been published
yet (ORT 1.27 emits MatMulNBitsMlp via a new fusion transformer whose
kernel only exists in plugin EP 0.2.0+). Pin CUDA and WebGPU python
installs back to 1.26.0 while keeping the C++ build at 1.27.0; CPU
python installs remain on 1.27.0.
ORT 1.27 ships the MatMulNBitsMlpFusion transformer, whose kernel only
exists in plugin EP 0.2.0+. The 0.2.0.dev20260611 wheel from the
ORT-Nightly Azure feed pairs cleanly with ORT 1.27 GA, so install it on
both Windows and macOS WebGPU jobs (the wheel is not on PyPI, hence the
explicit --index-url).
ORT-GPU 1.27 nupkgs link the CUDA 13 runtime (libcudart.so.13 /
libcublasLt.so.13), but CI runners still ship CUDA 12.x, so the provider
shared library fails to load. Pin the USE_CUDA branch in ortlib.cmake to
1.26.0 and align the Windows CUDA workflow's nuget download with the
same version. CPU and WebGPU paths stay on 1.27, so the WebGPU-only
SessionReleaseCapturedGraph feature is unaffected.
@qjia7 qjia7 closed this Jun 22, 2026
@qjia7 qjia7 reopened this Jun 22, 2026
Comment thread src/models/onnxruntime_api.h
@qjia7
qjia7 enabled auto-merge (squash) June 26, 2026 02:32
@baijumeswani
baijumeswani disabled auto-merge June 30, 2026 04:08
Comment thread cmake/ortlib.cmake Outdated
Keep the default downloaded ORT at 1.26.0. The ReleaseCapturedGraph
plumbing is already guarded by #if ORT_API_VERSION >= 27, so it
compiles to nothing against 1.26 headers and activates only when
ORT_HOME points at 1.27+. Drops the CUDA 1.26 override in
ortlib.cmake (now redundant with the default), the matching
explanatory comments in win-cuda-x64-build.yml, the per-EP wheel
selection in integration-pytest-step.yml, the cpu requirements.txt
bump to 1.27.0, and the cuda requirements.txt comment block.
@baijumeswani
baijumeswani enabled auto-merge (squash) June 30, 2026 06:13
@kunal-vaishnavi
kunal-vaishnavi disabled auto-merge June 30, 2026 10:58
@kunal-vaishnavi
kunal-vaishnavi merged commit ce9cf9b into main Jun 30, 2026
62 of 65 checks passed
@kunal-vaishnavi
kunal-vaishnavi deleted the per-graph-buffer-manager-webgpu branch June 30, 2026 10:58
tianleiwu pushed a commit that referenced this pull request Jul 11, 2026
Call ReleaseGraph on the ORT session in State::~State() to free
per-graph buffer resources in GPU. This pairs with the per-graph buffer
manager changes in ORT
microsoft/onnxruntime#28260.

Changes:
- Add OrtSession::ReleaseGraph wrapper calling the new C API
- Store graph_capture_session_ pointer in State during graph capture
runs
- Call ReleaseGraph in State destructor with the graph annotation ID

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
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.

4 participants