webgpu: Add session-level buffer pool for graph capture reuse - #28761
Conversation
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.
There was a problem hiding this comment.
Pull request overview
This PR adds a per-session WebGPU buffer pooling mechanism to improve graph-capture reuse across generator lifetimes by retaining a small number of recently retired per-graph buffer caches and seeding them into newly created per-graph BufferManager instances.
Changes:
- Introduces
webgpu::SessionBufferPoolto retain and recycle storage/uniform buffer caches across captured-graph lifetimes. - Adds a new provider option
ep.webgpuexecutionprovider.sessionBufferPoolGenerationsand parses/logs it in the WebGPU EP factory/config. - Wires pooling into graph-capture lifecycle: seed buffers on per-graph
BufferManagercreation and donate buffers onReleaseCapturedGraph.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/core/providers/webgpu/webgpu_provider_options.h | Adds config key for session buffer pool generations. |
| onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc | Parses sessionBufferPoolGenerations and logs configured value. |
| onnxruntime/core/providers/webgpu/webgpu_execution_provider.h | Adds config field and EP member for the session-level buffer pool. |
| onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc | Creates/clears pool and integrates donate/seed with graph capture lifecycle. |
| onnxruntime/core/providers/webgpu/session_buffer_pool.h | New pool type definition and slot structure for storage/uniform buffers. |
| onnxruntime/core/providers/webgpu/session_buffer_pool.cc | New implementation for donate/seed/clear and buffer releasing. |
| onnxruntime/core/providers/webgpu/buffer_manager.h | Exposes cache managers and adds extract/absorb APIs to cache interface. |
| onnxruntime/core/providers/webgpu/buffer_manager.cc | Implements extract/absorb for graph-mode cache managers to enable pooling. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- 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.
|
@hariharans29 @guschmue Please take a look, thanks. |
Verdict: Approve with comments (please address #1 and #2 before merge)Nice optimization — the layered approach (per-cache extract/absorb + per-session pool with bounded generations) is the right shape, LIFO consume / FIFO evict is the right pair, and the empirical phi4 numbers ( Two correctness questions worth answering before merge, plus some smaller observations. Please confirm before merge1. GPU work in-flight when donation happensThis is the question I most want answered. Looking at void WebGpuContext::ReleaseGraphResources(std::vector<webgpu::CapturedCommandInfo>& captured_commands) {
for (auto& command : captured_commands) {
if (command.bind_group != nullptr) {
wgpuBindGroupRelease(command.bind_group);
...
}
if (command.bind_group_layout != nullptr) {
wgpuBindGroupLayoutRelease(command.bind_group_layout);
...
}
}
}It releases CPU-side refs to bind groups and layouts. It does not call The bind groups hold strong refs to those buffers, so the buffer handles stay valid until the GPU finishes — that's fine for memory safety. But the buffer contents are not protected. If we donate those buffers, the next generator's In practice this probably works today because:
But neither is a guarantee. Two options: a. Add an explicit b. Document the contract in The existing comment in 2. Thread-safety contract of the pool
Either add a Smaller observations3. Default
|
|
Thanks for the thorough review @hariharans29! Going point by point. 1. GPU work in-flight at donation — no code change neededThe risk is narrower than it looks because of how WebGPU's queue model interacts with the existing
The only failure mode is if a future change either (a) moves part of the workflow to a different queue or (b) introduces a host-side 2. Thread safety of the pool — no code change neededFor the same reason as (1): Adding a 3. Default
|
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.
|
Filed #28970 to track the SessionBufferPool unit-test follow-up (point 6 from your review). |
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.
Verdict: ApproveThe Copilot review round addressed the three issues cleanly and the design is now in a shippable state. ConcurrentRunSupported=false is verified locally and the thread-safety reasoning in What changed since the prior review1. Strict option parsing in
|
Summary
SessionBufferPoolthat lets a session hold on to retired generator buffer caches (storage + uniform) and seed them into newly created generators.ep.webgpuexecutionprovider.sessionBufferPoolGenerationsto bound how many generations of retired buffers are kept (default1; set to0to disable).BufferManager's cache into the pool and absorb pooled buffers when a newBufferManageris created for the next generator.Motivation
With graph capture enabled, each generator owns its own per-graph
BufferManager. When the generator is destroyed (e.g., per-request in GenAI), the entire buffer cache is thrown away and the next generator must reallocate all storage and uniform buffers from scratch, increasing cold-start latency and GPU memory churn.By keeping a small pool of recently-retired buffer slots at the session level, the next generator can reuse them and skip reallocation entirely after the first cycle.
Test plan
--use_webgpu— clean build.lintrunner -areports no lint issues.verify_multi_gen.py: sequential and overlapping generators all produce matching, coherent output.verify_max_length_change.py: generators with varyingmax_lengthall coherent.storage hits=171 misses=0, uniform hits=296 misses=0, i.e., the pool actually engages and eliminates reallocation.Notes
SessionReleaseCapturedGraphfromState::~State()so the per-graphBufferManageris actually released and its buffers reach the pool.