ONNX Runtime CUDA Plugin EP 的 CUDA Graph 支持：捕获/重放架构与实战配置解析

ONNX Runtime CUDA Plugin EP 的 CUDA Graph 支持：捕获/重放架构与实战配置解析

【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime

本文以 ONNX Runtime 仓库中的 CUDA Plugin EP（onnxruntime/providers/cuda/plugin）为对象，深入解析其 CUDA Graph 捕获（capture）与重放（replay）能力的完整设计：包括 OrtEp 图捕获回调接口、warm-up→capture→replay 运行流程、每线程图状态管理、user_compute_stream组合模式、Arena 分配器集成，以及对应的 Python 测试与配置方法。读完本文，你将掌握如何在 CUDA Plugin EP 上开启 CUDA Graph 加速、理解其内部实现原理，并能复现官方测试用例验证效果。

背景：为什么 CUDA Plugin EP 需要自己的 CUDA Graph 支持

CUDA Graph 是 CUDA 运行时提供的一项关键性能优化：它把一段 GPU 操作序列录制（capture）为一张图（cudaGraph_t），实例化（instantiate）为可执行对象（cudaGraphExec_t）后，后续运行只需以极低的 CPU 开销重放（replay）这张图，从而大幅削减每次推理中 kernel launch 的 CPU 开销。

CUDA Plugin EP 是独立于 ONNX Runtime 主发行版发布的共享库（Linux 下为libonnxruntime_providers_cuda.so，Windows 下为onnxruntime_providers_cuda.dll），通过公开的 OrtEp C API 实现 CUDA EP，从而让 CUDA EP 的更新不依赖 ORT 的发布节奏。正因为插件 EP 走的是独立 API 通道，它必须自行实现一套完整的 CUDA Graph 捕获/重放管线，而不能依赖 ORT 核心内置的图逻辑。本文关联的完整设计文档位于 docs/cuda_plugin_ep/cuda_graph_for_cuda_plugin.md。

OrtEp 图捕获回调：插件与 ORT 核心的契约

OrtEp C API（自 ORT v1.26 起）为 EP 提供了四个图捕获回调，配合既有的OnRunStart/OnRunEnd生命周期回调驱动整个捕获工作流：

回调	签名	用途
IsGraphCaptureEnabled	bool(const OrtEp*)	报告图捕获是否启用
IsGraphCaptured	bool(const OrtEp*, int graph_annotation_id)	检查某个 annotation ID 对应的图是否已被捕获
ReplayGraph	OrtStatus*(OrtEp*, int graph_annotation_id)	启动一个先前捕获的图
GetGraphCaptureNodeAssignmentPolicy	OrtGraphCaptureNodeAssignmentPolicy(const OrtEp*)	指定节点分配验证的严格程度

从源码看，这些回调的挂载逻辑位于 cuda_ep.cc 的CudaEp构造函数中，且做了运行时的版本门控：插件可能被加载进比其编译版本更旧的 ORT 运行时，因此IsGraphCaptureEnabled、IsGraphCaptured、ReplayGraph、GetGraphCaptureNodeAssignmentPolicy与GetAvailableResource仅在协商出的运行时版本>= 26（即 ORT 1.26+）时才挂载，旧运行时下这些字段保持nullptr，仅禁用该可选能力而不会让 EP 失效。OnRunStart/OnRunEnd则按enable_cuda_graph配置条件挂载。

四个回调的声明位于 onnxruntime_ep_c_api.h，IExecutionProvider侧的虚函数GetGraphCaptureNodeAssignmentPolicy()位于 execution_provider.h。

架构与运行流程：warm-up → capture → replay

CUDA Graph 捕获要求捕获窗口内的所有 CUDA 操作都在同一流上录制，且捕获期间不允许发生新的内存分配（否则重放无法复现分配）。为此，插件 EP 将一次Session::Run()的生命周期切分为三个阶段：

Session::Run() │ ├─ Run 1..N (warmup): OnRunStart → kernel dispatch → OnRunEnd (increment counter) │ ├─ Run N+1 (capture): OnRunStart → cudaStreamBeginCapture → kernel dispatch │ → OnRunEnd → cudaStreamEndCapture → cudaGraphInstantiate → Replay │ └─ Run N+2+ (replay): IsGraphCaptured() → true → ReplayGraph() → cudaGraphLaunch (OnRunStart/OnRunEnd are NOT called during replay)

三个阶段的核心代码路径在 cuda_ep.cc 中：

• warm-up 阶段：OnRunStartImpl检查该 annotation ID 尚未捕获且已达到最小运行次数。未达到时，OnRunEndImpl调用context.cuda_graph.IncrementRegularRunCount(id)累加普通运行计数。
• 捕获阶段：当运行次数达到min_num_runs_before_cuda_graph_capture（默认 2）时，OnRunStartImpl先执行cudaSetDevice对齐设备、用cudaMemGetInfo记录捕获前 GPU 空闲内存水位，随后CaptureBegin(id)调用cudaStreamBeginCapture(stream_, cudaStreamCaptureModeThreadLocal)（见 cuda_graph_plugin.cc）；OnRunEndImpl则在 kernel 派发完成后调用CaptureEnd(id)，执行cudaStreamEndCapture→cudaGraphInstantiate得到cudaGraphExec_t存入图集合（见 cuda_graph_plugin.cc）。
• 捕获完成后立即重放一次：由于录制到捕获流的 CUDA 工作实际上不会真正在 GPU 上执行，OnRunEndImpl在CaptureEnd之后立即调用Replay(id, sync_stream)把捕获的图真正跑一遍，保证第一次捕获运行的输出正确。
• 重放阶段：后续运行中IsGraphCaptured(id)返回 true，ReplayGraphImpl直接调用CudaGraphManager::Replay(id, /*sync=*/false)执行cudaGraphLaunch，此路径不再经过OnRunStart/OnRunEnd（见 cuda_ep.cc）。

关键设计选择

原设计文档明确了以下几条核心决策，均能在源码中得到印证：

• 每线程独立图状态：每个线程持有自己专属的图cudaStream_t、CudaGraphManager和捕获簿记。CudaSyncStream::InitHandlesWithExternalStream()把线程的图流包装起来，使图捕获与 kernel 派发看到同一流。管理器以 annotation ID 为键存储捕获到的cudaGraphExec_t，因此同一线程可为不同输入形状维护多张图。
• 用户流与图捕获可组合：当同时提供user_compute_stream与图捕获时，每线程上下文直接采用用户自有流作为图流，捕获/重放与调用方驱动的 kernel 运行在同一流上；上下文记录owns_graph_stream = false，永不清理由用户持有的流，流生命周期完全归调用方。
• warm-up 运行（默认 2 次）：让内存分配在捕获开始前稳定下来，避免捕获窗口内出现分配。
• annotation ID 由运行配置驱动：图 annotation ID 从OrtRunOptions的"gpu_graph_id"键解析，-1表示跳过捕获，0为默认值。解析实现见CudaEp::GetGraphAnnotationId()（cuda_ep.cc），它通过GetRunConfigEntry读取并std::stoi转换，解析失败回退到默认值 0。
• 节点分配策略：GetGraphCaptureNodeAssignmentPolicyImpl固定返回OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES（cuda_ep.cc），与非插件版 CUDA EP 行为一致，允许形状推断类节点落在 CPU 上。
• 线程安全性：图状态与图流全部按线程存储；ORT 核心的DeviceStreamCollection缓存同样是线程亲和（thread-affine）的，图启用的运行可以在同一线程上回收流包装器而不泄露给其他线程。
• 运行结束同步：OnRunEndImpl遵循sync_stream标志，但对已重放的图不做二次同步（避免双同步），保持常规 EP 完成契约。

新增组件源码解析

设计文档提出并落地了四个核心组件，均位于 onnxruntime/core/providers/cuda/plugin/ 目录：

CudaGraphSet

CudaGraphSet（cuda_graph_plugin.h）是cudaGraphExec_t的哈希表存储，以 annotation ID 为键，拥有图执行资源：析构时Clear()逐个cudaGraphExecDestroy。Put()对重复 ID 抛异常提示换用不同 annotation ID，Get()对缺失 ID 抛异常——这些约束保证一张图在一个 ID 下只捕获一次。

CudaGraphManager

CudaGraphManager（cuda_graph_plugin.h）编排捕获生命周期，对外暴露CaptureBegin()/CaptureEnd()/Replay()/Reset()，并提供 warm-up 跟踪：IncrementRegularRunCount(id)按 ID 累加普通运行次数，IsGraphCaptureAllowed(id, min_runs)判断是否达到最小 warm-up 次数。CaptureBegin先同步流再以cudaStreamCaptureModeThreadLocal模式开始捕获，Replay对不存在的图返回ORT_INVALID_ARGUMENT、对启动/同步失败返回ORT_EP_FAIL并携带 CUDA 错误名与错误串（见 cuda_graph_plugin.cc）。

CudaEp::PerThreadContext

PerThreadContext（cuda_ep.cc）是每线程图流、CudaGraphManager与捕获前空闲内存水位（pre_capture_free_mem）的所有者。它的关键行为：

• 由thread_local缓存（PerThreadContextCache()）持有强引用，线程退出时自动析构并释放流与捕获的图执行对象；
• CudaEp只通过per_thread_context_caches_持有各线程缓存映射的弱引用，用于 EP 析构时清除自己的条目，且创建新上下文时会顺带清理已失效的缓存映射（见 cuda_ep.cc）；
• 析构顺序为先cuda_graph.Reset()（销毁图执行对象）再销毁流，避免在流销毁后重放。

CudaSyncStream::InitHandlesWithExternalStream

InitHandlesWithExternalStream()（cuda_stream_plugin.cc）包装外部（非自有）cudaStream_t用于注册与生命周期跟踪，设置owns_stream_ = false。当包装器不拥有库句柄时，迁移的 kernel 在派发时通过线程局部回退句柄（thread-local fallback handles）绑定 cuBLAS/cuDNN/cuBLASLt。

配置选项与使用示例

原文档给出的配置项如下表，均可通过 session 配置或 provider 选项设置：

Option Key	类型	默认值	说明
ep.cuda.enable_cuda_graph	bool	false	启用 CUDA graph 捕获/重放
ep.cuda.min_num_runs_before_cuda_graph_capture	int	2	捕获前的 warm-up 运行次数

同时兼容历史扁平别名enable_cuda_graph。参数解析实现位于 cuda_ep_factory.cc：read_session_config_bool优先读取 EP 作用域键（ep.cuda.*），再回退到历史扁平 session 配置名；min_num_runs_before_cuda_graph_capture走read_session_config_non_negative_int读取。CudaEp::Config中对应字段的默认值定义见 cuda_ep.h（enable_cuda_graph = false，min_num_runs_before_cuda_graph_capture = 2）。

Python 开启方式

以 Python API 为例，在创建InferenceSession时通过 provider 选项与 session 配置开启：

import onnxruntime as ort sess = ort.InferenceSession( "model.onnx", providers=[ ( "CUDAExecutionProvider", { "device_id": "0", "ep.cuda.enable_cuda_graph": "1", "ep.cuda.min_num_runs_before_cuda_graph_capture": "2", }, ), "CPUExecutionProvider", ], )

注意：CUDA Plugin EP 的构建与注册方式（含捆绑 wheel 与独立插件包的两种注册路径）详见 QUICK_START.md。默认构建使用传统源码内 CUDA EP（onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=OFF），只有传入--cmake_extra_defines "onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON"才会产出插件形态的onnxruntime_providers_cuda库。

运行配置驱动的多图管理：gpu_graph_id

除静态配置外，每次Run()还可通过OrtRunOptions的"gpu_graph_id"键选择要使用/捕获哪张图。这在处理变长输入形状（如不同 batch size、不同序列长度）时非常实用：为每个形状各自捕获一张图，运行时按形状挑选对应 ID 重放。Python 中通过RunOptions.add_run_config_entry("gpu_graph_id", "...")设置：

import onnxruntime as onnxrt ro1 = onnxrt.RunOptions() ro1.add_run_config_entry("gpu_graph_id", "1") session.run_with_iobinding(io1, ro1) # 形状 1 使用/捕获 graph 1 ro2 = onnxrt.RunOptions() ro2.add_run_config_entry("gpu_graph_id", "2") session.run_with_iobinding(io2, ro2) # 形状 2 使用/捕获 graph 2

官方测试test_cuda_graph_annotation_id（test_cuda_plugin_ep.py）正是用动态维度的 MatMul 模型，分别以gpu_graph_id=1（形状[2,3]@[3,4]）与gpu_graph_id=2（形状[4,5]@[5,6]）捕获两张图，随后用更新后的同形状数据分别重放并校验结果。

User Compute Stream + CUDA Graph 组合模式

调用方可以同时提供自有 CUDA 流（user_compute_streamprovider 选项，以指针形式传入）并启用图捕获。这一组合此前会被ORT_INVALID_ARGUMENT直接拒绝，现在已支持，且行为与捆绑（非插件）CUDA EP 对齐。关键源码路径：

• CudaEpFactory::CreateEpImpl不再拒绝该组合（见 cuda_ep_factory.cc），但设置user_compute_stream仍会强制启用 unified-stream 模式（use_ep_level_unified_stream = true），与捆绑 EP 一致。
• CreateSyncStreamForDeviceImpl的第一分支通过InitHandlesWithUserStream()包装用户流并挂载完整的 cuBLAS/cuDNN/cuBLASLt 句柄（cuda_ep.cc）。
• GetPerThreadContext()在has_user_compute_stream && enable_cuda_graph时，用use_external_stream标志构建以用户流为external_graph_stream的PerThreadContext（cuda_ep.cc），捕获与重放因此与 kernel 派发在同一流上。
• 上下文记录owns_graph_stream = false：析构时销毁已捕获的图执行对象，但绝不调用cudaStreamDestroy销毁用户流，流生命周期始终归调用方（cuda_ep.cc）。

值得注意的细节是：判断是否采用外部流用的是调用者意图标志（has_user_compute_stream）而非句柄是否非空——用户可能显式选择 CUDA 默认流（cudaStream_t(0)，即nullptr），这仍然是外部、用户所有的流，上下文同样不得销毁它。

由于用户只提供一个流，此模式天然是单流的；每线程图隔离依然生效——若同一 session 被多线程驱动，每个线程必须在它自己提供的流上驱动自己的捕获。

一个常见疑问：user_compute_stream 只对图启用生效吗？

阅读GetPerThreadContext()时容易产生疑问：为什么use_external_stream被has_user_compute_stream && enable_cuda_graph门控？这是否意味着用户计算流仅在图启用的运行中生效？答案是否定的：

• 用户计算流在图与非图运行中都被 kernel 使用。这发生在CreateSyncStreamForDeviceImpl的第一分支——它独立于enable_cuda_graph通过InitHandlesWithUserStream()包装config_.user_compute_stream（cuda_ep.cc）。
• enable_cuda_graph条件只影响PerThreadContext的图流。而PerThreadContext是纯图捕获对象：GetPerThreadContext()只从图路径触达（CreateSyncStreamForDeviceImpl的enable_cuda_graph分支、OnRunStart/OnRunEnd、IsGraphCaptured、ReplayGraph）。图未启用时根本不会构造PerThreadContext，其流所有权标志自然无关紧要。
• 因此该标志回答的是一个更窄的问题——"每线程的捕获/重放图流是否应采用（且不销毁）用户的流？"——它只在真正捕获图时有意义。

Arena 分配器集成：让捕获窗口内零分配

CUDA Graph 捕获的硬性要求是所有内存分配都发生在 warm-up 阶段，而不是捕获阶段。插件 Arena 分配器已落地并与图捕获路径集成。源码侧的具体机制：

• 默认 CUDA 设备分配来自插件托管的 Arena（CudaArenaAllocator，见 cuda_arena.cc）。warm-up 运行期间 Arena 增长到容纳所需全部 chunk；捕获与重放期间复用这些 chunk，不再触发cudaMalloc。
• kernel scratch/workspace 分配走 EP Arena：CudaKernel::GetScratchBuffer现在通过Info().GetAllocator()（即 EP arena）分配，而非每次调用原始cudaMallocAsync/cudaMalloc（见 cuda_kernel_adapter.h）。warm-up 后 Arena 达到稳态工作集，捕获运行的所有 scratch 请求都命中已预留 chunk，捕获窗口内设备空闲内存足迹保持稳定——这正是cudaMemGetInfo分配检测器对使用 scratch buffer 的图能够通过的原因，也与捆绑 CUDA EP 的行为一致（后者同样从Info().GetAllocator()获取 scratch）。
• 双流区分：GetScratchBuffer用KernelContext_GetSyncStream暴露的框架OrtSyncStream*对流打标（stream-tag），供 Arena 做跨流复用簿记；而 CUDA launch 与库调用仍使用KernelContext_GetGPUComputeStream返回的原始cudaStream_t。
• CUDA mempool 模式：配置arena.use_cuda_mempool=1时，CUDA 设备分配改由CudaMempoolOrtAllocator提供（见 cuda_mempool_allocator_plugin.h），其内部包装cudaMallocFromPoolAsync/cudaFreeAsync。这些异步分配/释放操作自 CUDA 11.4+ 起对 CUDA Graph 安全，会成为被捕获图拓扑的一部分。
• Pinned 分配同样由 Arena 支撑，但不参与流感知。
• 图流贯穿分配路径：PerThreadContext创建的图流经InitHandlesWithExternalStream()流入流感知 Arena 分配，保证 warm-up、捕获、重放使用同一cudaStream_t。
• 每次运行结束重置 chunk 关联：CudaSyncStream::OnSessionRunEndImpl()通过factory_.ResetDeviceArenaChunksUsingStream()在每次运行（含图启用运行）结束时重置 arena chunk 与流的关联，且执行于流包装器被回收进当前线程池桶之前。

分配发生在捕获期间的检测器是 Arena 集成后仍保留的最后一道防线：

• OnRunStartImpl在CaptureBegin之前通过cudaMemGetInfo把 GPU 空闲内存记录到每线程上下文（pre_capture_free_mem）；
• OnRunEndImpl在CaptureEnd后比较捕获后的空闲内存，若下降则打印 WARNING，建议用户增大min_num_runs_before_cuda_graph_capture让分配稳定下来（cuda_ep.cc）。

之所以保留此检查，是因为自定义 arena 选项、warm-up 不足或回归仍可能暴露捕获期间的分配问题。分配器集成细节可进一步参考 arena_allocator_migration_design.md。

并发运行支持

当宿主 ORT 运行时暴露KernelContext_GetSyncStream且 session 未被强制为 EP 级 unified-stream 模式时，CUDA Plugin EP 会宣告支持并发Session::Run()。判定实现在IsConcurrentRunSupportedImpl（cuda_ep.cc）：要求!config_.use_ep_level_unified_stream且协商的 ORT API 版本>= 28（KernelContext_GetSyncStream的最低版本）。具体机制：

• 图簿记（图流、图管理器、warm-up 次数、内存水位）仍按线程存放于PerThreadContext，避免跨线程共享捕获的图执行对象；
• 插件 kernel 通过KernelContext_GetSyncStream获取框架OrtSyncStream*，仅用于 scratch/workspace 分配的簿记；CUDA 工作仍通过KernelContext_GetGPUComputeStream的原始cudaStream_t启动；
• 带流标记（stream-tag）的 scratch chunk 让共享 Arena 对不同 CUDA 流上重叠运行的运行应用常规跨流复用规则；
• 当协商的 ORT API 版本不含KernelContext_GetSyncStream时，CudaKernel::GetScratchBuffer回退到空流标记，且IsConcurrentRunSupportedImpl返回 false。

验证：官方测试与复现步骤

Python 集成测试

聚焦验证脚本为 test_cuda_plugin_ep.py，其 CUDA graph 测试组覆盖：

测试	验证内容
test_cuda_graph_capture_and_replay	默认 arena 下的 warm-up + capture + replay 全流程（L2622）
test_cuda_graph_replay_with_updated_input	捕获后原地（in-place）更新输入再重放，结果正确（L2666）
test_cuda_graph_with_mempool	arena.use_cuda_mempool=1下的图捕获（L2704）
test_cuda_graph_annotation_id	通过gpu_graph_id运行配置管理多张图（不同形状）（L2760）
test_cuda_graph_second_device	非默认插件设备上的捕获/重放（多 GPU）（L2838）
test_cuda_graph_add_model	Arena 支撑的 Add 模型图捕获（L2891）

复现步骤（完整说明见 QUICK_START.md 的 Build Instructions 与 Running Tests 章节）：

# 1. 构建并部署插件（需传入 onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON） # 2. 指向插件库（Linux 示例） export ORT_CUDA_PLUGIN_PATH=/path/to/build/Release/libonnxruntime_providers_cuda.so # 3. 运行测试 cd onnxruntime/test/python/transformers export ORT_TEST_CUDA_PLUGIN_EP=1 python test_cuda_plugin_ep.py

以test_cuda_graph_capture_and_replay为例，其核心断言逻辑为：创建 MatMul 模型 → 用 provider 选项开启 CUDA graph 的 session → 第一次run_with_iobinding自动完成 warm-up + capture + 首次 replay 并校验A @ B结果 → 第二次运行走重放快路径再次校验。而test_cuda_graph_replay_with_updated_input则在OrtValue.update_inplace()原地写入新数据（同形状、不同数值）后重放，验证图捕获的是内存地址而非具体数值，这是 CUDA Graph 加速在真实推理场景可用的前提。

C++ 单元测试

onnxruntime/test/providers/cuda/plugin/cuda_plugin_user_stream_graph_test.cc 是覆盖user_compute_stream与enable_cuda_graph组合的 C++ 测试（以ORT_UNIT_TEST_HAS_CUDA_PLUGIN_EP门控），重点回归验证：

• 同时设置两个选项时 session 创建成功（对应被移除的工厂层校验）；
• 在用户流上捕获 + 重放结果正确；
• 在用户流上原地更新输入后重放结果正确。

针对最小支持 ORT 版本的验证

插件可能被加载进旧版 ORT 运行时（最小版本声明于 plugin-ep-cuda/MIN_ONNXRUNTIME_VERSION，目前为 1.24.4）。可用如下方式验证新构建的插件在最小运行时上的图能力：

pip install "onnxruntime==$(cat plugin-ep-cuda/MIN_ONNXRUNTIME_VERSION)" --force-reinstall export ORT_CUDA_PLUGIN_PATH=/path/to/build/Release/libonnxruntime_providers_cuda.so cd onnxruntime/test/python/transformers python test_cuda_plugin_ep.py

未来工作

设计文档明确指出一项后续计划：Profiling 集成。当前 CUDA Graph 重放绕过了 CUDA Plugin EP 的 profiler 路径，因为插件 EP 尚未实现OrtEp::CreateProfiler（源码中 profiler 相关回调以ENABLE_CUDA_PROFILING编译开关控制，见 cuda_profiler_plugin.h）。将图重放接入该路径，使得 CUPTI 等性能分析工具能够观测重放运行，是后续版本的工作方向。

小结

CUDA Plugin EP 的图捕获/重放能力是一套自包含、面向 OrtEp API 的完整实现：以四个图捕获回调为契约、以每线程PerThreadContext+CudaGraphManager为状态核心、以 Arena 分配器集成为正确性前提，并支持user_compute_stream组合与gpu_graph_id多图管理。开启方式只需两个配置项（ep.cuda.enable_cuda_graph与ep.cuda.min_num_runs_before_cuda_graph_capture），官方 Python 测试提供了从 warm-up 到多形状重放的完整验证矩阵，可作为上手与回归测试的起点。

【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime