Skip to content

[WebGPU] Enable graph capture for Gemma4-class decoder models#28955

Closed
feich-ms wants to merge 4 commits into
mainfrom
user/feich/webgpu-gemma4-graph-capture
Closed

[WebGPU] Enable graph capture for Gemma4-class decoder models#28955
feich-ms wants to merge 4 commits into
mainfrom
user/feich/webgpu-gemma4-graph-capture

Conversation

@feich-ms

@feich-ms feich-ms commented Jun 9, 2026

Copy link
Copy Markdown
Contributor

Note: This PR is intentionally kept as draft and will NOT be merged.
The GqaMaskReformattingToGraphCapture transformer is a temporary workaround for Gemma4 models that were built without enable_webgpu_graph=true (the genai builder does not yet have a Gemma4 dispatch entry). Once the genai team adds Gemma4 builder support and official models are converted with graph capture enabled, this ORT-side transformer becomes unnecessary. Until then, this branch serves as a local development tool for WebGPU graph capture on Gemma4.


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:

  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) — 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.

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

feich-ms and others added 2 commits June 4, 2026 19:53
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>

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

}

NodeArg& AddVec1Initializer(Graph& graph, const std::string& name_hint,
int data_type, const void* data, size_t byte_len) {

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.

Suggested change
int data_type, const void* data, size_t byte_len) {
int data_type, const void* data, size_t byte_len) {

@@ -0,0 +1,231 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
feich-ms and others added 2 commits June 11, 2026 11:20
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>
@feich-ms feich-ms closed this Jun 23, 2026
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.

2 participants