[WebGPU] Enable graph capture for Gemma4-class decoder models#28955
Closed
feich-ms wants to merge 4 commits into
Closed
[WebGPU] Enable graph capture for Gemma4-class decoder models#28955feich-ms wants to merge 4 commits into
feich-ms wants to merge 4 commits into
Conversation
Graph capture (GC) requires every kernel and buffer interaction in the
captured replay to be deterministic and self-contained. Several existing
WebGPU paths violated that — they relied on CPU-side branching, opset-
open registration, or buffer-bucket reuse. This change makes the EP
GC-clean for Gemma-class decoders without regressing the non-GC path.
Six coordinated changes, each gated so non-GC sessions are unaffected:
1. ConstantOfShape kernel + opset-by-opset registration
New onnxruntime/core/providers/webgpu/generator/constant_of_shape.{cc,h}.
Uses bitcast<output_value_t>(uniforms.value) so a single shader covers
all dtypes. Shape input marked OrtMemTypeCPU. Registered explicitly at
versions 9, 20, 21, 23, 24 because KernelRegistry::VerifyVersion treats
open-ended SinceVersion(N) as exact-match against node.SinceVersion().
2. GqaMaskReformattingToGraphCapture transformer (gated)
New onnxruntime/core/optimizer/gqa_mask_reformatting_to_graph_capture.{cc,h}.
Rewrites the int64 GQA-mask construction subgraph into a GC-replayable
form (Cast int64 to int32 + ReduceSum/Sub/ReduceMax). Registered in
graph_transformer_utils.cc only when
ep.webgpuexecutionprovider.enableGraphCapture == 1, so default sessions
see no transform.
3. Where int64 scalar path
tensor/where.{cc,h}: added is_scalar_ shader variant that reads the
condition word and extracts the selected bit, then reads scalar a/b.
Drives component=1 output for int64 (stored as vec2<u32>). Type
constraints widened to include int64_t. CacheHint disambiguates
broadcast x int64 variants.
4. Equal / binary-elementwise int64 path + widened type constraints
math/binary_elementwise_ops.{cc,h}: int64 shader path packs four scalar
reads per global_idx into a vec4 to keep output component=4 while
inputs are component=1 (int64 cannot vector-load). Equal registered
for int64 across versions 7-10, 11-12, 13-18, 19+.
5. captured_buffers_ quarantine in GraphCacheManager / GraphSimpleCacheManager
buffer_manager.cc: ReleaseBuffer now queues into pending_buffers_;
OnRefresh either returns them to buckets (Default state) or moves them
to captured_buffers_ when Capturing/Replaying. Destructor releases
both lists. This reverts the PR #27453 immediate-reuse optimization
only for the GC code path — non-GC behavior is unchanged.
6. PrepareIndirectDispatch helper + widened use_indirect_dispatch
contrib_ops/webgpu/bert/flash_attention.{cc,h}: new
PrepareIndirectDispatchProgram single-thread shader writes
[tile_count, num_heads, 1] from seqlen_k[0]+1. use_indirect_dispatch
is now true when GC is on OR the GPU-resident seqlens_k sentinel is 0,
so the dispatch shape no longer leaks through the CPU. Added an
explicit PrepareIndirectDispatch call on the kv_empty path because
CopyKVCache is skipped there under GC.
Verified on Gemma4 E2B Q4_K_M (opset 24, com.microsoft 1, 20 KV-shared
layers): GC=ON produces byte-identical output to GC=OFF and improves
single-batch decode throughput substantially. Non-GC sessions show no
behavioral change.
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Simplify the graph transformer without changing its behavior: - Replace TryResolveConstantInt (51 lines) with compact IsConstantInt - Merge AddInt64Vec1Initializer/AddInt32Vec1Initializer into AddVec1Initializer - Replace TryRemoveIfOrphan with batch RemoveChain - Inline chain matching directly in ApplyImpl with early-continue - Remove SeqlensKChain/TotalSeqLenChain structs and separate matchers Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
| } | ||
|
|
||
| NodeArg& AddVec1Initializer(Graph& graph, const std::string& name_hint, | ||
| int data_type, const void* data, size_t byte_len) { |
Contributor
There was a problem hiding this comment.
Suggested change
| int data_type, const void* data, size_t byte_len) { | |
| int data_type, const void* data, size_t byte_len) { |
Remove pending_buffers_ intermediate queue. Instead, track the current GraphCaptureState in a member and branch directly in ReleaseBuffer: Default mode does immediate bucket reuse (preserving PR #27453's optimization), Capturing/Replaying quarantines to captured_buffers_. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Tested end-to-end: graph capture produces correct output ("2+2 is 4")
with PR #27453's immediate buffer reuse. The quarantine was incorrectly
attributed as a fix during the original GC investigation. Phi-4 graph
capture also works without it, confirming immediate reuse is safe.
Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Graph capture (GC) requires every kernel and buffer interaction in the captured replay to be deterministic and self-contained. Several existing WebGPU paths violated that — they relied on CPU-side branching or opset-open registration. This change makes the EP GC-clean for Gemma-class decoders without regressing the non-GC path.
Five coordinated changes, each gated so non-GC sessions are unaffected:
ConstantOfShape kernel + opset-by-opset registration New onnxruntime/core/providers/webgpu/generator/constant_of_shape.{cc,h}. Uses bitcast<output_value_t>(uniforms.value) so a single shader covers all dtypes. Shape input marked OrtMemTypeCPU. Registered explicitly at versions 9, 20, 21, 23, 24 because KernelRegistry::VerifyVersion treats open-ended SinceVersion(N) as exact-match against node.SinceVersion().
GqaMaskReformattingToGraphCapture transformer (gated) — shrunk from 433 to 231 lines: simplified matching logic (inlined chain matching with early-continue, replaced TryResolveConstantInt with compact IsConstantInt, merged initializer helpers). Rewrites the int64 GQA-mask construction subgraph into a GC-replayable form (Cast int64 to int32 + ReduceSum/Sub/ReduceMax). Registered in graph_transformer_utils.cc only when ep.webgpuexecutionprovider.enableGraphCapture == 1, so default sessions see no transform.
Where int64 scalar path tensor/where.{cc,h}: added is_scalar_ shader variant that reads the condition word and extracts the selected bit, then reads scalar a/b. Drives component=1 output for int64 (stored as vec2). Type constraints widened to include int64_t. CacheHint disambiguates broadcast x int64 variants.
Equal / binary-elementwise int64 path + widened type constraints math/binary_elementwise_ops.{cc,h}: int64 shader path packs four scalar reads per global_idx into a vec4 to keep output component=4 while inputs are component=1 (int64 cannot vector-load). Equal registered for int64 across versions 7-10, 11-12, 13-18, 19+.
PrepareIndirectDispatch helper + widened use_indirect_dispatch contrib_ops/webgpu/bert/flash_attention.{cc,h}: new PrepareIndirectDispatchProgram single-thread shader writes [tile_count, num_heads, 1] from seqlen_k[0]+1. use_indirect_dispatch is now true when GC is on OR the GPU-resident seqlens_k sentinel is 0, so the dispatch shape no longer leaks through the CPU. Added an explicit PrepareIndirectDispatch call on the kv_empty path because CopyKVCache is skipped there under GC.
Note: GraphCacheManager buffer quarantine (previously listed as change #5) was verified as NOT load-bearing — graph capture produces correct output with PR #27453's immediate buffer reuse. The quarantine has been reverted.
Verified on Gemma4 E2B Q4_K_M (opset 24, com.microsoft 1, 20 KV-shared layers): GC=ON produces byte-identical output to GC=OFF and improves single-batch decode throughput substantially. Non-GC sessions show no behavioral change.