From f1d977f8703d68b0522659b4c558a424107f37ef Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 3 Jun 2026 17:31:45 +0800 Subject: [PATCH 1/4] webgpu: Add session-level buffer pool for graph capture reuse When graph capture is enabled, each generator instance owns a per-graph BufferManager whose cache is discarded when the generator is destroyed. For workloads that repeatedly create and destroy generators on the same session (e.g., GenAI's per-request generators), this means every new generator has to reallocate all storage and uniform buffers from scratch, inflating cold-start cost and GPU memory churn. This change introduces a SessionBufferPool owned by the session. When a retiring BufferManager is released, its cached storage and uniform buffers are donated to the pool; the next BufferManager seeded from the session absorbs those buffers, skipping reallocation entirely. The pool capacity is controlled by a new provider option "ep.webgpuexecutionprovider.sessionBufferPoolGenerations" (defaults to disabled). The pool evicts the oldest slot when full, keeping the freshest distribution of buffer shapes. Verified with GenAI multi-generator scripts on phi4: subsequent generators report zero cache misses for both storage and uniform caches and produce coherent output across max_length changes and overlapping generator lifetimes. --- .../core/providers/webgpu/buffer_manager.cc | 50 +++++++++++++ .../core/providers/webgpu/buffer_manager.h | 25 +++++++ .../providers/webgpu/session_buffer_pool.cc | 75 +++++++++++++++++++ .../providers/webgpu/session_buffer_pool.h | 54 +++++++++++++ .../webgpu/webgpu_execution_provider.cc | 22 +++++- .../webgpu/webgpu_execution_provider.h | 10 +++ .../webgpu/webgpu_provider_factory.cc | 14 ++++ .../webgpu/webgpu_provider_options.h | 1 + 8 files changed, 249 insertions(+), 2 deletions(-) create mode 100644 onnxruntime/core/providers/webgpu/session_buffer_pool.cc create mode 100644 onnxruntime/core/providers/webgpu/session_buffer_pool.h diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.cc b/onnxruntime/core/providers/webgpu/buffer_manager.cc index fadb425c2ad6d..45b0447d2d9be 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.cc +++ b/onnxruntime/core/providers/webgpu/buffer_manager.cc @@ -304,6 +304,25 @@ class GraphCacheManager : public IBufferCacheManager { // no-op - buffers are already in buckets_ } + std::vector> ExtractCachedBuffers() override { + std::vector> result; + for (auto& pair : buckets_) { + for (auto& buffer : pair.second) { + result.emplace_back(pair.first, buffer); + } + pair.second.clear(); + } + return result; + } + + void AbsorbCachedBuffers(std::vector>&& buffers) override { + for (auto& entry : buffers) { + if (entry.second) { + ReleaseBuffer(entry.second); + } + } + } + ~GraphCacheManager() { for (auto& pair : buckets_) { for (auto& buffer : pair.second) { @@ -386,6 +405,37 @@ class GraphSimpleCacheManager : public IBufferCacheManager { } } + std::vector> ExtractCachedBuffers() override { + // Donation is expected after captured commands have been released and any + // in-flight work has completed; all three containers therefore hold buffers + // no longer referenced by the device. + std::vector> result; + for (auto& pair : buffers_) { + for (auto& buffer : pair.second) { + result.emplace_back(pair.first, buffer); + } + pair.second.clear(); + } + buffers_.clear(); + for (auto& buffer : pending_buffers_) { + result.emplace_back(static_cast(wgpuBufferGetSize(buffer)), buffer); + } + pending_buffers_.clear(); + for (auto& buffer : captured_buffers_) { + result.emplace_back(static_cast(wgpuBufferGetSize(buffer)), buffer); + } + captured_buffers_.clear(); + return result; + } + + void AbsorbCachedBuffers(std::vector>&& buffers) override { + for (auto& entry : buffers) { + if (entry.second) { + buffers_[entry.first].emplace_back(entry.second); + } + } + } + protected: std::map> buffers_; std::vector pending_buffers_; diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.h b/onnxruntime/core/providers/webgpu/buffer_manager.h index 05a79726d71e4..b0309fda6d574 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.h +++ b/onnxruntime/core/providers/webgpu/buffer_manager.h @@ -4,6 +4,8 @@ #pragma once #include +#include +#include #include "core/providers/webgpu/webgpu_external_header.h" @@ -60,6 +62,24 @@ class IBufferCacheManager { // when a stream refresh is requested virtual void OnRefresh(GraphCaptureState graph_capture_state) = 0; + + // Extract all cached buffers from this manager, transferring ownership to the + // caller. The cache's internal containers are cleared (but bucket keys, if any, + // are preserved). Default returns empty; only graph-mode caches implement this. + virtual std::vector> ExtractCachedBuffers() { + return {}; + } + + // Accept buffers donated from another cache and take ownership of them. Caches + // that cannot store the buffers must release them via wgpuBufferRelease (the + // default below) to avoid leaks. + virtual void AbsorbCachedBuffers(std::vector>&& buffers) { + for (auto& entry : buffers) { + if (entry.second) { + wgpuBufferRelease(entry.second); + } + } + } }; // @@ -76,6 +96,11 @@ class BufferManager { void Download(WGPUBuffer src, void* dst, size_t size) const; void RefreshPendingBuffers(GraphCaptureState graph_capture_state) const; + // Direct access to the underlying cache managers. Used by SessionBufferPool to + // donate/seed buffers across per-graph BufferManager lifetimes. + IBufferCacheManager& StorageCache() const { return *storage_cache_; } + IBufferCacheManager& UniformCache() const { return *uniform_cache_; } + private: IBufferCacheManager& GetCacheManager(wgpu::BufferUsage usage) const; IBufferCacheManager& GetCacheManager(WGPUBuffer buffer) const; diff --git a/onnxruntime/core/providers/webgpu/session_buffer_pool.cc b/onnxruntime/core/providers/webgpu/session_buffer_pool.cc new file mode 100644 index 0000000000000..bad86a04a995a --- /dev/null +++ b/onnxruntime/core/providers/webgpu/session_buffer_pool.cc @@ -0,0 +1,75 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/session_buffer_pool.h" + +#include "core/providers/webgpu/buffer_manager.h" + +namespace onnxruntime { +namespace webgpu { + +namespace { +void ReleaseSlotBuffers(std::vector>& entries) { + for (auto& entry : entries) { + if (entry.second) { + wgpuBufferRelease(entry.second); + } + } + entries.clear(); +} +} // namespace + +SessionBufferPool::SessionBufferPool(size_t max_generations) + : max_generations_{max_generations} { + slots_.reserve(max_generations_); +} + +SessionBufferPool::~SessionBufferPool() { + Clear(); +} + +void SessionBufferPool::Donate(BufferManager& retiring_mgr) { + if (max_generations_ == 0) { + return; + } + + Slot slot; + slot.storage = retiring_mgr.StorageCache().ExtractCachedBuffers(); + slot.uniform = retiring_mgr.UniformCache().ExtractCachedBuffers(); + + if (slot.storage.empty() && slot.uniform.empty()) { + return; + } + + // Evict the oldest slot if at capacity so the freshest buffers (which most + // accurately reflect the current per-generator shape distribution) are kept. + while (slots_.size() >= max_generations_) { + auto& victim = slots_.front(); + ReleaseSlotBuffers(victim.storage); + ReleaseSlotBuffers(victim.uniform); + slots_.erase(slots_.begin()); + } + + slots_.emplace_back(std::move(slot)); +} + +void SessionBufferPool::SeedInto(BufferManager& new_mgr) { + if (slots_.empty()) { + return; + } + Slot slot = std::move(slots_.back()); + slots_.pop_back(); + new_mgr.StorageCache().AbsorbCachedBuffers(std::move(slot.storage)); + new_mgr.UniformCache().AbsorbCachedBuffers(std::move(slot.uniform)); +} + +void SessionBufferPool::Clear() { + for (auto& slot : slots_) { + ReleaseSlotBuffers(slot.storage); + ReleaseSlotBuffers(slot.uniform); + } + slots_.clear(); +} + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/session_buffer_pool.h b/onnxruntime/core/providers/webgpu/session_buffer_pool.h new file mode 100644 index 0000000000000..edec2b613dfd6 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/session_buffer_pool.h @@ -0,0 +1,54 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "core/providers/webgpu/webgpu_external_header.h" + +namespace onnxruntime { +namespace webgpu { + +class BufferManager; + +// SessionBufferPool retains buffers from retired per-graph BufferManagers so +// that subsequent generators on the same session can reuse them instead of +// allocating from the device. Scoped per WebGpuExecutionProvider (per session) +// because intermediate buffer shapes are model-dependent. +class SessionBufferPool { + public: + explicit SessionBufferPool(size_t max_generations); + + ~SessionBufferPool(); + + // Move freed buffers from a retiring per-graph BufferManager into the pool. + // When the pool is at capacity, the oldest slot is evicted (its buffers + // released to the device) so the freshest buffers are always retained. This + // lets the pool adapt when intermediate buffer shapes change between + // generators (for example when max_length differs). + void Donate(BufferManager& retiring_mgr); + + // Pre-populate a newly created per-graph BufferManager with one slot worth of + // pooled buffers (LIFO). No-op if the pool is empty. + void SeedInto(BufferManager& new_mgr); + + // Release all pooled buffers. Called on session teardown. + void Clear(); + + size_t Size() const { return slots_.size(); } + size_t Capacity() const { return max_generations_; } + + private: + struct Slot { + std::vector> storage; + std::vector> uniform; + }; + std::vector slots_; + size_t max_generations_; +}; + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 7e11ddf6b13a0..a588fee464260 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -582,6 +582,10 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id, multi_rotary_cache_concat_offset_{config.multi_rotary_cache_concat_offset}, prepack_allocator_{std::make_shared( [this]() -> const webgpu::BufferManager& { return context_.InitializerBufferManager(); }, false)} { + if (enable_graph_capture_ && config.session_buffer_pool_generations > 0) { + session_buffer_pool_ = std::make_unique( + config.session_buffer_pool_generations); + } if (config.enable_pix_capture) { #if defined(ENABLE_PIX_FOR_WEBGPU_EP) // set pix frame generator @@ -748,6 +752,11 @@ WebGpuExecutionProvider::~WebGpuExecutionProvider() { // but no entries in captured_graphs_ (edge case cleanup) per_graph_buffer_mgrs_.clear(); + // Release pooled buffers before the WebGpuContext is released. + if (session_buffer_pool_) { + session_buffer_pool_->Clear(); + } + WebGpuContextFactory::ReleaseContext(context_id_); } @@ -792,6 +801,9 @@ Status WebGpuExecutionProvider::OnRunStart(const onnxruntime::RunOptions& run_op webgpu::BufferCacheMode::Graph, webgpu::BufferCacheMode::GraphSimple, webgpu::BufferCacheMode::Disabled); + if (session_buffer_pool_) { + session_buffer_pool_->SeedInto(*it->second); + } } graph_buffer_mgr_active_ = true; @@ -877,8 +889,14 @@ Status WebGpuExecutionProvider::ReleaseCapturedGraph(int graph_annotation_id) { // 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); + // Release per-graph buffer manager (donate to session pool if enabled; otherwise destroy) + auto mgr_it = per_graph_buffer_mgrs_.find(graph_annotation_id); + if (mgr_it != per_graph_buffer_mgrs_.end()) { + if (session_buffer_pool_ && mgr_it->second) { + session_buffer_pool_->Donate(*mgr_it->second); + } + per_graph_buffer_mgrs_.erase(mgr_it); + } // Clean up run count tracking graph_id_to_run_count_.erase(graph_annotation_id); diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h index 92d1ebfd36c79..1dc7de0c21e9f 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h @@ -16,6 +16,7 @@ #include "core/graph/constants.h" #include "core/providers/providers.h" #include "core/providers/webgpu/buffer_manager.h" +#include "core/providers/webgpu/session_buffer_pool.h" #if defined(ENABLE_PIX_FOR_WEBGPU_EP) #include "core/providers/webgpu/webgpu_pix_frame_generator.h" @@ -46,6 +47,10 @@ struct WebGpuExecutionProviderConfig { bool enable_pix_capture{false}; // PIX capture is disabled by default bool enable_int64{false}; // int64 ops are not enabled by default uint32_t multi_rotary_cache_concat_offset{0}; // offset for concatenated multi rotary cache (0 = disabled) + // Number of generations of buffers to retain in the per-session pool for reuse + // across captured-graph lifetimes. 0 disables pooling. Default 1 caches one + // generator's worth of intermediate buffers. + size_t session_buffer_pool_generations{1}; std::vector force_cpu_node_names{}; }; @@ -144,6 +149,11 @@ class WebGpuExecutionProvider : public IExecutionProvider { // are isolated between different generators. std::unordered_map> per_graph_buffer_mgrs_; + // Per-session pool of buffers donated by retired per-graph BufferManagers, + // seeded into new per-graph BufferManagers to avoid device allocations for + // identically-shaped intermediate tensors across generators. + std::unique_ptr session_buffer_pool_; + // Store captured commands per graph annotation ID std::unordered_map> captured_graphs_; // Track which graph annotation IDs have completed capture diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index 16899370e47f1..d593bf0e8536f 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -61,6 +61,19 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) } } + if (std::string pool_generations_str; + config_options.TryGetConfigEntry(kSessionBufferPoolGenerations, pool_generations_str)) { + size_t pool_generations = 0; + auto result = std::from_chars(pool_generations_str.data(), + pool_generations_str.data() + pool_generations_str.size(), + pool_generations); + if (result.ec == std::errc{}) { + webgpu_ep_config.session_buffer_pool_generations = pool_generations; + } else { + ORT_THROW("Invalid sessionBufferPoolGenerations value: ", pool_generations_str, ". Must be a non-negative integer."); + } + } + std::string enable_int64_str; if (config_options.TryGetConfigEntry(kEnableInt64, enable_int64_str)) { if (enable_int64_str == kEnableInt64_ON) { @@ -122,6 +135,7 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) LOGS_DEFAULT(VERBOSE) << "WebGPU EP pix capture enable: " << webgpu_ep_config.enable_pix_capture; LOGS_DEFAULT(VERBOSE) << "WebGPU EP enable int64: " << webgpu_ep_config.enable_int64; LOGS_DEFAULT(VERBOSE) << "WebGPU EP multi rotary cache concat offset: " << webgpu_ep_config.multi_rotary_cache_concat_offset; + LOGS_DEFAULT(VERBOSE) << "WebGPU EP session buffer pool generations: " << webgpu_ep_config.session_buffer_pool_generations; return webgpu_ep_config; } diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_options.h b/onnxruntime/core/providers/webgpu/webgpu_provider_options.h index d2faccdb8c4a5..4cab81dafe9c7 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_options.h +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_options.h @@ -11,6 +11,7 @@ namespace options { constexpr const char* kPreferredLayout = "ep.webgpuexecutionprovider.preferredLayout"; constexpr const char* kEnableGraphCapture = "ep.webgpuexecutionprovider.enableGraphCapture"; +constexpr const char* kSessionBufferPoolGenerations = "ep.webgpuexecutionprovider.sessionBufferPoolGenerations"; constexpr const char* kEnableInt64 = "ep.webgpuexecutionprovider.enableInt64"; constexpr const char* kMultiRotaryCacheConcatOffset = "ep.webgpuexecutionprovider.multiRotaryCacheConcatOffset"; From 9a11120dbe0eca5a685a7cb999cf284ab1d9d9ba Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 3 Jun 2026 18:16:22 +0800 Subject: [PATCH 2/4] Address PR review: strict option parsing, non-const cache accessors - webgpu_provider_factory.cc: require std::from_chars to consume the full string for sessionBufferPoolGenerations so values like "1foo" are rejected instead of silently parsed as 1. - buffer_manager.h: drop const from StorageCache()/UniformCache() so the mutable cache references can no longer be obtained through a const BufferManager&. - session_buffer_pool.cc: drop slots_.reserve(max_generations_) to avoid a large up-front allocation when the option is set to an extreme value; slots grow on demand instead. --- onnxruntime/core/providers/webgpu/buffer_manager.h | 4 ++-- onnxruntime/core/providers/webgpu/session_buffer_pool.cc | 1 - .../core/providers/webgpu/webgpu_provider_factory.cc | 8 ++++---- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.h b/onnxruntime/core/providers/webgpu/buffer_manager.h index b0309fda6d574..9bd1f54711971 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.h +++ b/onnxruntime/core/providers/webgpu/buffer_manager.h @@ -98,8 +98,8 @@ class BufferManager { // Direct access to the underlying cache managers. Used by SessionBufferPool to // donate/seed buffers across per-graph BufferManager lifetimes. - IBufferCacheManager& StorageCache() const { return *storage_cache_; } - IBufferCacheManager& UniformCache() const { return *uniform_cache_; } + IBufferCacheManager& StorageCache() { return *storage_cache_; } + IBufferCacheManager& UniformCache() { return *uniform_cache_; } private: IBufferCacheManager& GetCacheManager(wgpu::BufferUsage usage) const; diff --git a/onnxruntime/core/providers/webgpu/session_buffer_pool.cc b/onnxruntime/core/providers/webgpu/session_buffer_pool.cc index bad86a04a995a..5d647f33b12d0 100644 --- a/onnxruntime/core/providers/webgpu/session_buffer_pool.cc +++ b/onnxruntime/core/providers/webgpu/session_buffer_pool.cc @@ -21,7 +21,6 @@ void ReleaseSlotBuffers(std::vector>& entries) { SessionBufferPool::SessionBufferPool(size_t max_generations) : max_generations_{max_generations} { - slots_.reserve(max_generations_); } SessionBufferPool::~SessionBufferPool() { diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index d593bf0e8536f..9986b2f4f94cd 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -64,10 +64,10 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) if (std::string pool_generations_str; config_options.TryGetConfigEntry(kSessionBufferPoolGenerations, pool_generations_str)) { size_t pool_generations = 0; - auto result = std::from_chars(pool_generations_str.data(), - pool_generations_str.data() + pool_generations_str.size(), - pool_generations); - if (result.ec == std::errc{}) { + const char* begin = pool_generations_str.data(); + const char* end = begin + pool_generations_str.size(); + auto result = std::from_chars(begin, end, pool_generations); + if (result.ec == std::errc{} && result.ptr == end) { webgpu_ep_config.session_buffer_pool_generations = pool_generations; } else { ORT_THROW("Invalid sessionBufferPoolGenerations value: ", pool_generations_str, ". Must be a non-negative integer."); From 3d5655840255a0fcb2d15374857d68a14862184b Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 10 Jun 2026 10:51:52 +0800 Subject: [PATCH 3/4] Document SessionBufferPool invariants and provider option Address review feedback with documentation-only changes: - SessionBufferPool: thread-safety contract notes that the pool relies on InferenceSession::session_mutex_ serializing Run and ReleaseCapturedGraph for non-concurrent EPs. Donate's preamble explains why donation is safe without an explicit GPU fence (Run/Release serialization, OnRunEnd flush, single-queue ordering) and what would invalidate that reasoning. SeedInto notes that the receiving manager may hold up to one slot's worth of stale buffer sizes until it is itself donated or destroyed. The early-exit in Donate calls out that caches are already empty when it returns. - WebGpuExecutionProviderConfig: clearer description of session_buffer_pool_generations, defining "generation" in plain language and explaining when the default of 1 eliminates reallocation. No behavior change. --- .../core/providers/webgpu/session_buffer_pool.cc | 15 +++++++++++++++ .../core/providers/webgpu/session_buffer_pool.h | 13 ++++++++++++- .../providers/webgpu/webgpu_execution_provider.h | 10 +++++++--- 3 files changed, 34 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/session_buffer_pool.cc b/onnxruntime/core/providers/webgpu/session_buffer_pool.cc index 5d647f33b12d0..1c0ffebab3590 100644 --- a/onnxruntime/core/providers/webgpu/session_buffer_pool.cc +++ b/onnxruntime/core/providers/webgpu/session_buffer_pool.cc @@ -28,6 +28,19 @@ SessionBufferPool::~SessionBufferPool() { } void SessionBufferPool::Donate(BufferManager& retiring_mgr) { + // Safe to take ownership of the retiring manager's buffers because: + // 1. InferenceSession::session_mutex_ serializes Run and + // ReleaseCapturedGraph for the WebGPU EP, so no Run is in progress on + // this session while we donate. + // 2. OnRunEnd flushes the queue (and, in ValidationMode::Basic+, fences + // via PopErrorScope) before Run returns, so prior captured-graph work + // has been submitted. + // 3. All work goes through a single WebGPU device queue, which orders + // future submits after the prior submitted work, so a subsequent + // generator's first Submit using these buffer handles happens-after + // any in-flight read/write on the GPU timeline. + // Donating across a different queue or after a host-side readback on + // donated buffers would invalidate (3) and require an explicit fence here. if (max_generations_ == 0) { return; } @@ -36,6 +49,8 @@ void SessionBufferPool::Donate(BufferManager& retiring_mgr) { slot.storage = retiring_mgr.StorageCache().ExtractCachedBuffers(); slot.uniform = retiring_mgr.UniformCache().ExtractCachedBuffers(); + // The caches are now empty either way; the retiring manager is about to be + // destroyed, so the early-exit only avoids pushing an empty slot. if (slot.storage.empty() && slot.uniform.empty()) { return; } diff --git a/onnxruntime/core/providers/webgpu/session_buffer_pool.h b/onnxruntime/core/providers/webgpu/session_buffer_pool.h index edec2b613dfd6..a4b7d038cce04 100644 --- a/onnxruntime/core/providers/webgpu/session_buffer_pool.h +++ b/onnxruntime/core/providers/webgpu/session_buffer_pool.h @@ -18,6 +18,14 @@ class BufferManager; // that subsequent generators on the same session can reuse them instead of // allocating from the device. Scoped per WebGpuExecutionProvider (per session) // because intermediate buffer shapes are model-dependent. +// +// Thread safety: not internally synchronized. Donate / SeedInto / Clear must +// not be called concurrently. The WebGPU EP relies on the fact that +// InferenceSession serializes Run and ReleaseCapturedGraph under +// session_mutex_ for execution providers where ConcurrentRunSupported() is +// false, which serializes the only call sites (OnRunStart -> SeedInto and +// ReleaseCapturedGraph -> Donate). Clear runs from the EP destructor after +// all Run / ReleaseCapturedGraph calls have returned. class SessionBufferPool { public: explicit SessionBufferPool(size_t max_generations); @@ -32,7 +40,10 @@ class SessionBufferPool { void Donate(BufferManager& retiring_mgr); // Pre-populate a newly created per-graph BufferManager with one slot worth of - // pooled buffers (LIFO). No-op if the pool is empty. + // pooled buffers (LIFO). No-op if the pool is empty. The receiving manager + // may end up holding at most one donated slot's worth of buffer sizes that + // the new generator never requests; those sit in the receiving manager's + // caches until that manager is itself donated or destroyed. void SeedInto(BufferManager& new_mgr); // Release all pooled buffers. Called on session teardown. diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h index 1dc7de0c21e9f..790e48c25dac9 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h @@ -47,9 +47,13 @@ struct WebGpuExecutionProviderConfig { bool enable_pix_capture{false}; // PIX capture is disabled by default bool enable_int64{false}; // int64 ops are not enabled by default uint32_t multi_rotary_cache_concat_offset{0}; // offset for concatenated multi rotary cache (0 = disabled) - // Number of generations of buffers to retain in the per-session pool for reuse - // across captured-graph lifetimes. 0 disables pooling. Default 1 caches one - // generator's worth of intermediate buffers. + // Number of retired generators' worth of buffers to keep across + // ReleaseCapturedGraph. A "generation" here is the buffer cache of one + // captured-graph BufferManager; donated buffers are reused by subsequent + // generators on the same session to avoid reallocating from the device. + // 0 disables pooling; the default 1 retains the most recent generation, + // which is sufficient to eliminate reallocation in steady-state per-request + // GenAI usage where one generator retires just before the next is created. size_t session_buffer_pool_generations{1}; std::vector force_cpu_node_names{}; }; From c7f6e58c3f848b8aaf11f527f46d43cfacd8d433 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 10 Jun 2026 11:19:02 +0800 Subject: [PATCH 4/4] Move sessionBufferPoolGenerations description to provider_options.h External consumers configure the pool through the option key in webgpu_provider_options.h, but until now that header carried no description and the key name alone could be misread as "number of retired generators". Add the same three-line description above the key so it is visible where the option is set, and restore the original concise comment above the internal field in webgpu_execution_provider.h. --- .../core/providers/webgpu/webgpu_execution_provider.h | 10 +++------- .../core/providers/webgpu/webgpu_provider_options.h | 3 +++ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h index 790e48c25dac9..1dc7de0c21e9f 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h @@ -47,13 +47,9 @@ struct WebGpuExecutionProviderConfig { bool enable_pix_capture{false}; // PIX capture is disabled by default bool enable_int64{false}; // int64 ops are not enabled by default uint32_t multi_rotary_cache_concat_offset{0}; // offset for concatenated multi rotary cache (0 = disabled) - // Number of retired generators' worth of buffers to keep across - // ReleaseCapturedGraph. A "generation" here is the buffer cache of one - // captured-graph BufferManager; donated buffers are reused by subsequent - // generators on the same session to avoid reallocating from the device. - // 0 disables pooling; the default 1 retains the most recent generation, - // which is sufficient to eliminate reallocation in steady-state per-request - // GenAI usage where one generator retires just before the next is created. + // Number of generations of buffers to retain in the per-session pool for reuse + // across captured-graph lifetimes. 0 disables pooling. Default 1 caches one + // generator's worth of intermediate buffers. size_t session_buffer_pool_generations{1}; std::vector force_cpu_node_names{}; }; diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_options.h b/onnxruntime/core/providers/webgpu/webgpu_provider_options.h index 4cab81dafe9c7..a0305cd1f01e6 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_options.h +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_options.h @@ -11,6 +11,9 @@ namespace options { constexpr const char* kPreferredLayout = "ep.webgpuexecutionprovider.preferredLayout"; constexpr const char* kEnableGraphCapture = "ep.webgpuexecutionprovider.enableGraphCapture"; +// Number of generations of buffers to retain in the per-session pool for reuse +// across captured-graph lifetimes. 0 disables pooling. Default 1 caches one +// generator's worth of intermediate buffers. constexpr const char* kSessionBufferPoolGenerations = "ep.webgpuexecutionprovider.sessionBufferPoolGenerations"; constexpr const char* kEnableInt64 = "ep.webgpuexecutionprovider.enableInt64"; constexpr const char* kMultiRotaryCacheConcatOffset = "ep.webgpuexecutionprovider.multiRotaryCacheConcatOffset";