Enable WebGPU graph capture for Gemma 4 decoder - #380
Merged
Conversation
Performance Comparison
|
|
The author of this PR, feich-ms, is not an activated member of this organization on Codecov. |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
justinchuby
reviewed
Jun 26, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
This PR aims to make Gemma 4 exports compatible with WebGPU EP graph capture by adjusting the ONNX graphs at export time (avoiding unsupported ops and CPU fallbacks), plus updating the GQA rewrite rule and the Gemma4 ORT GenAI example to support --ep webgpu / --dtype f16.
Changes:
- Add EP-capability branching (notably
enable_graph_capture) to emit INT32 inputs/const paths and reduce CPU fallbacks on WebGPU. - Update Gemma4 model export to avoid
Shape/ConstantOfShapein key graph-capture-sensitive areas, and to support internal per-layer embedding computation. - Extend the GroupQueryAttention rewrite rule to match constant dtypes against
attention_maskdtype, and update the example exporter CLI.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
src/mobius/tasks/_gemma4.py |
EP-capability gated INT32 inputs; decoder/embedding contract adjustments for internal per-layer embeddings |
src/mobius/models/gemma4.py |
Graph-capture export adjustments (empty-KV creation, per-layer embedding handling, seq-len derivation changes) |
src/mobius/models/base.py |
Graph-capture path casts attention_mask before ReduceSum for GQA context derivation |
src/mobius/rewrite_rules/_group_query_attention.py |
Make rewrite rule constants dtype-aware (INT32 vs INT64) to avoid type mismatches |
examples/gemma4_ort_genai.py |
Add --ep/--dtype options and propagate EP into genai_config generation |
feich-ms
force-pushed
the
user/feich/rewrite_gemma4_for_graph_capture
branch
from
July 1, 2026 04:17
8ac8069 to
260973c
Compare
feich-ms
force-pushed
the
user/feich/rewrite_gemma4_for_graph_capture
branch
2 times, most recently
from
July 1, 2026 09:50
fe43cfb to
2b07034
Compare
justinchuby
reviewed
Jul 1, 2026
feich-ms
marked this pull request as ready for review
July 2, 2026 09:05
justinchuby
reviewed
Jul 2, 2026
WebGPU has a 2 GiB per-buffer limit. The fused [V, L*D] embed_tokens_per_layer table (262144 × 8960 bfloat16 ≈ 4.70 GB) exceeds that limit and causes device loss. When targeting WebGPU EP: - Gemma4Config.split_per_layer_embedding is set True in build() - Gemma4TextModel creates both fused and split module lists; only the path called in forward() is realized as ONNX initializers - _compute_per_layer_inputs dispatches to 35 split [V,D] tables (~134 MB each) via ep_capabilities() at forward time - Gemma4EmbeddingModel skips per-layer computation (returns only inputs_embeds); the decoder owns the full per-layer pipeline - Gemma4Model.preprocess_weights splits embed_tokens_per_layer.weight and re-routes per_layer_projection weights to decoder.model.* - _build_decoder adds input_ids instead of per_layer_inputs for WebGPU Tested: model runs at 82.8 tok/s on WebGPU EP (no device crash). Non-WebGPU paths use the original fused table and are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
Slice with INT64 axes inputs produces CPU tensors under WebGPU graph capture, inserting MemcpyFromHost nodes that block capture. Gathering by scalar index (Gather on axis=2) stays on GPU and matches the shape expected by each decoder layer. Co-Authored-By: Claude <noreply@anthropic.com>
…en; correct buffer size comments to 256 MiB The ReduceSum over attention_mask already produces a [batch] tensor; the Squeeze to scalar was unnecessary and caused the Cast node to receive a CPU-assigned tensor, blocking WebGPU graph capture. Drop the Squeeze and cast [batch] directly, matching the reference model structure. Also update all buffer size comments from the incorrect "2 GiB" to the W3C WebGPU spec's actual maxBufferSize default of 256 MiB. Co-Authored-By: Claude <noreply@anthropic.com>
Vision/embedding/audio encoders contain ops (e.g. Slice with runtime INT64 axes) that run on CPU under graph capture, blocking the capture boundary. Add enable_graph_capture override to make_provider_options and _make_session_options, and pass False for these components so only the decoder runs with graph capture enabled. Co-Authored-By: Claude <noreply@anthropic.com>
…fig" This reverts commit 32280b5.
Add EpCapabilities.max_buffer_size (0 = no limit; 268_435_456 for WebGPU per the W3C spec default) so that the per-layer embedding split is driven by a capability rather than an EP name check. Gemma4Task.build() now computes the fused table size in bytes and sets split_per_layer_embedding = True only when it exceeds the EP's max_buffer_size. All downstream logic (gemma4.py, _gemma4.py) reads config.split_per_layer_embedding instead of calling ep_capabilities().name == "webgpu". Co-Authored-By: Claude <noreply@anthropic.com>
…can access it _compute_per_layer_inputs uses getattr(self.config, ...) but Gemma4TextModel.__init__ never stored config — only set self._dtype. Add self.config = config to unblock the three failing build_graph / weight_alignment tests. Co-Authored-By: Claude <noreply@anthropic.com>
…y / _gemma4.py Co-Authored-By: Claude <noreply@anthropic.com>
- gemma4.py: Fix total_seq_len to be scalar in graph-capture branch via Gather(reduce_sum, 0) instead of passing the [batch] vector to GQA - gemma4.py: Add assert before each chunk() call to catch wrong fused embedding shape early with a clear error message - _gemma4.py: Derive dtype_bytes from config.dtype.itemsize instead of hardcoding 2 (would undercount for fp32 builds) - _static_empty_kv.py: Fix BFLOAT16 handling — emit float32 Constant + explicit Cast instead of silently emitting wrong-dtype float32 constant; both CastLike and Cast variants updated - _group_query_attention.py: Guard RotaryAttentionToGQA.check() with an attention_mask presence check to fail the match cleanly rather than crashing with AttributeError in rewrite() Co-Authored-By: Claude <noreply@anthropic.com>
Replace onnx.helper / onnx.TensorProto / ir.serde.deserialize_model with the ir.Graph + GraphBuilder pattern used elsewhere in the repo (e.g. _htp_rank4_rmsnorm_test.py). Drops the onnx protobuf dependency from the test module entirely. Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
…q_len The previous fix added Gather(reduce_sum, 0) to extract a scalar from the [batch] ReduceSum output, believing GQA required a scalar for total_seq_len. This was wrong: GQA accepts a [batch] INT32 vector (confirmed by working models). The extra Gather was a CPU INT64 op that broke WebGPU graph capture partitioning — ORT refused to enable graph capture because not all nodes could be assigned to WebGpuExecutionProvider. Fix: cast reduce_sum directly to INT32, no Gather. Matches the pattern in all previously-working webgpu models. Co-Authored-By: Claude <noreply@anthropic.com>
Cast reduce_sum to INT32 before Gather(_, 0) so the scalar extraction is done entirely in INT32 (WebGPU-compatible). The previous attempt used Gather on the raw INT64 reduce_sum output, which WebGPU EP cannot handle, breaking graph capture partitioning. Sequence: reduce_sum [batch,INT64] → Cast(INT32) → Gather(0) → scalar INT32 GQAContext.total_seq_len is documented as scalar INT32; valid here because graph capture requires batch=1. Co-Authored-By: Claude <noreply@anthropic.com>
- Remove getattr guards on Gemma4Config.split_per_layer_embedding and hidden_size_per_layer_input: both are always-present fields (default False/0 in Gemma4Config), so direct attribute access is correct. - Change max_buffer_size default from 0 to None for clearer semantics (absent vs. zero-byte limit). - Clarify graph-capture comments: Shape outputs to CPU on any GPU EP (WebGPU and CUDA alike); ConstantOfShape is an additional WebGPU-only constraint. The StaticEmptyKV rule and total_seq_len branch are correct for any EP with enable_graph_capture=True. - Replace _IR_DTYPE_TO_NP dict with ir_dtype.numpy() in _static_zero_constant; handles bfloat16 natively via ml_dtypes, eliminating the float32 fallback + Cast. Co-Authored-By: Claude <noreply@anthropic.com>
Add EpCapabilities.apply_graph_capture_rewrite to explicitly control which EPs receive graph-capture rewrite rules (StaticEmptyKV and total_seq_len derivation). Currently enabled for WebGPU only, since CUDA EP's Shape kernel is already graph-capture-safe and does not need these rewrites. Co-Authored-By: Claude <noreply@anthropic.com>
Rename the flag to better fit the "capability" pattern — the EP *requires* these rewrites to support graph capture, rather than choosing to *apply* them. Co-Authored-By: Claude <noreply@anthropic.com>
feich-ms
force-pushed
the
user/feich/rewrite_gemma4_for_graph_capture
branch
from
July 8, 2026 00:41
19d1824 to
e36045b
Compare
…e line
The linter expected the multi-line form:
return op.Constant(
value=ir.tensor(np.zeros((1, 0, kv_hidden), dtype=ir_dtype.numpy()))
)
to be collapsed to a single line, since the total length fits within the
95-char line limit.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Fei Chen <feich@microsoft.com>
justinchuby
approved these changes
Jul 8, 2026
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.
Summary
Enable WebGPU graph capture for Gemma 4 decoder. All graph capture compatibility is handled at the ONNX export level in Mobius. Requires two ORT PRs: indirect dispatch support (microsoft/onnxruntime#29236) and INT64 support for Equal/Sub/Where/ReduceSum (microsoft/onnxruntime#29392).
Note: Graph capture is currently supported for the decoder only. Vision, embedding, and audio encoders each have ops that fall back to CPU under graph capture (investigated, to be addressed in a follow-up PR). When Mobius builds a WebGPU model with graph capture enabled, the generated
genai_config.jsonwill haveenableGraphCapture: "1"set for all sessions by default. Until the follow-up PR lands, you must manually editgenai_config.jsonto disableenableGraphCapture(set to"0") for thevision,embedding, andspeechsections, keeping it"1"only fordecoder.Changes
New rewrite rule:
static_empty_kv_rules(_static_empty_kv.py)Shape → Concat → ConstantOfShape → CastLike/Cast— used to build the empty[batch, 0, kv_hidden]KV tensor for Gemma4's shared-KV layers (layers 15–34)Shapeoutputs to CPU;ConstantOfShapeis unsupported by WebGPU EPConstant(zeros([1, 0, kv_hidden], dtype))— batch fixed to 1 (required by graph capture's static-shape requirement), dtype inferred from the model's activation dtypeCastLikeform (from onnxscript) and theCastform (post-quantization); BFLOAT16 handled by emitting a float32 constant followed by an explicit Cast (numpy has no native bfloat16)_optimizations.pywhenenable_graph_capture=True_static_empty_kv_test.py, built withir.Graph + GraphBuilderto match repo conventionsgemma4.pychanges[V, L*D]table (~4.7 GB for Gemma4 E2B) exceeds WebGPU's 256 MiBmaxBufferSizelimit, causing device loss. Whenconfig.split_per_layer_embeddingis set, L separate[V, D]tables (~128 MiB each) are used instead; shape asserted before eachchunk()call to catch layout drift earlySlice(combined, starts=[i], ends=[i+1], axes=[2])withGather(combined, Constant(value_int=i), axis=2)—Slicewith INT64 axis inputs runs on CPU under graph capture;Gatherwith a static scalar constant is GPU-residenttotal_seq_lenscalar withoutShape: in the graph-capture branch,Shape(attention_mask)outputs to CPU and breaks the capture boundary. Fix:total_seq_len = Cast(Gather(ReduceSum(attention_mask, axis=1), 0), INT32)—ReduceSumproduces[batch]INT32,Gather(..., 0)extracts the scalar (valid because graph capture requiresbatch=1). The non-capture branch continues to useGather(Shape(attention_mask), 1)unchanged.Gemma4TextModel.__init__now storesself.config = configso_compute_per_layer_inputscan access it_gemma4.pychangesGemma4Task.build()computes the fused table's byte size usingconfig.dtype.itemsizeand setsconfig.split_per_layer_embedding = fused_bytes > caps.max_buffer_size; when set,per_layer_inputsis omitted from the decoder graph andinput_idsis added insteadbase.pychangessplit_per_layer_embeddingflag, set to True by Gemma4Task.build() when the target EP's max_buffer_size is too small for the fused [V, L*D] per-layer embedding table, to split it into L separate [V, D] tables that each fit within the EP's buffer limit._execution_providers.pychangesEpCapabilities.max_buffer_size: new field (0= no limit). Set to268_435_456(256 MiB) for WebGPU per the W3C spec default. Drives the split-table decision inGemma4Task.build()instead of a hardcoded EP name check — any future EP with a tight buffer limit gets the split automaticallyORT WebGPU EP changes (separate PRs)
Test plan
_static_empty_kv_test.pypassConstantOfShape, 0Shape, 0Slice, 0SqueezeWebGpuExecutionProvider