Skip to content

[WebGPU] Deferred-dispatch to parallelize cold-start shader compilation - #29557

Merged
Hariharan Seshadri (hariharans29) merged 21 commits into
mainfrom
feature/webgpu-defer-dispatch
Aug 10, 2026
Merged

[WebGPU] Deferred-dispatch to parallelize cold-start shader compilation#29557
Hariharan Seshadri (hariharans29) merged 21 commits into
mainfrom
feature/webgpu-defer-dispatch

Conversation

@xiaofeihan1

@xiaofeihan1 Xiaofei Han (xiaofeihan1) commented Jul 6, 2026

Copy link
Copy Markdown
Contributor

Summary

Parallelize WebGPU cold-start pipeline compilation by starting CreateComputePipelineAsync without immediately waiting, recording subsequent dispatches, and then waiting, encoding, and submitting them together in bounded windows.

How it works

Previously, each cache-miss ProgramManager::Build() effectively serialized pipeline creation by calling Wait(CreateComputePipelineAsync(...)) before the operator could continue. A model with many first-run shaders therefore paid the compilation latency one pipeline at a time.

With this change:

  1. On a program-cache miss, ProgramManager::Build() starts CreateComputePipelineAsync and returns its future and callback state without waiting. It also builds an explicit wgpu::BindGroupLayout synchronously (from the input/output/uniform binding info) so the caller can create the bind group before the pipeline finishes compiling.
  2. WebGpuContext::Run() records the dispatch state needed later: the cache key, the already-created wgpu::BindGroup (which internally retains its bound buffers), dispatch dimensions, indirect buffer, and optional profiling metadata.
  3. Repeated cache misses for the same key within the active window share one in-flight pipeline build instead of compiling the same shader again.
  4. Once the number of recorded dispatches reaches maxNumPendingDispatches (16 by default), WaitForDeferredPipelineBuildsAndEncodeDispatches():
    • waits for all unique pending pipeline builds in the window;
    • inserts completed ProgramArtifacts into the program cache; and
    • calls LaunchComputePipeline() for every recorded dispatch in original order, including pipeline/bind-group setup and dispatch encoding.
  5. Flush() then submits the encoded work to the GPU. OnRunEnd() drains and submits the final partial window, and GPU-to-CPU download drains deferred work before recording its readback copy.

This overlaps independent pipeline compilations instead of waiting for each pipeline serially, reducing first-run and prefill latency.

Added complexity and correctness considerations

1. CPU profiling attribution changes

Previously, an operator's CPU duration included the synchronous shader/pipeline compilation performed by that operator. With deferred dispatch, most operators return quickly after starting or reusing an asynchronous build and recording their dispatch.

The operator that reaches the 16-dispatch window boundary performs the drain. Its CPU duration can therefore be much longer because it waits for outstanding pipeline builds and encodes the preceding dispatches. The same work may be charged to a run boundary for the final partial window. As a result, CPU profiling now shows many short operators and periodic long drain points rather than distributing compilation time across individual operators. GPU profiling metadata is captured at record time and replayed when each dispatch is encoded so dispatch accounting remains ordered.

2. Deferred buffer lifetime (cache the bind group, not raw buffers)

Previously, each bind group was created and consumed while its operator was executing, so its buffer references only needed to outlive that single dispatch.

After deferring dispatch encoding, the recorded dispatch must keep its buffers valid until the window is drained and encoded — even though ORT may recycle an intermediate tensor as soon as its last-use kernel returns. This is particularly important with storageBufferCacheMode=disabled, where GpuBufferAllocator::Free() reaches BufferManager::Release() and immediately calls wgpuBufferRelease().

Rather than tracking raw WGPUBuffer handles, each DeferredDispatch caches the already-created wgpu::BindGroup. A WebGPU bind group internally holds a reference to every buffer it binds, so caching the bind group is the buffer retention — no manual wgpuBufferReference/Release, and no dependency on the selected buffer-cache mode. The bind group (and its buffer references) is released via RAII when the dispatch is encoded or the window is cleared, including failure paths. Creating the bind group at record time requires the layout up front, which is why ProgramManager::Build() produces an explicit wgpu::BindGroupLayout synchronously (decoupled from the async pipeline compile).

3. Asynchronous callback lifetime

The async pipeline callback writes the compiled pipeline into a heap-allocated PipelineCallbackContext that owns the wgpu::ComputePipeline by value. Only that context needs a stable address — Dawn holds a raw pointer to it as callback state until the future completes — so it is the sole heap allocation required for a pending build. The remaining in-flight build state (PendingPipelineBuild) is stored inline in the dispatch window and may move freely as the window grows. The callback is noexcept and reports failures through Status rather than letting exceptions cross the Dawn callback boundary.

Performance

Phi-4 Mini prefill latency on WebGPU/D3D12:

Scenario Serial pipeline waits Deferred dispatch Improvement
Warm driver shader cache (NVIDIA) ~540 ms ~380 ms ~29.6%
Cold, no driver cache ~210 ms ~120 ms ~42.9%

Validation

  • Built onnxruntime_providers_webgpu, the WebGPU shared library, and the onnxruntime_webgpu wheel in RelWithDebInfo.
  • Verified Phi-4 Mini easy and prefill-500 execution with the default 16-dispatch window.
  • Verified storageBufferCacheMode=disabled with window sizes 16 and 4096.
  • Verified GPU-to-CPU readback ordering by draining deferred compute before encoding the copy.
  • Verified asynchronous build failure cleanup, profiling bookkeeping, indirect dispatch, and graph-capture resource paths.

@github-actions github-actions Bot 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.

You can commit the suggested changes from lintrunner.

Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc Outdated
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.h Outdated

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 WebGPU EP “deferred-dispatch” mode for the first non-graph-captured prefill run to reduce cold-start latency by overlapping shader pipeline compilation (CreateComputePipelineAsync) with dispatch recording, then flushing recorded work in windows.

Changes:

  • Adds deferred-dispatch gating to WebGpuExecutionProvider::OnRunStart/OnRunEnd, including run-option/env toggles and routing through a reserved graph-mode buffer manager annotation (-2).
  • Extends ProgramManager::Build to support an asynchronous mode returning a wgpu::Future plus a callback context, enabling later wait/commit.
  • Implements deferred dispatch recording and windowed flush logic in WebGpuContext (record dispatches on cache misses, then wait/encode/submit per window).

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
onnxruntime/core/providers/webgpu/webgpu_execution_provider.h Adds EP-level flags tracking deferred-dispatch pending/active state.
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc Enables deferred-dispatch on eligible first prefill run and drains it at run end; adds run-option/env gating.
onnxruntime/core/providers/webgpu/webgpu_context.h Introduces deferred-dispatch data structures (pending builds + recorded dispatches) and flush APIs.
onnxruntime/core/providers/webgpu/webgpu_context.cc Implements deferred-dispatch recording in Run() and windowed wait/encode/submit in FlushDeferred*.
onnxruntime/core/providers/webgpu/program_manager.h Adds async Build-mode API and callback context type for CreateComputePipelineAsync.
onnxruntime/core/providers/webgpu/program_manager.cc Implements async Build-mode (returning future + owned callback context) and a new ProgramArtifact ctor used at flush time.

Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc Outdated
Comment thread onnxruntime/core/providers/webgpu/program_manager.cc Outdated
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc Outdated
The first prefill run of a WebGPU session serially compiles one shader
compute pipeline per cache-miss program (WGSL generation + shader module +
CreateComputePipelineAsync + Wait). On a cold start this serial compilation
dominates the first-token latency.

This change adds a "deferred-dispatch" path for the first non-graph-captured
prefill run. Instead of synchronously compiling and launching each program,
Run() issues asynchronous pipeline compilation (CreateComputePipelineAsync
without waiting) and records the dispatch. Once a window of dispatches
(maxNumPendingDispatches) has been recorded, that window's pipelines -- which
Dawn's worker pool has been compiling concurrently -- are waited on and then
encoded + submitted. Windowing preserves buffer recycling and CPU/GPU overlap
while several shader compilations proceed in parallel.

The deferred run is routed through a graph-mode buffer manager (Graph /
GraphSimple) under a reserved annotation id (-2) so the recorded WGPUBuffer
handles stay valid until the window is flushed -- the same stability guarantee
graph capture relies on.

Enabled by default, but only under a graph-capture-enabled session: there,
prefill uses annotation -1 and GroupQueryAttention keeps total_sequence_length
on the GPU, which avoids a CPU read-back that would otherwise force an early
flush and break deferral. It can be disabled per run via the run-option
"ep.webgpu.defer_dispatch"="0" or the env var ORT_WEBGPU_DEFER_DISPATCH=0.
…ompiles, review fixes

- Enable deferred-dispatch by default for the first prefill run of any session
  (no longer require graph capture); still skips runs that will be graph-captured
  and honors ep.webgpu.defer_dispatch=0 / ORT_WEBGPU_DEFER_DISPATCH=0.
- Dedup repeated cache keys within a window so a program compiles only once
  (deferred_inflight_builds_); duplicates resolve the pipeline from cache at flush.
- FlushDeferredWindow: on compile failure, release uniform buffers and drop the
  window instead of leaking recorded buffers/pointers.
- Use std::make_unique for the async pipeline callback context.
- Trim duplicated DEFER-DISPATCH banner comments; keep one canonical description.
A mid-run GPU->CPU readback (BufferManager::Download, used by MemcpyToHost) only
flushed the command encoder, not the deferred-dispatch queue. When a value the
readback depends on was produced by a still-deferred compute (e.g. GroupQueryAttention's
CPU-side total_sequence_length input when graph capture is off), the readback observed
stale data -- silently wrong for short prompts and a hard failure
(total_sequence_length must be positive, got 0) for longer prefills.

Flush any pending deferred dispatches before the readback so it observes correct
results. Deferred-dispatch stays enabled, so later ops resume batching; in practice
only the few ops before the readback are flushed early, preserving most of the
cold-start speedup.
Drop the ep.webgpu.defer_dispatch run-option, ORT_WEBGPU_DEFER_DISPATCH env var, and the will_be_captured annotation parsing. Defer now activates on the first run of any non-graph-capture session and is skipped entirely under graph capture (mutually exclusive with record/replay). Removes the now-unused env_var.h include.
@xiaofeihan1
Xiaofei Han (xiaofeihan1) force-pushed the feature/webgpu-defer-dispatch branch from 48e08a1 to 88b7180 Compare July 15, 2026 07:37

@xiaofeihan1 Xiaofei Han (xiaofeihan1) left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

add comments

Comment thread onnxruntime/core/providers/webgpu/webgpu_context.h Outdated
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc

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

Copilot reviewed 7 out of 7 changed files in this pull request and generated 5 comments.

Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc Outdated
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc Outdated
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.h
Comment thread onnxruntime/core/providers/webgpu/program_manager.cc
Retain deferred dispatch buffers until encoding, clarify in-window pipeline build sharing, remove an unused include, and clean up ProgramArtifact construction.
@xiaofeihan1
Xiaofei Han (xiaofeihan1) marked this pull request as ready for review July 17, 2026 05:48
@qjia7

Copy link
Copy Markdown
Contributor

Thanks for tackling this — the cold-start serialization of CreateComputePipelineAsync is a real bottleneck and the measured prefill wins (~30–43%) are compelling. I agree with the core direction: defer the pipeline-build completion so compilations overlap within a batch, and let the GPU start working per-window instead of waiting until the very end.

My concern is that the current implementation is heavier than it needs to be. Deferring the entire dispatch encoding (not just the compilation) is what pulls in:

  • manual wgpuBufferAddRef / wgpuBufferRelease retention on every bind/indirect buffer (DeferredDispatch::RetainBuffers/ReleaseBuffers + the custom move ctor),
  • correctness reasoning tied to buffer-cache mode (storageBufferCacheMode=disabled),
  • new drain points in BufferManager::Download, WebGpuKernel::PrePack, and OnRunEnd,
  • a second encode loop that largely duplicates the existing graph-capture Replay() path.

Proposal: cache the bind group, not the raw buffers

Instead of recording raw WGPUBuffer handles and re-encoding later, I'd like us to record, per dispatch:

  1. the CreateComputePipelineAsync future/promise,
  2. the already-created WGPUBindGroup,
  3. the dispatch workgroup size (or the indirect buffer), and
  4. profiling info.

When the batch reaches the max window: wait for all promises → then SetPipeline / SetBindGroup / dispatch one-by-one → submit. Continue with the next batch the same way.

The key advantage: a WGPU bind group internally holds references to the buffers it binds, so caching the bind group is the buffer retention. That means:

  • no manual AddRef/Release and no DeferredDispatch buffer bookkeeping,
  • no dependency on cache mode — this is correct because uniform buffers (easy mode) aren't reused/destroyed within a batch, and storage-buffer bucket reuse doesn't affect correctness or get destroyed within a batch, and WebGPU enforces the lifetime regardless,
  • no changes to Download / PrePack / OnRunEnd buffer logic, since buffer management is untouched.

The one prerequisite: explicit bind group layout

The blocker today is line ~917:

WGPUBindGroupLayout bind_group_layout =
    program_artifact.compute_pipeline.GetBindGroupLayout(0).MoveToCHandle();

We currently derive the layout from the compiled pipeline, so the bind group can't be built until the promise resolves. To create+cache the bind group at record time, we need to switch to an explicit WGPUBindGroupLayout constructed from the binding info we already have (inputs + outputs + optional uniform — already encoded in bind_buffers_segments). Build() would create that layout synchronously, pass it to CreateComputePipelineAsync, and hand it back so Run() can create the bind group immediately. This is a contained change inside Build() / LaunchComputePipeline and doesn't touch buffer management. The only correctness responsibility is that the explicit layout matches the shader's binding declarations (binding indices, visibility, storage vs. uniform, read-only vs. read-write) — all of which ShaderHelper/the segment logic already knows.

Bonus: this unifies with the graph-capture path

CapturedCommandInfo already stores essentially this exact record (pipeline + cached bind group + dispatch size / indirect buffer + profiling), and Replay() already does the "set pipeline, set bind group, dispatch one-by-one, submit in windows" loop. So this approach can reuse CapturedCommandInfo and the Replay() encode/cleanup machinery (including ReleaseGraphResources for the cached bind groups) instead of introducing a parallel DeferredDispatch mechanism with different lifetime rules.

One thing this does not change

Profiling attribution still shifts — encoding happens at the batch boundary, so the operator that triggers the drain absorbs the wait+encode CPU time. That's inherent to deferral, not specific to either implementation, so let's just document it either way.

Net: I think we can get the same parallel-compilation win with fewer lines, no buffer-lifetime special-casing, no cache-mode assumptions, and no changes to Download/PrePack/OnRunEnd — by caching the bind group behind an explicit layout and reusing the existing capture/replay path. Would you be open to reworking it in this direction?

Create explicit bind group and pipeline layouts so deferred dispatches can retain ready bind groups instead of raw buffers. Share command encoding with graph replay and use RAII for bind group lifetime.

@github-actions github-actions Bot 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.

You can commit the suggested changes from lintrunner.

Comment thread onnxruntime/core/providers/webgpu/program_manager.h Outdated
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc Outdated
PipelineCallbackContext now owns the compiled pipeline by value, so only that heap-allocated context (not the whole build) needs a stable address. PendingPipelineBuild is stored inline via std::optional (drops unique_ptr<PendingPipelineBuild>) and no longer carries a redundant compute_pipeline field; the pipeline is delivered through callback_context->pipeline. Context allocation moves from ProgramManager::Build to the caller, so Build takes a PipelineCallbackContext& and just wires it to Dawn.

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

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc Outdated
Comment thread onnxruntime/core/providers/webgpu/program_manager.cc Outdated
- Allow the bind-group binding count to equal maxBindingsPerBindGroup (< -> <=) in CreateBindGroup and CreatePipelineLayout; the count reaching the limit is valid.
- Fix CreatePipelineLayout declaration indentation.
@xiaofeihan1

Copy link
Copy Markdown
Contributor Author

Thanks for the detailed proposal — a few updates on where things landed:

  • Bind-group caching behind an explicit layout is already in the current revision. ProgramManager::Build() now creates the wgpu::BindGroupLayout synchronously from the segment info and hands it back, so Run() builds and caches the wgpu::BindGroup at record time. There is no more raw-WGPUBuffer retention, no manual AddRef/Release, and no storageBufferCacheMode special-casing — the bind group holds its buffer references.
  • I also simplified the in-flight build ownership: PipelineCallbackContext now owns the compiled pipeline by value, so only that small heap object needs a stable address, and PendingPipelineBuild is stored inline.

On reusing CapturedCommandInfo + Replay() and dropping DeferredDispatch: I kept them as separate types on purpose. At record time a cache-miss dispatch has no ready pipeline handle — it is produced later by the async CreateComputePipelineAsync callback — so the record has to carry the cache key plus the in-flight build, which CapturedCommandInfo (a copyable, replay-ready value that already holds a resolved pipeline) does not model. Folding the async build state into it would also make it move-only and pin build-only fields onto every long-lived replay command.

That said, you're right that the encode / timestamp / window loop is duplicated between the deferred drain and Replay(). I'm happy to extract that shared per-dispatch loop into a common helper used by both paths, keeping the two record types distinct. Would that address the duplication concern?

@qjia7 Jiajia Qin (qjia7) 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.

Two suggestions on the deferred-dispatch structure:

  1. Move WaitForDeferredPipelineBuildsAndEncodeDispatches into Flush.

It feels more natural for Flush to own the "drain everything queued" contract. Requiring callers to remember to wait on deferred builds before flushing is a footgun, and any future flush path that forgets the
wait would silently drop work. Making it Flush's responsibility gives us a clean invariant: after Flush returns, no deferred dispatches remain.

  1. Merge DeferredDispatch into CapturedCommandInfo.

The two structs describe the same dispatch at different lifecycle stages (pipeline-pending vs. pipeline-ready). I'd suggest:

  • Delete DeferredDispatch.
  • In CapturedCommandInfo, change compute_pipeline to std::optional<...> and add std::optional pending_build plus the program key.
  • deferred_dispatches_ stores CapturedCommandInfo directly.

At flush time, for each entry check program_mgr_->Get(key).pipeline:

  • If already available (another dispatch of the same program resolved it), populate compute_pipeline from the cache and drop pending_build.
  • Otherwise wait on pending_build, then write the pipeline back into both program_mgr_->Get(key) and compute_pipeline.

Benefits: single storage/replay path for both capture mode and deferred-dispatch mode, and duplicate dispatches of the same program during the wait window naturally share one pipeline via the cache instead of
racing on separate PendingPipelineBuilds.

Comment thread onnxruntime/core/providers/webgpu/webgpu_context.h
Merge deferred dispatch storage into CapturedCommandInfo, resolve pending pipelines during flush, and share command dispatch logic with graph replay.
@xiaofeihan1

Copy link
Copy Markdown
Contributor Author

Jiajia Qin (@qjia7) Thanks for your suggestion!

Move WaitForDeferredPipelineBuildsAndEncodeDispatches into Flush.

Done. I moved it into Flush method and rename it to EncodeDeferredDispatches. But for Download of buffer manager, we also keep first EncodeDeferredDispatches and then flush to ensure the record buffer of command encoder is by order.

Merge DeferredDispatch into CapturedCommandInfo.
Done

@hariharans29

Copy link
Copy Markdown
Member

Review: PR #29557 — [WebGPU] Deferred-dispatch to parallelize cold-start shader compilation (head a0b5436)

Author: @xiaofeihan1 (Intel WebGPU EP). Branch feature/webgpu-defer-dispatchmain. 19 commits, +404 / −183 across 8 files. CI: 86 / 87 checks OK. Reviewers: @qjia7 (Intel, actively engaged — substantive design feedback shaped the current shape), @hariharans29 (requested). 3 participants + Copilot AI (multiple review rounds). 42 conversation entries.

Verdict: approve on the current head, subject to one small lint fix. Substantive perf win (Phi-4 Mini prefill: −29.6% warm driver cache, −42.9% cold), correctness argument holds, and the design has iterated cleanly through @qjia7's review — the current revision is materially simpler than the initial submission: bind-group RAII instead of manual AddRef/Release, unified CapturedCommandInfo state machine across deferred + graph-capture paths, drain-owned-by-Flush contract, and in-window pipeline-build dedup. The complexity that remains is intrinsic to the async-compile-then-encode design, not incidental.


What this PR does

Turns CreateComputePipelineAsync from a serialization bottleneck into a parallelism opportunity by:

  1. Non-blocking ProgramManager::Build() — starts the async pipeline compile and returns immediately with wgpu::Future + PipelineCallbackContext&, no Wait(). Also builds an explicit wgpu::BindGroupLayout synchronously from the input/output/uniform segment info (via new CreatePipelineLayout()) so the caller can build the bind group without waiting for the pipeline.
  2. Record-time bind group creation in WebGpuContext::Run() — the already-created wgpu::BindGroup is stored in a CapturedCommandInfo entry, alongside dispatch dims, indirect buffer, and optional profiling metadata. The bind group internally retains its bound buffers, so buffer lifetime is preserved without any manual refcount work.
  3. In-window build dedup — repeated cache misses for the same key within one window share one in-flight PipelineCallbackContext; only the first record owns pending_build, later records just carry program_key and resolve via the program cache after the wait.
  4. Window-boundary drain — when deferred_dispatches_ reaches max_num_pending_dispatches_ (16), WaitForDeferredPipelineBuilds() waits on all unique pending builds, populates the program cache, promotes every recorded command to state 4 (pipeline resolved), then DispatchCommand() encodes them in original order. Flush() then submits.
  5. Sync-point orderingUpload, MemCpy, Download, FillZero, and PrePack each call EncodeDeferredDispatches() before recording their copies/fills so GPU-side ordering is preserved.

Two new types:

  • PipelineCallbackContext — heap-allocated (only object that needs a stable address; Dawn holds a raw pointer to it until the future completes). Owns wgpu::ComputePipeline by value.
  • PendingPipelineBuild — inline in CapturedCommandInfo (moves freely as the window grows). Owns PipelineCallbackContext via unique_ptr, plus the future, the layout, the program name, and the shape-uniform ranks.

Design evolution — significant improvement over the initial submission

@qjia7's pass-1 review flagged that the initial approach:

  • did manual wgpuBufferAddRef/wgpuBufferRelease on every bind/indirect buffer,
  • had correctness reasoning tied to storageBufferCacheMode=disabled,
  • added new drain points in Download, PrePack, OnRunEnd,
  • and duplicated the graph-capture Replay() encode/window loop.

The current revision addresses each point:

Concern Resolution
Manual buffer AddRef/Release Gone. wgpu::BindGroup retains its bound buffers implicitly (da769f6).
storageBufferCacheMode dependency Gone. Bind-group retention is buffer-cache-mode-independent.
Duplicated encode/window loop CapturedCommandInfo extended with program_key + pending_build? + compute_pipeline? so one struct represents both deferred and captured dispatches (cb0943f). One encode loop (DispatchCommand) serves both paths.
Wait responsibility on callers WaitForDeferredPipelineBuilds() folded into Flush(); renamed pre-flush entry point to EncodeDeferredDispatches(). Flush now owns the "no deferred work remains after return" invariant.
Bind-group layout coupling Explicit wgpu::BindGroupLayout built synchronously from segment info in new ProgramManager::CreatePipelineLayout().
Flush failure handling in graph capture New OnRunEnd cleanup path: on Flush failure, calls CaptureEnd, releases graph resources, clears captured commands, drops graph_buffer_mgr_active_, pops error scope, returns status. Prevents dangling captured-command entries.

That's a lot of clean-up done in review. The current diff feels appropriately sized for the value.


Correctness spot-checks

1. Lifetime of the async-build callback context

struct PipelineCallbackContext {
  wgpu::ComputePipeline pipeline;
  Status status;
};
struct PendingPipelineBuild {
  std::unique_ptr<PipelineCallbackContext> callback_context;  // heap; stable address
  wgpu::Future future;
  // ...
};

Dawn holds a raw pointer to the PipelineCallbackContext heap object as callback state until the future completes. PendingPipelineBuild owns via unique_ptr, so moving the enclosing CapturedCommandInfo in deferred_dispatches_.push_back(...) reallocation does not invalidate Dawn's pointer. ✓ Callback is noexcept and reports via status, so exceptions do not cross the Dawn boundary. ✓

2. The 5-state CapturedCommandInfo state machine

The header comment enumerates:

  1. Program-cache hit → pending_build = empty, compute_pipeline = set.
  2. First cache-miss for a key → pending_build = set, compute_pipeline = empty.
  3. Later cache-miss same key → pending_build = empty, compute_pipeline = empty, program_key set for cache lookup after wait.
  4. Post-WaitForDeferredPipelineBuildspending_build = empty, compute_pipeline = set.
  5. Captured / replay → pending_build = empty, compute_pipeline = set. Never depends on pending_build.

DispatchCommand() requires compute_pipeline to be set. State 3 (both empty) is transient and resolves to state 4 before encoding. Explicit invariant, well-documented, and enforced at the encode site. ✓

3. Buffer/download drain ordering

// Download:
ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches());
EnforceBufferUnmapped(context_, src);
// ...
auto& command_encoder = context_.GetCommandEncoder();
context_.EndComputePass();
command_encoder.CopyBufferToBuffer(src, 0, staging_buffer, 0, buffer_size);
ORT_THROW_IF_ERROR(context_.Flush(*this));

Deferred dispatches are encoded into the current compute pass first, then the pass is ended, then the copy is recorded. Because a single command encoder preserves record order in the eventual submit, GPU-side ordering is guaranteed. Applied consistently to Upload, MemCpy, Download, FillZero. ✓

4. Explicit layout matches shader binding declarations

CreatePipelineLayout() derives the layout from the same inputs_segments, outputs_segments, and shape_uniform_ranks that ShaderHelper::GenerateSourceCode() uses to emit the shader's @group(0) @binding(N) declarations, plus the UniformVariables().length > 0 check to decide whether to add the uniform binding. Both the layout and the shader come from the same source of truth, so declarative binding-order drift is by construction impossible. ✓ Also asserts binding <= maxBindingsPerBindGroup to catch device-limit overflow.

5. In-window dedup semantics

When a second cache-miss for the same key lands in the window, FindPendingPipelineBuild(key) returns the earlier command's build, and the current command sets program_key only (states 2 → 3 → 4). Only one CreateComputePipelineAsync per unique key per window. Saves both Dawn round-trips and driver-side compile time. ✓

6. Flush() return-type change

Flush() now returns Status. All internal callers updated:

Public API breakage bounded to within the WebGPU EP. External consumers should be unaffected. ✓

7. PrePack drain

The comment explains the invariant: "PrePack has no OnRunEnd hook, so finish its deferred pipeline builds and submit any encoded GPU work before returning and allowing ORT to release the original initializer tensor." Correct — the initializer's tensor buffer is the source of the copy, and ORT may recycle it as soon as PrePack returns.


Nits (all non-blocking)

  1. Lint fix: reviewdog flagged #include <vector> missing in program_manager.h (line 71, std::vector<int>& parameter). One-line addition. Should land before merge.
  2. Hardcoded window size max_num_pending_dispatches_ = 16. Author validated 16 and 4096. Consider exposing via WebGpuContextConfig or env override for tuning across model shapes / driver contexts. Non-blocking.
  3. Extensibility of EncodeDeferredDispatches injection points: today five sync-points call it (Upload, MemCpy, Download, FillZero, PrePack). A future sync-point that forgets to call it would silently miss GPU ordering. Consider either (a) having EndComputePass itself invoke EncodeDeferredDispatches when the window is non-empty, or (b) a comment + assertion in EndComputePass that flags "unencoded deferred work + about-to-record-copy" as an error in debug builds.
  4. State-machine assertion: DispatchCommand() requires compute_pipeline to be set. Add a debug_assert(command.compute_pipeline.has_value()) or ORT_ENFORCE at the top of DispatchCommand if not already present — cheap defense against future refactors accidentally emitting a state-3 command past the wait.
  5. Direct unit tests: 19 commits and 42 conversation entries, but no new unit tests. The 5-state lifecycle, in-window dedup ("multiple misses for same key produce one build"), and Download/PrePack drain ordering are all testable with a scripted WebGpuContext fixture. Not blocking (perf integration testing covers happy-path), but a regression suite would harden future refactors.
  6. Perf numbers: Phi-4 Mini on D3D12 only. Cross-backend numbers (Metal, Vulkan) would strengthen the story — Vulkan/Metal drivers with poorer parallel-compile scheduling may see different (still positive but smaller) wins. Nice-to-have for the PR description.
  7. 19-commit history: messy from iterative review, but auto-squash on merge collapses it. Author notes are consistent with what's in the diff.

Merge state

  • CI: 86 / 87 checks OK on a0b5436. One outstanding. Clean.
  • Approvals: none yet. @qjia7 has driven substantive review through multiple rounds and appears satisfied on the latest commit. @hariharans29 requested.
  • Copilot AI: five review rounds, most recent generated no new comments.
  • Labels: none. ep:WebGPU would help retention.

Bottom line

Real perf win (large, especially for cold-start), correctness argument holds up under the 5-state lifecycle and the sync-point drain audit, and the design has been meaningfully simplified through iterative review. @qjia7's pass-1 concerns (manual refcount, cache-mode coupling, duplicated encode loop) are all resolved. The current shape is well-scoped for the value it delivers.

Action items:

  1. Add missing #include <vector> in program_manager.h per reviewdog.
  2. @qjia7 or @hariharans29 — please stamp on the current head.
  3. Merge on green after the lint fix.
  4. Follow-ups (post-merge, non-blocking):
    • Expose max_num_pending_dispatches_ via config for cross-model tuning.
    • Add direct unit tests for the 5-state lifecycle and in-window dedup.
    • Cross-backend perf data (Metal, Vulkan) to broaden the perf story.

Ready to merge once the include is added and one Microsoft-side reviewer signs off.

@xiaofeihan1

Copy link
Copy Markdown
Contributor Author

Hariharan Seshadri (@hariharans29) Thanks for the detailed review. I addressed the merge-blocking lint item by adding the missing direct #include <vector> to program_manager.h in dce3ff5.

A couple of the non-blocking suggestions are already covered in the current revision:

  • maxNumPendingDispatches is exposed as a WebGPU EP provider option and validated in webgpu_provider_factory.cc.
  • DispatchCommand() already enforces command.compute_pipeline.has_value() with ORT_ENFORCE.

Validation completed:

  • lintrunner -a passed.
  • onnxruntime_providers_webgpu built and linked successfully in RelWithDebInfo.

@qjia7 Jiajia Qin (qjia7) 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.

Overall looks good. Just some nits.

Comment thread onnxruntime/core/providers/webgpu/webgpu_kernel.cc Outdated
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.cc
@hariharans29
Hariharan Seshadri (hariharans29) merged commit 49193db into main Aug 10, 2026
88 of 89 checks passed
@hariharans29
Hariharan Seshadri (hariharans29) deleted the feature/webgpu-defer-dispatch branch August 10, 2026 21:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants