Skip to content

webgpu: Add session-level buffer pool for graph capture reuse - #28761

Merged
qjia7 merged 4 commits into
microsoft:mainfrom
qjia7:webgpu-session-buffer-pool
Jun 10, 2026
Merged

webgpu: Add session-level buffer pool for graph capture reuse#28761
qjia7 merged 4 commits into
microsoft:mainfrom
qjia7:webgpu-session-buffer-pool

Conversation

@qjia7

@qjia7 qjia7 commented Jun 3, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Introduces SessionBufferPool that lets a session hold on to retired generator buffer caches (storage + uniform) and seed them into newly created generators.
  • Adds provider option ep.webgpuexecutionprovider.sessionBufferPoolGenerations to bound how many generations of retired buffers are kept (default 1; set to 0 to disable).
  • Wires the WebGPU EP to donate a retiring BufferManager's cache into the pool and absorb pooled buffers when a new BufferManager is created for the next generator.
  • The pool is only created when graph capture is enabled AND the option is > 0, so non-graph-capture sessions are unaffected.

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

  • Build ORT (Windows, D3D12) with --use_webgpu — clean build.
  • lintrunner -a reports no lint issues.
  • Verified end-to-end with GenAI on phi4 + WebGPU graph capture using two scripts:
    • verify_multi_gen.py: sequential and overlapping generators all produce matching, coherent output.
    • verify_max_length_change.py: generators with varying max_length all coherent.
  • With diagnostic prints (since removed), confirmed that after the first generator donates buffers, subsequent generators report storage hits=171 misses=0, uniform hits=296 misses=0, i.e., the pool actually engages and eliminates reallocation.

Notes

  • Pairs with a GenAI-side change that invokes SessionReleaseCapturedGraph from State::~State() so the per-graph BufferManager is actually released and its buffers reach the pool.

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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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::SessionBufferPool to retain and recycle storage/uniform buffer caches across captured-graph lifetimes.
  • Adds a new provider option ep.webgpuexecutionprovider.sessionBufferPoolGenerations and parses/logs it in the WebGPU EP factory/config.
  • Wires pooling into graph-capture lifecycle: seed buffers on per-graph BufferManager creation and donate buffers on ReleaseCapturedGraph.

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.

Comment thread onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc Outdated
Comment thread onnxruntime/core/providers/webgpu/buffer_manager.h Outdated
Comment thread onnxruntime/core/providers/webgpu/session_buffer_pool.cc
- 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.
@qjia7

qjia7 commented Jun 4, 2026

Copy link
Copy Markdown
Contributor Author

@hariharans29 @guschmue Please take a look, thanks.

@guschmue guschmue added the ep:WebGPU ort-web webgpu provider label Jun 4, 2026
@hariharans29

Copy link
Copy Markdown
Member

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 (storage hits=171 misses=0) demonstrate the pool engages as designed. The Copilot round already cleaned up strict option parsing and the cache accessors.

Two correctness questions worth answering before merge, plus some smaller observations.


Please confirm before merge

1. GPU work in-flight when donation happens

This is the question I most want answered. Looking at WebGpuContext::ReleaseGraphResources (onnxruntime/core/providers/webgpu/webgpu_context.cc#L922):

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 OnSubmittedWorkDone, device.Poll(Wait), or any equivalent fence. So when ReleaseCapturedGraph then calls session_buffer_pool_->Donate(...) a few lines later, the GPU may still be executing the last Replay's commands, which reference buffers in captured_buffers_ / buffers_ / pending_buffers_ of the per-graph caches.

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 SeedInto returns them to its cache, the next OnRunStart starts dispatching, and now we have two queue submissions reading/writing the same memory with no fence between them. WebGPU's data-race rules put the burden on the application here.

In practice this probably works today because:

  • State::~State() (the GenAI-side change) runs after the user has read output tokens, which typically forces a queue sync.
  • WebGPU drivers may serialize submissions from the same queue in ways that happen to be safe.

But neither is a guarantee. Two options:

a. Add an explicit device.Poll(Wait) / OnSubmittedWorkDone inside ReleaseGraphResources (or in SessionBufferPool::Donate before extraction). This is the bulletproof fix and has very little cost given how rarely ReleaseCapturedGraph runs.

b. Document the contract in SessionBufferPool::Donate as "caller must ensure all GPU work referencing this manager's buffers has completed before calling" and add a matching contract check at the GenAI call site.

The existing comment in GraphSimpleCacheManager::ExtractCachedBuffers ("Donation is expected after captured commands have been released and any in-flight work has completed") is the right intent but it's a comment, not a check or an enforced invariant. Option (a) is what I'd push for.

2. Thread-safety contract of the pool

SessionBufferPool::Donate, SeedInto, and Clear all touch slots_ with no locking. The contract for the pool is implicitly "single-threaded" — but the call paths are:

  • SeedInto — from OnRunStart (Run thread).
  • Donate — from ReleaseCapturedGraph, which is called from State::~State() per the PR description. The destructor runs on whichever thread holds the last unique_ptr<State>. In typical GenAI usage that's the Python/main thread that also calls Run, but it's not enforced.
  • Clear — from EP destructor.

WebGpuExecutionProvider::ConcurrentRunSupported() returns false, so OnRunStart is serialized within the EP. But State::~State() can race with OnRunStart if the GenAI consumer destroys one generator's State on a background thread while another generator's Run is starting.

Either add a std::mutex inside SessionBufferPool (cheap — these calls are rare), or document the single-thread invariant on SessionBufferPool itself and on ReleaseCapturedGraph in the EP header. Today there's no documentation either way.


Smaller observations

3. Default session_buffer_pool_generations = 1 is a behavior change

Existing graph-capture users — even those who don't know about this option — will start retaining 1 generation's worth of intermediate buffers across ReleaseCapturedGraph. For a large transformer that's not trivial GPU memory. The empirical phi4 numbers (storage hits=171) suggest this is order ~hundreds of buffers per generation.

I'd lean toward defaulting to 0 (opt-in) for safety, with the doc saying "set to 1 to enable buffer reuse across generators" — same way enable_graph_capture is opt-in. The fact that the pool only allocates when graph capture is on does narrow the blast radius, but it's still a memory-footprint regression for an existing feature.

Worth at least flagging in the PR description.

4. Shape-drift accumulation within a slot

The pool caps slot count (max_generations), but within an absorbed slot, buffers from sizes the new generator doesn't actually request will sit in the receiving manager's caches until that manager itself is donated or destroyed. With max_generations=1 this self-corrects each cycle (the next donation flushes the prior slot entirely). With higher generations, stale sizes can accumulate proportional to shape variation across generators.

Not a leak — Clear() and the eventual manager destruction release them — but worth a one-line comment on SeedInto so a future reader understands the bound is "1 slot's worth of donated buffer sizes per active slot", not "exactly N matched buffers".

5. Donate early-exit happens after Extract clears the cache

slot.storage = retiring_mgr.StorageCache().ExtractCachedBuffers();
slot.uniform = retiring_mgr.UniformCache().ExtractCachedBuffers();

if (slot.storage.empty() && slot.uniform.empty()) {
  return;
}

This is fine — the manager is about to be destroyed anyway, so leaving its caches empty has no observable effect. But a reader might wonder whether the early-exit should be before the extract. It's not — extract has to happen first to know if there's anything to donate. A one-line "// caches are now empty either way" comment would prevent the next reviewer asking the same question.

6. No unit tests for the pool

The phi4 end-to-end validation is real and valuable, but it's empirical and platform-specific. A unit test with a mock IBufferCacheManager exercising Donate → SeedInto roundtrip, Donate eviction at capacity, and Clear would lock in the contract and make the pool tractable to modify safely later. Not blocking, but a session-level resource pool with no dedicated test feels light for a maintainability standpoint.

7. BufferManager::StorageCache() / UniformCache() non-const accessors

These widen BufferManager's public surface for one specific consumer (SessionBufferPool). A friend declaration or a narrower API like void DonateTo(SessionBufferPool&) / void SeedFrom(SessionBufferPool&) on BufferManager would keep the cache-manager pointers encapsulated. Style nit — not blocking.

8. Provider option naming

sessionBufferPoolGenerations — fine, consistent with sibling enableGraphCapture, enableInt64, multiRotaryCacheConcatOffset. The unit "generations" is a bit jargon-y for an external option. Mentally I read it as "how many retired generators' worth of buffers to keep". Worth a sentence in the option description (which I see is in the comment above session_buffer_pool_generations in the header — that's reachable through code but not through whatever user-facing docs list provider options).


Things I checked and found OK

  • Destruction order: per_graph_buffer_mgrs_.clear() then session_buffer_pool_->Clear() then WebGpuContextFactory::ReleaseContext(context_id_). Pool buffers released before the device goes away. Good.
  • GraphCacheManager::AbsorbCachedBuffers routes through ReleaseBuffer, which inserts unknown sizes into a new bucket. So an extract-then-absorb roundtrip preserves layout regardless of bucket configuration.
  • GraphSimpleCacheManager::ExtractCachedBuffers correctly drains buffers_, pending_buffers_, and captured_buffers_, queries wgpuBufferGetSize for the latter two (consistent with how OnRefresh populates buffers_).
  • Default IBufferCacheManager::AbsorbCachedBuffers releases buffers via wgpuBufferRelease rather than leaking — right safe default for cache modes that can't absorb.
  • LIFO SeedInto and FIFO eviction — right pair for hot reuse with shape-drift adaptation. Comment captures the rationale well.
  • This PR is a no-op for users until the companion GenAI change lands (since no one currently calls ReleaseCapturedGraph from State::~State()). Worth linking the GenAI PR in the description so it's clear the two need to land together for the feature to work.

Bottom line

The design is sound and the empirical validation is encouraging. The two real questions are the GPU fence at donation time and the threading contract on the pool — answering those (preferably with code, alternatively with documentation) is what I'd want before approving the merge. Everything else is polish.

@qjia7

qjia7 commented Jun 10, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the thorough review @hariharans29! Going point by point.

1. GPU work in-flight at donation — no code change needed

The risk is narrower than it looks because of how WebGPU's queue model interacts with the existing session_mutex_ and OnRunEnd paths. Three invariants close it:

  1. session_mutex_ serializes Run and ReleaseCapturedGraph for non-concurrent EPs. WebGpuExecutionProvider::ConcurrentRunSupported() == false, so both InferenceSession::Run (inference_session.cc:3320-3323) and InferenceSession::ReleaseCapturedGraph (inference_session.cc:3898-3907) acquire session_mutex_. They cannot overlap on the host.

  2. OnRunEnd submits and (in validation mode) fences before returning. WebGpuExecutionProvider::OnRunEnd calls context_.Flush(BufferManager()) (webgpu_context.cc:734device_queue_.Submit). With ValidationMode::Basic or higher it then calls PopErrorScope(), which is Wait(device_.PopErrorScope(...)) (webgpu_context.cc:719-732, Wait = instance_.WaitAny(f, UINT64_MAX)). That is a CPU-side wait for the device callback, which only fires after all prior queue work has drained. So with validation enabled, the donation point is already fully fenced.

  3. WebGPU same-queue ordering protects the no-validation case too. Everything in WebGpuContext goes through the single device_queue_. WebGPU guarantees that submits on the same queue execute in submission order on the GPU timeline. So even without a host-side fence, Generator A's submitted-but-not-yet-completed Replay commands happen-before Generator B's subsequent first Submit — including any read or write of the donated buffers. The buffer handles are reused, but the GPU never observes overlapping access to them.

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 MapAsync/readback on donated buffers between donation and reuse. I'll add a one-line invariant comment to Donate capturing this so a future reader gets the reasoning without having to re-derive it.

2. Thread safety of the pool — no code change needed

For the same reason as (1): ConcurrentRunSupported() == false plus the session_mutex_ locking in both Run and ReleaseCapturedGraph means OnRunStart (which calls SeedInto) and ReleaseCapturedGraph (which calls Donate) are mutually exclusive on the same session. Clear() only runs from the EP destructor, which the session destroys after all Run/ReleaseCapturedGraph calls have returned. The pool is touched on the same thread the session serializes everything else on.

Adding a std::mutex would be cheap but redundant. I'll add a header comment on SessionBufferPool documenting the session_mutex_ invariant so this isn't an implicit contract.

3. Default session_buffer_pool_generations = 1no code change

This isn't a regression — it's a deferred release. The first generator's buffers are kept for the lifetime of the next generator instead of being freed at ~State(). As soon as the next generator is destroyed (in steady-state GenAI usage, that's per-request), its donation evicts the prior slot and the old buffers go back to the device. So peak memory is bounded by ~1 generation's worth of buffers extra, not a growing footprint, and the carve-out only applies when graph capture is already enabled (which itself is opt-in via enableGraphCapture). I'd rather keep the default that delivers the benefit out-of-the-box for the opt-in feature.

4. Shape-drift accumulation within a slot — will add comment

Fair, will add a one-liner to SeedInto explaining the bound is "≤1 donated slot's worth of stale sizes per active manager".

5. Donate early-exit after Extractwill add comment

Will add a // caches are now empty either way line above the early-exit so the next reader doesn't ask the same question.

6. No unit test for the pool — will defer to a follow-up

Agreed it would improve maintainability. I'd prefer to land this PR with the existing end-to-end coverage (verify_multi_gen.py / verify_max_length_change.py) and add a dedicated session_buffer_pool_test.cc in a follow-up once the GenAI-side State::~State() change lands and the feature is actually exercised in production. Happy to file an issue to track it.

7. BufferManager::StorageCache() / UniformCache() widening public surface — deferring

Agreed this is a style nit. Reshaping to DonateTo(SessionBufferPool&) / SeedFrom(SessionBufferPool&) would invert the dependency (BufferManager would need to know about SessionBufferPool) and is a larger restructure than the current accessor pair. Will leave as-is for this PR and reconsider if the surface grows.

8. Provider option naming/docs — will expand the header doc

Will expand the comment above session_buffer_pool_generations in webgpu_execution_provider.h to say "number of retired generators' worth of buffers to keep across ReleaseCapturedGraph". The option key itself stays sessionBufferPoolGenerations for consistency with the existing enableGraphCapture / enableInt64 naming.


Planned follow-up commit covers 1, 2, 4, 5, 8 (all doc-only — no code change beyond comments). Will push shortly and ping you again.

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.
@qjia7

qjia7 commented Jun 10, 2026

Copy link
Copy Markdown
Contributor Author

Filed #28970 to track the SessionBufferPool unit-test follow-up (point 6 from your review).

@qjia7
qjia7 requested a review from hariharans29 June 10, 2026 03:00
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.
@hariharans29

Copy link
Copy Markdown
Member

Verdict: Approve

The 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 SessionBufferPool::Donate holds together. A few small observations remain, all non-blocking.


What changed since the prior review

1. Strict option parsing in webgpu_provider_factory.cc — good, and actually stricter than the surrounding code

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.");
}

This is the gold-standard pattern. Worth noting that the existing multiRotaryCacheConcatOffset parser, a few lines above, only checks result.ec and would silently accept "12abc" as 12:

// pre-existing, NOT this PR
auto result = std::from_chars(..., offset_value);
if (result.ec == std::errc{}) {
  webgpu_ep_config.multi_rotary_cache_concat_offset = offset_value;
} else {
  ORT_THROW(...);
}

Same lax pattern in ParseWebGpuContextConfig for kDeviceId and kWebGpuInstance. Not your problem to fix here, but since you've now established the strict pattern in the same file, a follow-up PR to tighten the older ones would harmonize things. Worth filing as a small cleanup issue.

2. Non-const StorageCache() / UniformCache() accessors — necessary and correctly scoped

// Direct access to the underlying cache managers. Used by SessionBufferPool to
// donate/seed buffers across per-graph BufferManager lifetimes.
IBufferCacheManager& StorageCache() { return *storage_cache_; }
IBufferCacheManager& UniformCache() { return *uniform_cache_; }

The earlier round had these const which wouldn't compile with the mutating Extract/Absorb calls. The fix is the right shape — exposes only the two caches the pool actually needs (storage + uniform), leaving query_resolve_cache_ and default_cache_ private. Good encapsulation.

3. The Donate-safety comment block is now the standout doc in this PR

// 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.

Verified locally:

  • (1) WebGpuExecutionProvider::ConcurrentRunSupported() const override { return false; } — InferenceSession serializes everything via session_mutex_. ✓
  • (2) OnRunEnd calls context_.Flush(BufferManager()) and ends with context_.PopErrorScope() in Basic+ validation. ✓
  • (3) Single queue ordering is a WebGPU spec guarantee for submits to the same queue. ✓

The "what would break invariants" note (different queue, host readback) is exactly the right risk callout for the next person who tries to extend this.


Re-confirmation of the core design

Donate timing in ReleaseCapturedGraph is safe

if (cmd_it != captured_graphs_.end()) {
  if (!cmd_it->second.empty()) {
    context_.ReleaseGraphResources(cmd_it->second);   // releases bind_group/layout only
  }
  captured_graphs_.erase(cmd_it);
}
...
if (session_buffer_pool_ && mgr_it->second) {
  session_buffer_pool_->Donate(*mgr_it->second);
}

Looked at ReleaseGraphResources — it only releases wgpuBindGroupRelease / wgpuBindGroupLayoutRelease, not buffer refs. The WGPUBuffer handles being donated are still alive (the BufferManager's caches hold the canonical refs), and the queue internally retains refs on submitted buffers until GPU completion regardless of host-side bind group release. The donation can't free a buffer the GPU still needs.

Eviction policy

Donate evicts FIFO (oldest), SeedInto pops LIFO (newest). At default max_generations_ = 1 this is moot, but at capacity > 1 the policy is "keep the freshest N, hand out the newest first." Consistent and defensible for the "buffer shapes drift between generators" motivation — older slots are increasingly stale.

GraphSimpleCacheManager::ExtractCachedBuffers covers pending_buffers_ and captured_buffers_ too

Worth calling out as a positive — a naive implementation would only drain buffers_ and silently leak whatever was in pending_buffers_ / captured_buffers_ at donation time. The comment in the implementation captures the safety reasoning ("all three containers therefore hold buffers no longer referenced by the device").

The reasoning depends on donation happening only after OnRunEnd (which flushes), which is true for the call site. If someone ever calls ExtractCachedBuffers mid-run, this would be unsafe — a one-line ORT_ENFORCE or comment in the header virtual saying "only safe to call when no Run is in progress" would harden the interface. Not blocking.

Default AbsorbCachedBuffers releases — correct fallback

The base class default in buffer_manager.h:

virtual void AbsorbCachedBuffers(std::vector<std::pair<size_t, WGPUBuffer>>&& buffers) {
  for (auto& entry : buffers) {
    if (entry.second) {
      wgpuBufferRelease(entry.second);
    }
  }
}

Right semantics — "if a future cache mode can't store these, don't leak." Currently only Graph and GraphSimple actually override, but those are the only ones used with graph capture, so the default is purely defensive against future additions. Fine.


Remaining observations (non-blocking)

a) slots_ is a std::vector with erase(begin()) for FIFO

while (slots_.size() >= max_generations_) {
  auto& victim = slots_.front();
  ReleaseSlotBuffers(victim.storage);
  ReleaseSlotBuffers(victim.uniform);
  slots_.erase(slots_.begin());
}

O(N) on slots_.size() per eviction. At default capacity 1 this is irrelevant; even at expected real-world values (≤ a handful) it's negligible. If anyone ever sets this to a large value, std::deque<Slot> with pop_front() would be the more honest container. Nit-only.

b) Donate's max_generations_ == 0 early-return is correct but slightly under-documented

if (max_generations_ == 0) {
  return;
}

When max_generations_ == 0, the SessionBufferPool isn't created at all in the EP ctor:

if (enable_graph_capture_ && config.session_buffer_pool_generations > 0) {
  session_buffer_pool_ = std::make_unique<webgpu::SessionBufferPool>(...);
}

So the early-return is dead code in practice (the unique_ptr would be null and the call site is guarded by if (session_buffer_pool_)). Fine as defense-in-depth, but it could go away to make the invariant "if a pool exists, capacity > 0" more obvious. Optional.

c) Empty-slot suppression in Donate

slot.storage = retiring_mgr.StorageCache().ExtractCachedBuffers();
slot.uniform = retiring_mgr.UniformCache().ExtractCachedBuffers();
if (slot.storage.empty() && slot.uniform.empty()) {
  return;
}

Good — avoids pool slots filling with empty entries when a generator did almost nothing. Worth noting that this means an unusual case (capacity=2, one full + one empty donation in sequence) won't evict the older full one, which is the right behavior here. Quietly correct.

d) Header thread-safety note is precise and load-bearing

The header comment on SessionBufferPool:

// 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, ...

This is exactly the right level of detail for a class that has no internal locking and isn't general-purpose. If someone ever reuses SessionBufferPool from a different call site, they'll know to either keep the same serialization invariant or wrap it. Good.

e) GenAI dependency

Pairs with a GenAI-side change that invokes SessionReleaseCapturedGraph from State::~State() so the per-graph BufferManager is actually released and its buffers reach the pool.

Worth checking: what happens with current GenAI (without the paired change) once this lands? Answer from reading the code: nothing — ReleaseCapturedGraph is the only donation path, so without the GenAI change the pool stays empty and behavior is unchanged from before. So this lands safely ahead of the GenAI-side change. A one-line comment in the PR description confirming "this PR is safe to land independently; pool stays empty without the paired GenAI change" would save a maintainer cycle.


Bottom line

Approve. The three Copilot review points were addressed correctly, the thread-safety reasoning checks out against the actual code (ConcurrentRunSupported() == false, OnRunEnd flushes, single queue ordering), and the design has been further hardened (strict option parsing, pending_buffers_ / captured_buffers_ now also drained on extraction). The only suggested follow-up worth mentioning explicitly is tightening the older from_chars call sites in the same file to match the new strict pattern.

@qjia7
qjia7 merged commit 50c9510 into microsoft:main Jun 10, 2026
91 of 92 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ep:WebGPU ort-web webgpu provider

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants