[Plugin EP] Port graph capture/replay APIs

include/onnxruntime/core/framework/execution_provider.h

@@ -292,6 +292,15 @@ class IExecutionProvider {
     return Status::OK();
   }
 
+  /**
+     Get the node assignment validation policy for graph capture.
+     When graph capture is enabled, ORT validates that nodes are assigned to EPs
+     in a way compatible with graph capture. This tells ORT which policy to apply.
+   */
+  virtual OrtGraphCaptureNodeAssignmentPolicy GetGraphCaptureNodeAssignmentPolicy() const {
+    return OrtGraphCaptureNodeAssignmentPolicy_ALL_NODES_ON_EP;
+  }
+
   /**
      Called when session creation is complete
      This provides an opportunity for execution providers to optionally synchronize and

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -1389,6 +1389,10 @@ struct Env : detail::Base<OrtEnv> {
                      const std::vector<Value>& dst_tensors,
                      OrtSyncStream* stream) const;  ///< Wraps OrtApi::CopyTensors
 
+  /// Wraps OrtApi::CopyTensors
+  /// Copies only one src tensor to another dst tensor.
+  Status CopyTensor(const OrtValue* src_tensor, OrtValue* dst_tensor, OrtSyncStream* stream) const;
+
   /// \brief Wraps OrtApi::SetPerSessionThreadPoolCallbacks
   /// Stores work callbacks on the Env for per-session thread pools.
   /// Only affects sessions created after this call. Does not affect global thread pools.

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -1055,6 +1055,11 @@ inline Status Env::CopyTensors(const std::vector<Value>& src_tensors,
   return Status(status);
 }
 
+inline Status Env::CopyTensor(const OrtValue* src_tensor, OrtValue* dst_tensor, OrtSyncStream* stream) const {
+  OrtStatus* status = GetApi().CopyTensors(p_, &src_tensor, &dst_tensor, stream, 1);
+  return Status(status);
+}
+
 inline UnownedAllocator Env::CreateSharedAllocator(const OrtEpDevice* ep_device, OrtDeviceMemoryType mem_type,
                                                    OrtAllocatorType allocator_type,
                                                    const OrtKeyValuePairs* allocator_options) {

include/onnxruntime/core/session/onnxruntime_ep_c_api.h

@@ -2027,6 +2027,23 @@ typedef enum OrtEpDataLayout {
   OrtEpDataLayout_Default = OrtEpDataLayout_NCHW,
 } OrtEpDataLayout;
 
+/**
+ * \brief Node assignment policies for graph capture validation.
+ *
+ * When graph capture is enabled, ORT validates that nodes are assigned to EPs in a way that is
+ * compatible with graph capture. An EP can specify which validation policy ORT should apply.
+ *
+ * \since Version 1.26.
+ */
+typedef enum OrtGraphCaptureNodeAssignmentPolicy {
+  /** All nodes in the main graph must be assigned to this EP. No CPU fallback is allowed. */
+  OrtGraphCaptureNodeAssignmentPolicy_ALL_NODES_ON_EP = 0,
+
+  /** Compute nodes must be on this EP. CPU nodes are allowed for shape computation as long as
+   *  no memory copy nodes exist. */
+  OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES = 1,
+} OrtGraphCaptureNodeAssignmentPolicy;
+
 /**
  * \brief The OrtEp struct provides functions to implement for an execution provider.
  * \since Version 1.22.
@@ -2346,6 +2363,101 @@ struct OrtEp {
    */
   ORT_API2_STATUS(CreateProfiler, _In_ OrtEp* this_ptr,
                   _Outptr_result_maybenull_ OrtEpProfilerImpl** profiler);
+
+  /** \brief Indicate whether the graph capturing mode (e.g., CUDA graph) is enabled for the provider.
+   *
+   * Graph capture allows an EP to record a sequence of device (e.g., GPU) operations during an initial run and replay
+   * them on subsequent runs, bypassing per-kernel CPU launch overhead.
+   *
+   * Applications enable graph capture via EP-specific provider options (e.g., `enable_cuda_graph=1`
+   * for the CUDA EP). An EP should return true from this function if it has been configured to enable
+   * graph capture/replay.
+   *
+   * **ORT graph capture/replay summary:**
+   * During OrtSession initialization, ORT calls OrtEp::IsGraphCaptureEnabled() on each EP in the order specified during
+   * provider registration with the session. If an EP returns true, ORT validates that the graph is suitable for
+   * graph capture, and if so, caches the EP for graph capture during the next run. The graph validation ensures
+   * that there are no control flow nodes and that node-to-EP assignments are compatible with the policy specified
+   * by the EP via OrtEp::GetGraphCaptureNodeAssignmentPolicy().
+   * Note that an OrtSession only supports graph capture for one EP (i.e., the first EP to claim support).
+   *
+   * During the first call to OrtApi::Run() for the OrtSession, ORT performs multiple internal runs of the model
+   * until the EP indicates that the graph has been captured by returning `true` from `OrtEp::IsGraphCaptured()`.
+   * If the EP is unable to capture the graph within 8 runs, the call to OrtApi::Run() returns an error OrtStatus.
+   * Each internal run invokes `OrtEp::OnRunStart()`, normal execution, and `OrtEp::OnRunEnd()`. EPs should use
+   * these run callbacks to track the number of necessary warm-up runs and begin/end graph capture when ready.
+   *
+   * After successful graph capture, subsequent calls to OrtApi::Run() skip normal execution and ORT instead calls
+   * `OrtEp::ReplayGraph()` directly.
+   *
+   * Applications can capture and replay multiple graphs (e.g., one per distinct input shape) by setting the
+   * `"gpu_graph_id"` run config entry via `OrtApi::AddRunConfigEntry()` to different integer values. ORT passes
+   * the value as the `graph_annotation_id` parameter to `OrtEp::IsGraphCaptured()` and `OrtEp::ReplayGraph()`.
+   *
+   * \param[in] this_ptr The OrtEp instance.
+   * \return true if graph capture mode is enabled, false otherwise.
+   *
+   * \note Implementation of this function is optional. If set to NULL, ORT assumes graph capture is not enabled.
+   * \note If this function returns true, `OrtEp::IsGraphCaptured` and `OrtEp::ReplayGraph` must also be implemented.
+   *       If either is NULL, ORT will log a warning and ignore this EP for graph capture.
+   *
+   * \since Version 1.26.
+   */
+  ORT_API_T(bool, IsGraphCaptureEnabled, _In_ const OrtEp* this_ptr);
+
+  /** \brief Indicate whether a graph has been captured and instantiated.
+   *
+   * ORT calls this before each `Session::Run()`. If true, ORT calls `ReplayGraph()` instead of
+   * normal execution. After a run where this returns false, ORT automatically retries until it
+   * returns true (handling warm-up runs transparently).
+   *
+   * \param[in] this_ptr The OrtEp instance.
+   * \param[in] graph_annotation_id Identifies which captured graph to query.
+   *            Applications can set this value via `OrtApi::AddRunConfigEntry()` with the key `"gpu_graph_id"`.
+   *            The default value is 0 when the run config entry is not set.
+   *            Setting different IDs allows the EP to capture and manage multiple graphs (e.g., one per
+   *            distinct input shape). A value of -1 means graph capture/replay should be skipped for this run.
+   * \return true if the graph has been captured, false otherwise.
+   *
+   * \note This function must be implemented if `OrtEp::IsGraphCaptureEnabled` is implemented and may return true.
+   *
+   * \since Version 1.26.
+   */
+  ORT_API_T(bool, IsGraphCaptured, _In_ const OrtEp* this_ptr, _In_ int graph_annotation_id);
+
+  /** \brief Run the instantiated (captured) graph.
+   *
+   * Called by ORT instead of normal execution when `IsGraphCaptured()` returns true.
+   *
+   * \param[in] this_ptr The OrtEp instance.
+   * \param[in] graph_annotation_id Identifies which captured graph to replay.
+   *            Applications can set this value via `OrtApi::AddRunConfigEntry()` with the key `"gpu_graph_id"`.
+   *            The default value is 0 when the run config entry is not set.
+   *            A value of -1 means graph replay should be skipped for this run.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \note This function must be implemented if `OrtEp::IsGraphCaptureEnabled` is implemented and may return true.
+   *
+   * \since Version 1.26.
+   */
+  ORT_API2_STATUS(ReplayGraph, _In_ OrtEp* this_ptr, _In_ int graph_annotation_id);
+
+  /** \brief Get the node assignment validation policy for graph capture.
+   *
+   * When graph capture is enabled, ORT validates that nodes are assigned to EPs in a way that is
+   * compatible with graph capture. This function tells ORT which validation policy to apply.
+   *
+   * \param[in] this_ptr The OrtEp instance.
+   * \return The node assignment policy for graph capture.
+   *
+   * \note Implementation of this function is optional. If set to NULL, ORT uses
+   *       OrtGraphCaptureNodeAssignmentPolicy_ALL_NODES_ON_EP (strictest validation).
+   *
+   * \since Version 1.26.
+   */
+  ORT_API_T(OrtGraphCaptureNodeAssignmentPolicy, GetGraphCaptureNodeAssignmentPolicy,
+            _In_ const OrtEp* this_ptr);
 };
 
 /** \brief The function signature that ORT will call to create OrtEpFactory instances.

onnxruntime/core/providers/cuda/cuda_execution_provider.h

@@ -124,6 +124,9 @@ class CUDAExecutionProvider : public IExecutionProvider {
   bool IsGraphCaptureEnabled() const override;
   bool IsGraphCaptured(CudaGraphAnnotation_t graph_annotation_id) const override;
   Status ReplayGraph(CudaGraphAnnotation_t graph_annotation_id) override;
+  OrtGraphCaptureNodeAssignmentPolicy GetGraphCaptureNodeAssignmentPolicy() const override {
+    return OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES;
+  }
   void RegisterStreamHandlers(IStreamCommandHandleRegistry& stream_handle_registry, AllocatorMap& allocators) const override;
   OrtDevice GetOrtDeviceByMemType(OrtMemType mem_type) const override;
   std::vector<AllocatorPtr> CreatePreferredAllocators() override;

onnxruntime/core/providers/cuda/plugin/cuda_ep.cc

@@ -58,6 +58,12 @@
   Compile = nullptr;
   ReleaseNodeComputeInfos = nullptr;
 
+  // Graph capture/replay
+  IsGraphCaptureEnabled = IsGraphCaptureEnabledImpl;
+  IsGraphCaptured = IsGraphCapturedImpl;
+  ReplayGraph = ReplayGraphImpl;
+  GetGraphCaptureNodeAssignmentPolicy = GetGraphCaptureNodeAssignmentPolicyImpl;
+
   const OrtApi& ort_api = factory_.GetOrtApi();
   Ort::Status log_status(ort_api.Logger_LogMessage(&logger_, ORT_LOGGING_LEVEL_INFO,
                                                    "CUDA Plugin EP created",
@@ -304,5 +310,30 @@
   EXCEPTION_TO_STATUS_END
 }
 
+bool ORT_API_CALL CudaEp::IsGraphCaptureEnabledImpl(const OrtEp* /*this_ptr*/) noexcept {
+  // TODO: forward to EpImpl()->IsGraphCaptureEnabled()
+  return false;
+}
+
+/*static*/
+bool ORT_API_CALL CudaEp::IsGraphCapturedImpl(const OrtEp* /*this_ptr*/, int /*graph_annotation_id*/) noexcept {
+  // TODO: forward to EpImpl()->IsGraphCaptured(graph_annotation_id)
+  return false;
+}
+
+/*static*/
+OrtStatus* ORT_API_CALL CudaEp::ReplayGraphImpl(OrtEp* /*this_ptr*/, int /*graph_annotation_id*/) noexcept {
+  // TODO: forward to EpImpl()->ReplayGraph(graph_annotation_id)
+  return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED,
+                                    "Graph capture replay is not yet supported in the CUDA plugin EP.");
+}
+
+/*static*/
+OrtGraphCaptureNodeAssignmentPolicy ORT_API_CALL CudaEp::GetGraphCaptureNodeAssignmentPolicyImpl(
+    const OrtEp* /*this_ptr*/) noexcept {
+  // TODO: forward to EpImpl()->GetGraphCaptureNodeAssignmentPolicy()
+  return OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES;
+}
+
 }  // namespace cuda_plugin
 }  // namespace onnxruntime

onnxruntime/core/providers/cuda/plugin/cuda_ep.h

@@ -61,6 +61,17 @@ class CudaEp : public onnxruntime::ep::adapter::Ep {
 
   static OrtStatus* ORT_API_CALL SyncImpl(OrtEp* this_ptr) noexcept;
 
+  static bool ORT_API_CALL IsGraphCaptureEnabledImpl(const OrtEp* this_ptr) noexcept;
+
+  static bool ORT_API_CALL IsGraphCapturedImpl(const OrtEp* this_ptr,
+                                               int graph_annotation_id) noexcept;
+
+  static OrtStatus* ORT_API_CALL ReplayGraphImpl(OrtEp* this_ptr,
+                                                 int graph_annotation_id) noexcept;
+
+  static OrtGraphCaptureNodeAssignmentPolicy ORT_API_CALL GetGraphCaptureNodeAssignmentPolicyImpl(
+      const OrtEp* this_ptr) noexcept;
+
   CudaEpFactory& factory_;
   std::string name_;
   Config config_;

onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h

@@ -356,6 +356,11 @@ namespace Dml
             return m_impl->ReplayGraph(graph_annotation_id);
         }
 
+        OrtGraphCaptureNodeAssignmentPolicy GetGraphCaptureNodeAssignmentPolicy() const override
+        {
+            return OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES;
+        }
+
     private:
         ComPtr<ExecutionProviderImpl> m_impl;
     };

onnxruntime/core/providers/js/js_execution_provider.h

@@ -72,6 +72,9 @@ class JsExecutionProvider : public IExecutionProvider {
   bool IsGraphCaptureEnabled() const override;
   bool IsGraphCaptured(int graph_annotation_id) const override;
   Status ReplayGraph(int graph_annotation_id) override;
+  OrtGraphCaptureNodeAssignmentPolicy GetGraphCaptureNodeAssignmentPolicy() const override {
+    return OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES;
+  }
 
  private:
   bool IsGraphCaptureAllowed() const;

onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc

@@ -1280,7 +1280,9 @@ void NvExecutionProvider::HandleCudaGraphStart(cudaStream_t stream, bool require
 }
 
 bool NvExecutionProvider::IsGraphCaptureEnabled() const {
-  return cuda_graph_enable_;
+  // Return false so that ORT's framework does not cache this EP for ORT-managed graph capture/replay.
+  // NvTensorRTRTX manages CUDA graph capture/replay internally.
+  return false;
 }
 
 bool NvExecutionProvider::IsGraphCaptured(int graph_annotation_id) const {

onnxruntime/core/providers/webgpu/ep/ep.cc

@@ -40,6 +40,10 @@ Ep::Ep(std::unique_ptr<IExecutionProvider> impl, Factory& factory, const OrtLogg
   CreateSyncStreamForDevice = nullptr;          // Not stream aware
   GetCompiledModelCompatibilityInfo = nullptr;  // Not a compiled EP
   IsConcurrentRunSupported = IsConcurrentRunSupportedImpl;
+  IsGraphCaptureEnabled = IsGraphCaptureEnabledImpl;
+  IsGraphCaptured = IsGraphCapturedImpl;
+  ReplayGraph = ReplayGraphImpl;
+  GetGraphCaptureNodeAssignmentPolicy = GetGraphCaptureNodeAssignmentPolicyImpl;
 }
 
 // OrtEp interface implementations
@@ -253,6 +257,34 @@ OrtStatus* ORT_API_CALL Ep::IsConcurrentRunSupportedImpl(_In_ OrtEp* /*this_ptr*
   return nullptr;
 }
 
+bool ORT_API_CALL Ep::IsGraphCaptureEnabledImpl(_In_ const OrtEp* this_ptr) noexcept {
+  auto* ep = static_cast<const Ep*>(this_ptr);
+  return ep->EpImpl()->IsGraphCaptureEnabled();
+}
+
+bool ORT_API_CALL Ep::IsGraphCapturedImpl(_In_ const OrtEp* this_ptr, _In_ int graph_annotation_id) noexcept {
+  auto* ep = static_cast<const Ep*>(this_ptr);
+  return ep->EpImpl()->IsGraphCaptured(graph_annotation_id);
+}
+
+OrtStatus* ORT_API_CALL Ep::ReplayGraphImpl(_In_ OrtEp* this_ptr, _In_ int graph_annotation_id) noexcept {
+  EXCEPTION_TO_RETURNED_STATUS_BEGIN
+  auto* ep = static_cast<Ep*>(this_ptr);
+  auto status = ep->EpImpl()->ReplayGraph(graph_annotation_id);
+  if (!status.IsOK()) {
+    return Api().ort.CreateStatus(static_cast<OrtErrorCode>(status.Code()),
+                                  status.ErrorMessage().c_str());
+  }
+  return nullptr;
+  EXCEPTION_TO_RETURNED_STATUS_END
+}
+
+OrtGraphCaptureNodeAssignmentPolicy ORT_API_CALL Ep::GetGraphCaptureNodeAssignmentPolicyImpl(
+    _In_ const OrtEp* this_ptr) noexcept {
+  auto* ep = static_cast<const Ep*>(this_ptr);
+  return ep->EpImpl()->GetGraphCaptureNodeAssignmentPolicy();
+}
+
 OrtStatus* ORT_API_CALL Ep::CreateAllocatorImpl(_In_ OrtEp* this_ptr,
                                                 _In_ const OrtMemoryInfo* memory_info,
                                                 _Outptr_result_maybenull_ OrtAllocator** allocator) noexcept {

onnxruntime/core/providers/webgpu/ep/ep.h

@@ -67,6 +67,17 @@ class Ep : public onnxruntime::ep::adapter::Ep {
   static OrtStatus* ORT_API_CALL IsConcurrentRunSupportedImpl(_In_ OrtEp* this_ptr,
                                                               _Out_ bool* is_concurrent_run_supported) noexcept;
 
+  static bool ORT_API_CALL IsGraphCaptureEnabledImpl(_In_ const OrtEp* this_ptr) noexcept;
+
+  static bool ORT_API_CALL IsGraphCapturedImpl(_In_ const OrtEp* this_ptr,
+                                               _In_ int graph_annotation_id) noexcept;
+
+  static OrtStatus* ORT_API_CALL ReplayGraphImpl(_In_ OrtEp* this_ptr,
+                                                 _In_ int graph_annotation_id) noexcept;
+
+  static OrtGraphCaptureNodeAssignmentPolicy ORT_API_CALL GetGraphCaptureNodeAssignmentPolicyImpl(
+      _In_ const OrtEp* this_ptr) noexcept;
+
   Factory& factory_;
   const OrtLogger& logger_;
   Config config_{};

onnxruntime/core/providers/webgpu/webgpu_execution_provider.h

@@ -97,6 +97,9 @@ class WebGpuExecutionProvider : public IExecutionProvider {
   bool IsGraphCaptureEnabled() const override;
   bool IsGraphCaptured(int graph_annotation_id) const override;
   Status ReplayGraph(int graph_annotation_id) override;
+  OrtGraphCaptureNodeAssignmentPolicy GetGraphCaptureNodeAssignmentPolicy() const override {
+    return OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES;
+  }
   webgpu::BufferManager& BufferManager() const;
   AllocatorPtr PrepackAllocator() const { return prepack_allocator_; }
   std::span<const std::string> GetForceCpuNodeNames() const { return force_cpu_node_names_; }