[WebGPU] Deferred-dispatch to parallelize cold-start shader compilation - #29557
Conversation
There was a problem hiding this comment.
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::Buildto support an asynchronous mode returning awgpu::Futureplus 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. |
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.
48e08a1 to
88b7180
Compare
Xiaofei Han (xiaofeihan1)
left a comment
There was a problem hiding this comment.
add comments
Retain deferred dispatch buffers until encoding, clarify in-window pipeline build sharing, remove an unused include, and clean up ProgramArtifact construction.
|
Thanks for tackling this — the cold-start serialization of 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:
Proposal: cache the bind group, not the raw buffersInstead of recording raw
When the batch reaches the max window: wait for all promises → then 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:
The one prerequisite: explicit bind group layoutThe 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 Bonus: this unifies with the graph-capture path
One thing this does not changeProfiling 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.
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.
- Allow the bind-group binding count to equal maxBindingsPerBindGroup (< -> <=) in CreateBindGroup and CreatePipelineLayout; the count reaching the limit is valid. - Fix CreatePipelineLayout declaration indentation.
|
Thanks for the detailed proposal — a few updates on where things landed:
On reusing That said, you're right that the encode / timestamp / window loop is duplicated between the deferred drain and |
Jiajia Qin (qjia7)
left a comment
There was a problem hiding this comment.
Two suggestions on the deferred-dispatch structure:
- 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.
- 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.
Merge deferred dispatch storage into CapturedCommandInfo, resolve pending pipelines during flush, and share command dispatch logic with graph replay.
|
Jiajia Qin (@qjia7) Thanks for your suggestion!
Done. I moved it into Flush method and rename it to
|
Review: PR #29557 — [WebGPU] Deferred-dispatch to parallelize cold-start shader compilation (head
|
| 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:
- Program-cache hit →
pending_build = empty,compute_pipeline = set. - First cache-miss for a key →
pending_build = set,compute_pipeline = empty. - Later cache-miss same key →
pending_build = empty,compute_pipeline = empty,program_keyset for cache lookup after wait. - Post-
WaitForDeferredPipelineBuilds→pending_build = empty,compute_pipeline = set. - 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:
- buffer_manager.cc →
ORT_THROW_IF_ERROR(context_.Flush(*this)) - webgpu_execution_provider.cc:
OnRunEndcaptures the status and runs a graph-capture cleanup path on failure before returning it - webgpu_kernel.cc:
PrePackmerges the flush status into its own return
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)
- Lint fix: reviewdog flagged
#include <vector>missing in program_manager.h (line 71,std::vector<int>¶meter). One-line addition. Should land before merge. - Hardcoded window size
max_num_pending_dispatches_ = 16. Author validated 16 and 4096. Consider exposing viaWebGpuContextConfigor env override for tuning across model shapes / driver contexts. Non-blocking. - Extensibility of
EncodeDeferredDispatchesinjection 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) havingEndComputePassitself invokeEncodeDeferredDispatcheswhen the window is non-empty, or (b) a comment + assertion inEndComputePassthat flags "unencoded deferred work + about-to-record-copy" as an error in debug builds. - State-machine assertion:
DispatchCommand()requirescompute_pipelineto be set. Add adebug_assert(command.compute_pipeline.has_value())orORT_ENFORCEat the top ofDispatchCommandif not already present — cheap defense against future refactors accidentally emitting a state-3 command past the wait. - 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
WebGpuContextfixture. Not blocking (perf integration testing covers happy-path), but a regression suite would harden future refactors. - 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.
- 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.
@qjia7has driven substantive review through multiple rounds and appears satisfied on the latest commit.@hariharans29requested. - Copilot AI: five review rounds, most recent generated no new comments.
- Labels: none.
ep:WebGPUwould 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:
- Add missing
#include <vector>in program_manager.h per reviewdog. @qjia7or@hariharans29— please stamp on the current head.- Merge on green after the lint fix.
- 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.
- Expose
Ready to merge once the include is added and one Microsoft-side reviewer signs off.
|
Hariharan Seshadri (@hariharans29) Thanks for the detailed review. I addressed the merge-blocking lint item by adding the missing direct A couple of the non-blocking suggestions are already covered in the current revision:
Validation completed:
|
Jiajia Qin (qjia7)
left a comment
There was a problem hiding this comment.
Overall looks good. Just some nits.
49193db
into
main
Summary
Parallelize WebGPU cold-start pipeline compilation by starting
CreateComputePipelineAsyncwithout 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 callingWait(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:
ProgramManager::Build()startsCreateComputePipelineAsyncand returns its future and callback state without waiting. It also builds an explicitwgpu::BindGroupLayoutsynchronously (from the input/output/uniform binding info) so the caller can create the bind group before the pipeline finishes compiling.WebGpuContext::Run()records the dispatch state needed later: the cache key, the already-createdwgpu::BindGroup(which internally retains its bound buffers), dispatch dimensions, indirect buffer, and optional profiling metadata.maxNumPendingDispatches(16 by default),WaitForDeferredPipelineBuildsAndEncodeDispatches():ProgramArtifacts into the program cache; andLaunchComputePipeline()for every recorded dispatch in original order, including pipeline/bind-group setup and dispatch encoding.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, whereGpuBufferAllocator::Free()reachesBufferManager::Release()and immediately callswgpuBufferRelease().Rather than tracking raw
WGPUBufferhandles, eachDeferredDispatchcaches the already-createdwgpu::BindGroup. A WebGPU bind group internally holds a reference to every buffer it binds, so caching the bind group is the buffer retention — no manualwgpuBufferReference/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 whyProgramManager::Build()produces an explicitwgpu::BindGroupLayoutsynchronously (decoupled from the async pipeline compile).3. Asynchronous callback lifetime
The async pipeline callback writes the compiled pipeline into a heap-allocated
PipelineCallbackContextthat owns thewgpu::ComputePipelineby 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 isnoexceptand reports failures throughStatusrather than letting exceptions cross the Dawn callback boundary.Performance
Phi-4 Mini prefill latency on WebGPU/D3D12:
Validation
onnxruntime_providers_webgpu, the WebGPU shared library, and theonnxruntime_webgpuwheel inRelWithDebInfo.easyandprefill-500execution with the default 16-dispatch window.storageBufferCacheMode=disabledwith window sizes 16 and 4096.