[CUDA Plugin EP] Expose kernel sync stream for scratch allocation - #29244
Merged
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends the kernel-context C API to expose the framework OrtSyncStream* for the current kernel invocation, and updates the CUDA plugin EP to use that stream for stream-aware scratch allocation bookkeeping so it can safely advertise concurrent Session::Run() when supported by the host runtime.
Changes:
- Adds
OrtApi::KernelContext_GetSyncStream(plus C++ and adapter wrappers) to retrieve the framework stream wrapper associated with a kernel context. - Updates the CUDA plugin kernel adapter to associate scratch/workspace allocations with the framework stream (instead of a null stream tag).
- Re-enables CUDA plugin EP concurrent-run support when the host runtime supports the new API and unified-stream mode is not forced.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
include/onnxruntime/core/session/onnxruntime_c_api.h |
Adds the public C API entry and documentation for KernelContext_GetSyncStream. |
onnxruntime/core/session/ort_apis.h |
Declares the new OrtApis implementation entry point. |
onnxruntime/core/session/onnxruntime_c_api.cc |
Wires the new function pointer into the OrtApi table. |
onnxruntime/core/session/custom_ops.cc |
Implements KernelContext_GetSyncStream by returning the kernel’s framework compute stream wrapper. |
include/onnxruntime/core/session/onnxruntime_cxx_api.h |
Adds Ort::KernelContext::GetSyncStream() declaration. |
include/onnxruntime/core/session/onnxruntime_cxx_inline.h |
Implements the C++ wrapper calling into the C API. |
include/onnxruntime/ep/adapter/op_kernel.h |
Adds version-gated adapter access to GetSyncStream() for plugin kernels. |
onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h |
Tracks/uses the framework stream wrapper for stream-aware scratch allocation tagging. |
onnxruntime/core/providers/cuda/plugin/cuda_ep.cc |
Gates IsConcurrentRunSupported on API availability and unified-stream configuration. |
onnxruntime/test/shared_lib/custom_op_utils.cc |
Extends shared-lib custom-op tests to exercise GetSyncStream(). |
docs/cuda_plugin_ep/cuda_plugin_ep_design.md |
Updates plugin design docs for the new gated stream-tagged scratch capability. |
docs/cuda_plugin_ep/cuda_graph_for_cuda_plugin.md |
Updates CUDA graph docs to reflect stream-tagged scratch allocation and concurrent-run conditions. |
docs/cuda_plugin_ep/arena_allocator_migration_design.md |
Updates allocator migration design docs to reflect stream-tagged scratch allocation and compatibility behavior. |
3 tasks
Base automatically changed from
tlwu/20260623/cuda_plugin_ep_cuda_graph_stream
to
main
June 25, 2026 00:20
Tianlei Wu (tianleiwu)
force-pushed
the
tlwu/concurrent_stream
branch
from
June 25, 2026 00:45
e3d1cc9 to
39c5fe2
Compare
Tianlei Wu (tianleiwu)
requested review from
Copilot and
Dmitri Smirnov (yuslepukhin)
June 25, 2026 20:38
Add ComputeStreamScope RAII guard around CudaKernel::Compute so the thread-local current_cuda_stream/current_framework_stream are scoped to a single Compute invocation (save/restore for nested calls). This prevents kernel constructors that call GetScratchBuffer(..., nullptr) from tagging arena chunks with a stale framework stream pointer whose lifetime ended with a previous Session::Run(). Also clarify the GetScratchBuffer comment that Stream(ctx) likewise records the framework-stream mapping.
Dmitri Smirnov (yuslepukhin)
approved these changes
Jun 26, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR adds a kernel-context C API accessor for the framework
OrtSyncStream*and uses it in the CUDA plugin EP so scratch allocations can be tagged with the actual compute stream selected for the kernel. It is stacked on #29221 and turns the previously documented concurrent multi-stream limitation into a gated capability: older runtimes keep the conservative fallback, while runtimes with the new API can safely advertise concurrent runs when EP-level unified stream mode is not forced.Summary of Changes
Public API and Adapters
include/onnxruntime/core/session/onnxruntime_c_api.hKernelContext_GetSyncStreamto expose the borrowed framework stream for stream-aware allocation and synchronization bookkeeping.onnxruntime/core/session/custom_ops.ccOpKernelContext::GetComputeStream()inside ORT core.onnxruntime/core/session/ort_apis.handonnxruntime/core/session/onnxruntime_c_api.ccinclude/onnxruntime/core/session/onnxruntime_cxx_api.handinclude/onnxruntime/core/session/onnxruntime_cxx_inline.hOrt::KernelContext::GetSyncStream()wrapper.include/onnxruntime/ep/adapter/op_kernel.hCUDA Plugin EP
OrtStreamAdapterstream arguments.KernelContext_GetSyncStreamis available and EP-level unified stream mode is not forced.Tests and Docs
Ort::KernelContext::GetSyncStream().Why a C API is needed
The implementation of
KernelContext_GetSyncStreamis intentionally small, but the API boundary is the important part. ORT core can safely castOrtKernelContext*back toonnxruntime::OpKernelContext*because it owns both the opaque C handle and the private C++ implementation. A plugin kernel should not perform that cast directly: it would make the plugin depend on ORT-core private C++ layout, vtables, and exact build compatibility.The new API keeps that private cast inside ORT core and gives plugin kernels a stable ABI entry point:
This also lets the plugin use runtime version gating. When loaded by an older ORT runtime that does not expose the API, the adapter returns null, scratch allocation uses the conservative fallback, and concurrent runs are not advertised.
Testing
lintrunner -aninja -C build/cu130_plugin/Debug onnxruntime_providers_cuda_pluginninja -C build/cu130_plugin/Debug onnxruntime_shared_lib_testcd build/cu130_plugin/Debug && ./onnxruntime_shared_lib_test --gtest_filter=CApiTest.custom_op_handler --gtest_color=noChecklist