Per-graph buffer manager for WebGPU multi-graph capture

Replace the single shared graph_buffer_mgr_ with per-annotation-ID buffer
managers so that multiple generators with different prompts each get isolated
buffer caches. This prevents cross-generator buffer corruption when graph
capture is enabled.

Changes:
- Add per_graph_buffer_mgrs_ map keyed by graph annotation ID
- Route BufferManager() to the correct per-graph instance during capture/replay
- Change GpuBufferAllocator from const BufferManager& to pointer with
  SetBufferManager() for dynamic routing
- Add ReleaseGraph API through the full stack (EP base class, C API,
  InferenceSession, plugin EP, EP adapter) so callers can free captured
  graph resources when a generator is destroyed
- WebGPU EP ReleaseGraph cleans up captured commands, buffer manager,
  and tracking state for the given annotation ID

Author: qjia7

include/onnxruntime/core/framework/execution_provider.h

@@ -292,6 +292,14 @@ class IExecutionProvider {
     return Status::OK();
   }
 
+  /**
+     Release a previously captured graph and its associated resources.
+     Called when the caller no longer needs the captured graph for the given annotation ID.
+   */
+  virtual common::Status ReleaseGraph(int /*graph_annotation_id*/) {
+    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

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -7435,6 +7435,21 @@ struct OrtApi {
   ORT_API2_STATUS(SetPerSessionThreadPoolCallbacks, _Inout_ OrtEnv* env,
                   _In_ const OrtThreadPoolCallbacksConfig* config);
 
+  /** \brief Release a previously captured graph and its associated GPU resources.
+   *
+   * When graph capture is enabled, the EP records GPU commands during initial runs and replays them
+   * on subsequent runs. This function releases the captured commands and associated GPU buffers
+   * for a specific graph annotation ID, freeing GPU memory.
+   *
+   * \param[in] session The OrtSession instance.
+   * \param[in] graph_annotation_id The annotation ID of the captured graph to release.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.26.
+   */
+  ORT_API2_STATUS(SessionReleaseGraph, _In_ OrtSession* session, _In_ int graph_annotation_id);
+
   /// @}
 };
 

include/onnxruntime/core/session/onnxruntime_ep_c_api.h

@@ -2514,6 +2514,20 @@ struct OrtEp {
    */
   ORT_API2_STATUS(ReplayGraph, _In_ OrtEp* this_ptr, _In_ int graph_annotation_id);
 
+  /** \brief Release a previously captured graph and its associated resources.
+   *
+   * Called when the caller no longer needs the captured graph for the given annotation ID.
+   * This allows the EP to free GPU buffers and other resources tied to this graph.
+   *
+   * \param[in] this_ptr The EP instance.
+   * \param[in] graph_annotation_id The annotation ID of the graph to release.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.26.
+   */
+  ORT_API2_STATUS(ReleaseGraph, _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

onnxruntime/core/providers/webgpu/allocator.cc

@@ -15,7 +15,7 @@ GpuBufferAllocator::GpuBufferAllocator(const BufferManager& buffer_manager, bool
                                                : OrtAllocatorType::OrtDeviceAllocator,
                         WebGpuDevice,
                         OrtMemTypeDefault)),
-      buffer_manager_{buffer_manager},
+      buffer_manager_{&buffer_manager},
       mapped_at_creation_{is_read_only_allocator && buffer_manager.SupportsUMA()} {
 }
 
@@ -29,12 +29,12 @@ void* GpuBufferAllocator::Alloc(size_t size) {
   wgpu::BufferUsage usage = mapped_at_creation_ ? wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapWrite
                                                 : wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect;
 
-  return buffer_manager_.Create(size, usage);
+  return buffer_manager_->Create(size, usage);
 }
 
 void GpuBufferAllocator::Free(void* p) {
   if (p != nullptr) {
-    buffer_manager_.Release(static_cast<WGPUBuffer>(p));
+    buffer_manager_->Release(static_cast<WGPUBuffer>(p));
     stats_.num_allocs--;
   }
 }
@@ -43,5 +43,9 @@ void GpuBufferAllocator::GetStats(AllocatorStats* stats) {
   *stats = stats_;
 }
 
+void GpuBufferAllocator::SetBufferManager(const BufferManager& buffer_manager) {
+  buffer_manager_ = &buffer_manager;
+}
+
 }  // namespace webgpu
 }  // namespace onnxruntime

onnxruntime/core/providers/webgpu/allocator.h

@@ -24,9 +24,12 @@ class GpuBufferAllocator : public IAllocator {
   virtual void Free(void* p) override;
   void GetStats(AllocatorStats* stats) override;
 
+  // Update the buffer manager used for allocations (for per-graph buffer isolation)
+  void SetBufferManager(const BufferManager& buffer_manager);
+
  private:
   AllocatorStats stats_;
-  const BufferManager& buffer_manager_;
+  const BufferManager* buffer_manager_;
   bool mapped_at_creation_;
 };
 

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

@@ -43,6 +43,7 @@ Ep::Ep(std::unique_ptr<IExecutionProvider> impl, Factory& factory, const OrtLogg
   IsGraphCaptureEnabled = IsGraphCaptureEnabledImpl;
   IsGraphCaptured = IsGraphCapturedImpl;
   ReplayGraph = ReplayGraphImpl;
+  ReleaseGraph = ReleaseGraphImpl;
   GetGraphCaptureNodeAssignmentPolicy = GetGraphCaptureNodeAssignmentPolicyImpl;
 }
 
@@ -279,6 +280,18 @@ OrtStatus* ORT_API_CALL Ep::ReplayGraphImpl(_In_ OrtEp* this_ptr, _In_ int graph
   EXCEPTION_TO_RETURNED_STATUS_END
 }
 
+OrtStatus* ORT_API_CALL Ep::ReleaseGraphImpl(_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()->ReleaseGraph(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);

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

@@ -75,6 +75,9 @@ class Ep : public onnxruntime::ep::adapter::Ep {
   static OrtStatus* ORT_API_CALL ReplayGraphImpl(_In_ OrtEp* this_ptr,
                                                  _In_ int graph_annotation_id) noexcept;
 
+  static OrtStatus* ORT_API_CALL ReleaseGraphImpl(_In_ OrtEp* this_ptr,
+                                                  _In_ int graph_annotation_id) noexcept;
+
   static OrtGraphCaptureNodeAssignmentPolicy ORT_API_CALL GetGraphCaptureNodeAssignmentPolicyImpl(
       _In_ const OrtEp* this_ptr) noexcept;
 

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

@@ -10,6 +10,7 @@
 #include "core/framework/execution_provider.h"
 #include "core/framework/config_options.h"
 #include "core/providers/webgpu/webgpu_provider_factory_creator.h"
+#include "core/providers/webgpu/webgpu_execution_provider.h"
 #include "core/providers/webgpu/webgpu_context.h"
 #include "core/providers/webgpu/allocator.h"
 
@@ -134,10 +135,13 @@ OrtStatus* ORT_API_CALL Factory::CreateEpImpl(
   static_cast<WebGpuExecutionProvider*>(webgpu_ep.get())->SetEpLogger(logger);
   auto factory = static_cast<Factory*>(this_ptr);
   const int context_id = webgpu_ep->GetDeviceId();
+  auto* webgpu_ep_ptr = static_cast<WebGpuExecutionProvider*>(webgpu_ep.get());
+  auto device_alloc = std::make_shared<webgpu::GpuBufferAllocator>(webgpu_ep_ptr->BufferManager(), false);
+  webgpu_ep_ptr->SetDeviceAllocator(device_alloc.get());
   Ep::Config webgpu_ep_config{
       CPUAllocator::DefaultInstance(),                                                                                              // CPU allocator
-      std::make_shared<webgpu::GpuBufferAllocator>(WebGpuContextFactory::GetContext(context_id).BufferManager(), false),            // default device allocator
-      std::make_shared<webgpu::GpuBufferAllocator>(WebGpuContextFactory::GetContext(context_id).InitializerBufferManager(), true),  // initializer device allocator
+      device_alloc,                                                                                                                 // default device allocator
+      std::make_shared<webgpu::GpuBufferAllocator>(WebGpuContextFactory::GetContext(context_id).InitializerBufferManager(), true),   // initializer device allocator
   };
   *ep = new Ep(std::move(webgpu_ep), *factory, *logger, webgpu_ep_config);
   return nullptr;

onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc

@@ -581,7 +581,7 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id,
   // If graph capture is enabled, create a dedicated buffer manager for graph mode
   if (enable_graph_capture_) {
     // Create buffer manager for graph capture mode with appropriate cache modes
-    graph_buffer_mgr_ = webgpu::BufferManagerFactory::Create(
+    graph_default_buffer_mgr_ = webgpu::BufferManagerFactory::Create(
         context_,
         webgpu::BufferCacheMode::Graph,
         webgpu::BufferCacheMode::GraphSimple,
@@ -599,11 +599,13 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id,
 }
 
 std::vector<AllocatorPtr> WebGpuExecutionProvider::CreatePreferredAllocators() {
+  auto default_alloc = std::make_unique<webgpu::GpuBufferAllocator>(BufferManager(), false);
+  default_gpu_allocator_ = default_alloc.get();
   return {
       // allocator for initializers
       std::make_unique<webgpu::GpuBufferAllocator>(context_.InitializerBufferManager(), true),
       // default allocator
-      std::make_unique<webgpu::GpuBufferAllocator>(BufferManager(), false),
+      std::move(default_alloc),
   };
 }
 
@@ -733,11 +735,17 @@ std::optional<bool> WebGpuExecutionProvider::ShouldConvertDataLayoutForOp(std::s
 }
 
 WebGpuExecutionProvider::~WebGpuExecutionProvider() {
-  // Release all resources associated with the captured graph
-  if (!captured_commands_.empty()) {
-    context_.ReleaseGraphResources(captured_commands_);
+  // Release all captured graphs and their associated resources
+  std::vector<int> graph_ids;
+  graph_ids.reserve(captured_graph_ids_.size());
+  for (int id : captured_graph_ids_) {
+    graph_ids.push_back(id);
   }
-  // The graph_buffer_mgr_ will be automatically cleaned up by unique_ptr
+  for (int id : graph_ids) {
+    (void)ReleaseGraph(id);
+  }
+  // Also release any per-graph buffer managers for graphs that were never fully captured
+  per_graph_buffer_mgrs_.clear();
 
   WebGpuContextFactory::ReleaseContext(context_id_);
 }
@@ -772,8 +780,24 @@ Status WebGpuExecutionProvider::OnRunStart(const onnxruntime::RunOptions& run_op
                   *graph_annotation_str);
     }
 
+    // Create a per-graph buffer manager on first encounter of each annotation ID
+    if (graph_annotation_id != -1) {
+      if (per_graph_buffer_mgrs_.find(graph_annotation_id) == per_graph_buffer_mgrs_.end()) {
+        per_graph_buffer_mgrs_[graph_annotation_id] = webgpu::BufferManagerFactory::Create(
+            context_,
+            webgpu::BufferCacheMode::Graph,
+            webgpu::BufferCacheMode::GraphSimple,
+            webgpu::BufferCacheMode::Disabled);
+      }
+      // Route allocator to this graph's buffer manager
+      if (default_gpu_allocator_) {
+        default_gpu_allocator_->SetBufferManager(*per_graph_buffer_mgrs_[graph_annotation_id]);
+      }
+    }
+
     if (graph_annotation_id != -1 && IsGraphCaptureAllowed() && !IsGraphCaptured(graph_annotation_id)) {
-      context_.CaptureBegin(&captured_commands_, *graph_buffer_mgr_);
+      auto& commands = captured_graphs_[graph_annotation_id];
+      context_.CaptureBegin(&commands, *per_graph_buffer_mgrs_[graph_annotation_id]);
     }
     m_current_graph_annotation_id = graph_annotation_id;
   }
@@ -787,7 +811,7 @@ Status WebGpuExecutionProvider::OnRunEnd(bool /* sync_stream */, const onnxrunti
   if (IsGraphCaptureEnabled() && !IsGraphCaptured(m_current_graph_annotation_id)) {
     if (m_current_graph_annotation_id != -1 && IsGraphCaptureAllowed()) {
       context_.CaptureEnd();
-      is_graph_captured_ = true;
+      captured_graph_ids_.insert(m_current_graph_annotation_id);
       ORT_RETURN_IF_ERROR(ReplayGraph(m_current_graph_annotation_id));
     } else {
       IncrementRegularRunCountBeforeGraphCapture();
@@ -808,6 +832,11 @@ Status WebGpuExecutionProvider::OnRunEnd(bool /* sync_stream */, const onnxrunti
   }
 #endif  // ENABLE_PIX_FOR_WEBGPU_EP
 
+  // Reset allocator to default buffer manager after run completes
+  if (default_gpu_allocator_ && graph_default_buffer_mgr_) {
+    default_gpu_allocator_->SetBufferManager(*graph_default_buffer_mgr_);
+  }
+
   if (context_.ValidationMode() >= ValidationMode::Basic) {
     return context_.PopErrorScope();
   } else {
@@ -820,7 +849,7 @@ bool WebGpuExecutionProvider::IsGraphCaptureEnabled() const {
 }
 
 bool WebGpuExecutionProvider::IsGraphCaptured(int graph_annotation_id) const {
-  return is_graph_captured_ && graph_annotation_id != -1;
+  return graph_annotation_id != -1 && captured_graph_ids_.count(graph_annotation_id) > 0;
 }
 
 Status WebGpuExecutionProvider::ReplayGraph(int graph_annotation_id) {
@@ -829,27 +858,59 @@ Status WebGpuExecutionProvider::ReplayGraph(int graph_annotation_id) {
   if (session_profiler_ && session_profiler_->Enabled()) {
     context_.StartProfiling();
   }
-  context_.Replay(captured_commands_, *graph_buffer_mgr_);
+  context_.Replay(captured_graphs_.at(graph_annotation_id), *per_graph_buffer_mgrs_.at(graph_annotation_id));
   if (session_profiler_ && session_profiler_->Enabled()) {
     // Session-level profiling: collect into profiler's own events storage.
     context_.CollectProfilingData(session_profiler_->GpuEvents());
   }
   return Status::OK();
 }
 
+Status WebGpuExecutionProvider::ReleaseGraph(int graph_annotation_id) {
+  // Release captured commands
+  auto cmd_it = captured_graphs_.find(graph_annotation_id);
+  if (cmd_it != captured_graphs_.end()) {
+    if (!cmd_it->second.empty()) {
+      context_.ReleaseGraphResources(cmd_it->second);
+    }
+    captured_graphs_.erase(cmd_it);
+  }
+
+  // Remove from captured set
+  captured_graph_ids_.erase(graph_annotation_id);
+
+  // Release per-graph buffer manager (destroys cached buffers)
+  per_graph_buffer_mgrs_.erase(graph_annotation_id);
+
+  // Clean up run count tracking
+  graph_id_to_run_count_.erase(graph_annotation_id);
+
+  return Status::OK();
+}
+
 webgpu::BufferManager& WebGpuExecutionProvider::BufferManager() const {
-  if (graph_buffer_mgr_) {
-    return *graph_buffer_mgr_;
+  // Use per-graph buffer manager if one exists for the current annotation ID
+  if (m_current_graph_annotation_id != 0 && m_current_graph_annotation_id != -1) {
+    auto it = per_graph_buffer_mgrs_.find(m_current_graph_annotation_id);
+    if (it != per_graph_buffer_mgrs_.end()) {
+      return *it->second;
+    }
+  }
+  // Fall back to default graph buffer manager (warmup runs) or context buffer manager
+  if (graph_default_buffer_mgr_) {
+    return *graph_default_buffer_mgr_;
   } else {
     return context_.BufferManager();
   }
 }
 
 bool WebGpuExecutionProvider::IsGraphCaptureAllowed() const {
-  return regular_run_count_before_graph_capture_ >= min_num_runs_before_cuda_graph_capture_;
+  auto it = graph_id_to_run_count_.find(m_current_graph_annotation_id);
+  int run_count = (it != graph_id_to_run_count_.end()) ? it->second : 0;
+  return run_count >= min_num_runs_before_cuda_graph_capture_;
 }
 
 void WebGpuExecutionProvider::IncrementRegularRunCountBeforeGraphCapture() {
-  ++regular_run_count_before_graph_capture_;
+  ++graph_id_to_run_count_[m_current_graph_annotation_id];
 }
 }  // namespace onnxruntime

onnxruntime/core/providers/webgpu/webgpu_execution_provider.h

@@ -7,6 +7,8 @@
 #include <span>
 #include <string>
 #include <memory>
+#include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "core/framework/execution_provider.h"
@@ -97,11 +99,14 @@ class WebGpuExecutionProvider : public IExecutionProvider {
   bool IsGraphCaptureEnabled() const override;
   bool IsGraphCaptured(int graph_annotation_id) const override;
   Status ReplayGraph(int graph_annotation_id) override;
+  Status ReleaseGraph(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_; }
+  // Set the device allocator pointer so we can call SetBufferManager on it during OnRunStart/OnRunEnd
+  void SetDeviceAllocator(webgpu::GpuBufferAllocator* allocator) { default_gpu_allocator_ = allocator; }
   std::span<const std::string> GetForceCpuNodeNames() const { return force_cpu_node_names_; }
   uint32_t MultiRotaryCacheConcatOffset() const { return multi_rotary_cache_concat_offset_; }
 
@@ -126,24 +131,35 @@ class WebGpuExecutionProvider : public IExecutionProvider {
   bool enable_graph_capture_ = false;
   bool enable_int64_ = false;
   uint32_t multi_rotary_cache_concat_offset_ = 0;
-  bool is_graph_captured_ = false;
-  int regular_run_count_before_graph_capture_ = 0;
+  std::unordered_map<int, int> graph_id_to_run_count_;
   const int min_num_runs_before_cuda_graph_capture_ = 1;  // Required regular runs before graph capture for any necessary allocations.
   int m_current_graph_annotation_id = 0;
 
 #if defined(ENABLE_PIX_FOR_WEBGPU_EP)
   std::unique_ptr<WebGpuPIXFrameGenerator> pix_frame_generator_ = nullptr;
 #endif  // ENABLE_PIX_FOR_WEBGPU_EP
 
-  // Buffer manager specifically for graph capture mode
-  std::unique_ptr<webgpu::BufferManager> graph_buffer_mgr_ = nullptr;
+  // Default buffer manager for graph capture mode (used during warmup runs
+  // and as the stable reference target for GpuBufferAllocator)
+  std::unique_ptr<webgpu::BufferManager> graph_default_buffer_mgr_ = nullptr;
 
-  // Store captured commands directly in the EP instead of in WebGpuContext
-  std::vector<webgpu::CapturedCommandInfo> captured_commands_;
+  // Per-graph buffer managers keyed by annotation ID.
+  // Each captured graph gets its own buffer manager so that buffer caches
+  // are isolated between different generators.
+  std::unordered_map<int, std::unique_ptr<webgpu::BufferManager>> per_graph_buffer_mgrs_;
+
+  // Store captured commands per graph annotation ID
+  std::unordered_map<int, std::vector<webgpu::CapturedCommandInfo>> captured_graphs_;
+  // Track which graph annotation IDs have completed capture
+  std::unordered_set<int> captured_graph_ids_;
 
   // Allocator for prepacked weights (uses buffers without mapping)
   AllocatorPtr prepack_allocator_;
 
+  // Raw pointer to the default GPU allocator (owned by the framework via CreatePreferredAllocators)
+  // Used to swap the buffer manager for per-graph isolation
+  webgpu::GpuBufferAllocator* default_gpu_allocator_ = nullptr;
+
 #if defined(ORT_USE_EP_API_ADAPTERS)
   std::unique_ptr<onnxruntime::ep::adapter::Logger> ep_logger_;
 #endif