Enable WebGPU graph capture for Gemma 4 decoder (#380)

## 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](microsoft/onnxruntime#29236))
and INT64 support for Equal/Sub/Where/ReduceSum
([microsoft/onnxruntime#29392](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.json` will have `enableGraphCapture: "1"` set
for **all** sessions by default. Until the follow-up PR lands, you must
manually edit `genai_config.json` to disable `enableGraphCapture` (set
to `"0"`) for the `vision`, `embedding`, and `speech` sections, keeping
it `"1"` only for `decoder`.

## Changes

### New rewrite rule: `static_empty_kv_rules` (`_static_empty_kv.py`)

- **Pattern**: `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)
- **Why it breaks graph capture**: `Shape` outputs to CPU;
`ConstantOfShape` is unsupported by WebGPU EP
- **Fix**: replace with a static `Constant(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 dtype
- Two variants cover the `CastLike` form (from onnxscript) and the
`Cast` form (post-quantization); BFLOAT16 handled by emitting a float32
constant followed by an explicit Cast (numpy has no native bfloat16)
- Rule applied automatically in `_optimizations.py` when
`enable_graph_capture=True`
- 6 unit tests in `_static_empty_kv_test.py`, built with `ir.Graph +
GraphBuilder` to match repo conventions

### `gemma4.py` changes

- **Split per-layer embedding tables**: the fused `[V, L*D]` table (~4.7
GB for Gemma4 E2B) exceeds WebGPU's 256 MiB `maxBufferSize` limit,
causing device loss. When `config.split_per_layer_embedding` is set, L
separate `[V, D]` tables (~128 MiB each) are used instead; shape
asserted before each `chunk()` call to catch layout drift early
- **Gather-based per-layer extraction**: replaced `Slice(combined,
starts=[i], ends=[i+1], axes=[2])` with `Gather(combined,
Constant(value_int=i), axis=2)` — `Slice` with INT64 axis inputs runs on
CPU under graph capture; `Gather` with a static scalar constant is
GPU-resident
- **`total_seq_len` scalar without `Shape`**: 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)` — `ReduceSum` produces `[batch]` INT32,
`Gather(..., 0)` extracts the scalar (valid because graph capture
requires `batch=1`). The non-capture branch continues to use
`Gather(Shape(attention_mask), 1)` unchanged.
- `Gemma4TextModel.__init__` now stores `self.config = config` so
`_compute_per_layer_inputs` can access it

### `_gemma4.py` changes

- **Split table routing**: `Gemma4Task.build()` computes the fused
table's byte size using `config.dtype.itemsize` and sets
`config.split_per_layer_embedding = fused_bytes > caps.max_buffer_size`;
when set, `per_layer_inputs` is omitted from the decoder graph and
`input_ids` is added instead

### `base.py` changes

- Added the `split_per_layer_embedding` flag, 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.py` changes

- **`EpCapabilities.max_buffer_size`**: new field (`0` = no limit). Set
to `268_435_456` (256 MiB) for WebGPU per the [W3C spec
default](https://www.w3.org/TR/webgpu/#typedefdef-gpusize64). Drives the
split-table decision in `Gemma4Task.build()` instead of a hardcoded EP
name check — any future EP with a tight buffer limit gets the split
automatically

### ORT WebGPU EP changes (separate PRs)

- **Indirect dispatch**
([microsoft/onnxruntime#29236](microsoft/onnxruntime#29236)):
prerequisite for graph capture correctness
- **INT64 for Equal/Sub/Where/ReduceSum**
([microsoft/onnxruntime#29392](microsoft/onnxruntime#29392)):
these ops were forcing CPU fallback when inputs were INT64, preventing
graph capture

## Test plan
- [x] 6 unit tests in `_static_empty_kv_test.py` pass
- [x] Exported INT4 Gemma4 WebGPU decoder: 0 `ConstantOfShape`, 0
`Shape`, 0 `Slice`, 0 `Squeeze`
- [x] All decoder nodes assigned to `WebGpuExecutionProvider`
- [x] End-to-end inference with graph capture ON: **90+ tok/s** (INT4
WebGPU) and OFF: **70+ tok/s**, coherent output
- [x] Exported a CPU INT4 decoder: coherent output vs WebGPU, but with
poor perf
- [x] Confirmed the mode applied the rewrite rules produce bit-identical
output compared to the original model that didn't apply the rewrite
rules on WebGPU ep.

---------

Signed-off-by: Fei Chen <feich@microsoft.com>
Signed-off-by: Copilot <copilot@github.com>
Co-authored-by: Claude Opus 4 (1M context) <noreply@anthropic.com>
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: 5 people

src/mobius/_configs/_base.py

@@ -1488,6 +1488,11 @@ class Gemma4Config(VisionLanguageConfig):
       causal attention (only ``None`` and ``"vision"`` are accepted).
     """
 
+    # 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.
+    split_per_layer_embedding: bool = False
+
     @classmethod
     def from_transformers(cls, config, parent_config=None) -> Gemma4Config:
         base = ArchitectureConfig.from_transformers(config, parent_config)

src/mobius/_execution_providers.py

@@ -97,6 +97,20 @@ class EpCapabilities:
             devices.  ``True`` only for WebGPU (consumer GPU); ``False`` for
             CUDA / CPU / DML / TRT-RTX where the runtime can handle large
             pre-allocations.
+        max_buffer_size: Maximum allowed size in bytes for a single model
+            weight buffer on this EP.  ``None`` means no limit.  When non-zero,
+            large weight tensors (e.g. fused per-layer embedding tables) must
+            be split into chunks that each fit within this bound.  WebGPU's
+            W3C spec default ``maxBufferSize`` is 268,435,456 bytes (256 MiB).
+        requires_graph_capture_rewrite: Whether this EP requires rewrite rules
+            to make models compatible with graph capture (e.g. replacing
+            ``Shape`` / ``ConstantOfShape`` with static alternatives for
+            shared-KV layer models like Gemma4).  Not all EPs with
+            ``enable_graph_capture`` need this — e.g. CUDA EP's ``Shape``
+            kernel is already registered inside the CUDA partition and is
+            graph-capture-safe.  Set ``True`` only for EPs that cannot execute
+            ``Shape`` / ``ConstantOfShape`` under graph capture (currently
+            WebGPU).
     """
 
     name: str
@@ -112,6 +126,8 @@ class EpCapabilities:
     enable_graph_capture: bool = False
     supports_past_present_share_buffer: bool = False
     cap_kv_buffer_max_length: bool = False
+    max_buffer_size: int | None = None
+    requires_graph_capture_rewrite: bool = False
 
     def __post_init__(self) -> None:
         if not self.supports_fused_rope and self.qkv_pack_dtypes:
@@ -281,6 +297,9 @@ def _register_builtins() -> None:
             enable_graph_capture=True,
             supports_past_present_share_buffer=True,
             cap_kv_buffer_max_length=True,
+            # W3C WebGPU spec default maxBufferSize (https://www.w3.org/TR/webgpu/)
+            max_buffer_size=268_435_456,  # 256 MiB
+            requires_graph_capture_rewrite=True,
         ),
         EpCapabilities(
             name="trt-rtx",

src/mobius/_optimizations.py

@@ -68,6 +68,7 @@
     separate_rope_rules,
     skip_layer_norm_rules,
     skip_norm_rules,
+    static_empty_kv_rules,
     unpack_qkv_rules,
 )
 
@@ -317,6 +318,16 @@ def _get_optimization_passes(
     if not caps.supports_rank4_rmsnorm:
         lower.append(("HtpRank4RMSNorm", list(htp_rank4_rmsnorm_rules())))
 
+    # --- Graph-capture rewrite ---
+    # Supports graph capture for shared-KV layer models on WebGPU EP
+    # (currently Gemma4). Pattern-based: only fires on models that emit the
+    # dynamic Shape → ConstantOfShape → Cast empty-KV pattern. Gated by
+    # requires_graph_capture_rewrite rather than enable_graph_capture because
+    # not all graph-capture EPs need it — e.g. CUDA EP's Shape kernel is
+    # already graph-capture-safe.
+    if caps.requires_graph_capture_rewrite:
+        lower.append(("StaticEmptyKV", list(static_empty_kv_rules())))
+
     return fuse, lower
 
 

src/mobius/models/gemma4.py

@@ -1478,6 +1478,7 @@ class Gemma4TextModel(nn.Module):
 
     def __init__(self, config: Gemma4Config):
         super().__init__()
+        self.config = config
         self._dtype = config.dtype
 
         embed_scale = math.sqrt(config.hidden_size)
@@ -1559,14 +1560,30 @@ def __init__(self, config: Gemma4Config):
         if self._per_layer_dim:
             self._num_layers = config.num_hidden_layers
             vocab_per_layer = getattr(config, "vocab_size_per_layer_input", 0)
-            # Single fused [V, L*D] table. Requires ORT >= 1.27 for CUDA
-            # Gather int64 index support (onnxruntime#28107).
+            # Fused [V, L*D] table — used when split_per_layer_embedding is False.
+            # Requires ORT >= 1.27 for CUDA Gather int64 index support (onnxruntime#28107).
             self.embed_tokens_per_layer = Gemma3TextScaledWordEmbedding(
                 vocab_per_layer,
                 self._num_layers * self._per_layer_dim,
                 config.pad_token_id,
                 embed_scale=float(self._per_layer_dim**0.5),
             )
+            # Split [V, D] tables — used when split_per_layer_embedding is True
+            # (i.e. the fused table exceeds the EP's max_buffer_size, e.g. WebGPU's
+            # 256 MiB limit; ~128 MiB each vs ~4.7 GB fused).
+            # Only the table actually called in forward() is realized as an
+            # ONNX initializer, so the unused one adds no graph weight.
+            self.embed_tokens_per_layer_split = nn.ModuleList(
+                [
+                    Gemma3TextScaledWordEmbedding(
+                        vocab_per_layer,
+                        self._per_layer_dim,
+                        config.pad_token_id,
+                        embed_scale=float(self._per_layer_dim**0.5),
+                    )
+                    for _ in range(self._num_layers)
+                ]
+            )
             self.per_layer_model_projection = Linear(
                 config.hidden_size,
                 config.num_hidden_layers * self._per_layer_dim,
@@ -1605,17 +1622,26 @@ def _compute_per_layer_inputs(
                 masked_ids,
             )
 
-        fused_emb = self.embed_tokens_per_layer(op, masked_ids)
-        fused_emb = op.Reshape(
-            fused_emb,
-            op.Constant(value_ints=[0, 0, self._num_layers, self._per_layer_dim]),
-        )
+        if self.config.split_per_layer_embedding:
+            # L separate Gathers on [V, D] tables — each fits within the EP's
+            # max_buffer_size (e.g. WebGPU's 256 MiB limit).
+            per_layer_embs = [
+                op.Unsqueeze(self.embed_tokens_per_layer_split[i](op, masked_ids), [2])
+                for i in range(self._num_layers)
+            ]
+            fused_emb = op.Concat(*per_layer_embs, axis=2)  # [B, S, L, D]
+        else:
+            fused_emb = self.embed_tokens_per_layer(op, masked_ids)
+            fused_emb = op.Reshape(
+                fused_emb,
+                op.Constant(value_ints=[0, 0, self._num_layers, self._per_layer_dim]),
+            )
 
         combined = op.Add(proj, fused_emb)
         combined = op.Mul(combined, float(0.5**0.5))
 
         return [
-            op.Squeeze(op.Slice(combined, starts=[i], ends=[i + 1], axes=[2]), [2])
+            op.Gather(combined, op.Constant(value_int=i), axis=2)
             for i in range(self._num_layers)
         ]
 
@@ -1721,17 +1747,25 @@ def forward(
             # seqlens_k[b] = sum(attention_mask[b]) - 1  (last valid KV idx)
             # total_seq_len = attention_mask.shape[1]     (past + current)
             one_i32 = op.Constant(value_int=1)
+            reduce_sum = op.ReduceSum(attention_mask, [1], keepdims=0)
             seqlens_k = op.Cast(
-                op.Sub(
-                    op.ReduceSum(attention_mask, [1], keepdims=0),
-                    one_i32,
-                ),
-                to=ir.DataType.INT32,
-            )
-            total_seq_len = op.Cast(
-                op.Gather(op.Shape(attention_mask), 1),
+                op.Sub(reduce_sum, one_i32),
                 to=ir.DataType.INT32,
             )
+            if caps.requires_graph_capture_rewrite:
+                # Support graph capture for shared-KV layer models on WebGPU EP.
+                # Derive total_seq_len from reduce_sum (already computed) as a
+                # scalar INT32 via Gather index 0 (valid because graph capture
+                # requires batch=1).
+                total_seq_len = op.Gather(
+                    op.Cast(reduce_sum, to=ir.DataType.INT32),
+                    op.Constant(value_int=0),
+                )
+            else:
+                total_seq_len = op.Cast(
+                    op.Gather(op.Shape(attention_mask), 1),
+                    to=ir.DataType.INT32,
+                )
 
             # Per-layer-type GQA contexts with appropriate cos/sin caches
             # and local_window_size for sliding layers.
@@ -1903,6 +1937,7 @@ def preprocess_weights(
                 state_dict.pop(key, None)
         # HF's model.embed_tokens_per_layer.weight [V, L*D] maps directly
         # to our fused embedding table — no splitting needed.
+        # (For WebGPU, splitting is handled by _Gemma4DecoderModel.preprocess_weights.)
         # Map HF expert weight names and fold router scale
         _remap_moe_expert_weights(state_dict, self.config)
         return super().preprocess_weights(state_dict)
@@ -1966,7 +2001,23 @@ def forward(
     def preprocess_weights(
         self, state_dict: dict[str, torch.Tensor]
     ) -> dict[str, torch.Tensor]:
-        return vlm_decoder_weights(state_dict, tie=self.config.tie_word_embeddings)
+        state_dict = vlm_decoder_weights(state_dict, tie=self.config.tie_word_embeddings)
+        # For WebGPU: split the fused [V, L*D] per-layer embedding into L separate [V, D] tables.
+        per_layer_dim = self.config.hidden_size_per_layer_input
+        if per_layer_dim and self.config.split_per_layer_embedding:
+            fused_key = "model.embed_tokens_per_layer.weight"
+            if fused_key in state_dict:
+                num_layers = self.config.num_hidden_layers
+                fused = state_dict.pop(fused_key)
+                assert fused.shape[1] == num_layers * per_layer_dim, (
+                    f"{fused_key} dim 1 expected {num_layers * per_layer_dim} "
+                    f"({num_layers} layers x {per_layer_dim} per_layer_dim), "
+                    f"got {fused.shape[1]}"
+                )
+                chunks = fused.chunk(num_layers, dim=1)
+                for i, chunk in enumerate(chunks):
+                    state_dict[f"model.embed_tokens_per_layer_split.{i}.weight"] = chunk
+        return state_dict
 
 
 class _Gemma4VisionEncoderModel(nn.Module):
@@ -2183,6 +2234,12 @@ def forward(
         if not self._per_layer_dim:
             return outputs
 
+        # When split_per_layer_embedding is set, the per-layer computation runs
+        # inside the decoder using split [V, D] tables.  The embedding model only
+        # emits inputs_embeds in that case.
+        if self.config.split_per_layer_embedding:
+            return outputs
+
         # Compute per-layer input embeddings (moved from the decoder).
         # 1. Project hidden states → [B, S, L*D] and scale by hidden_size**-0.5
         proj = self.per_layer_model_projection(op, hidden)
@@ -2737,6 +2794,33 @@ def preprocess_weights(
         # Map HF expert weight names and fold router scale
         _remap_moe_expert_weights(renamed, self.config)
 
+        # For WebGPU: the fused [V, L*D] embed_tokens_per_layer exceeds the 256 MiB
+        # per-buffer limit.  Split it into L separate [V, D] tables in the decoder.
+        # The per_layer_projection weights also live in the decoder (not embedding).
+        if self.config.split_per_layer_embedding:
+            fused_key = "embedding.embed_tokens_per_layer.weight"
+            if fused_key in renamed:
+                num_layers = self.config.num_hidden_layers
+                per_layer_dim = self.config.hidden_size_per_layer_input
+                fused = renamed.pop(fused_key)
+                assert fused.shape[1] == num_layers * per_layer_dim, (
+                    f"{fused_key} dim 1 expected {num_layers * per_layer_dim} "
+                    f"({num_layers} layers x {per_layer_dim} per_layer_dim), "
+                    f"got {fused.shape[1]}"
+                )
+                chunks = fused.chunk(num_layers, dim=1)
+                for i, chunk in enumerate(chunks):
+                    renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.weight"] = chunk
+            # Re-route the projection weights from embedding.* → decoder.model.*
+            for k in list(renamed.keys()):
+                if k.startswith(
+                    (
+                        "embedding.per_layer_model_projection.",
+                        "embedding.per_layer_projection_norm.",
+                    )
+                ):
+                    renamed[k.replace("embedding.", "decoder.model.", 1)] = renamed.pop(k)
+
         return renamed
 
 

src/mobius/rewrite_rules/__init__.py

@@ -39,6 +39,7 @@
     "separate_rope_rules",
     "skip_layer_norm_rules",
     "skip_norm_rules",
+    "static_empty_kv_rules",
     "unpack_qkv_rules",
 ]
 
@@ -56,4 +57,5 @@
 from mobius.rewrite_rules._separate_rope import separate_rope_rules
 from mobius.rewrite_rules._skip_layer_norm import skip_layer_norm_rules
 from mobius.rewrite_rules._skip_norm import skip_norm_rules
+from mobius.rewrite_rules._static_empty_kv import static_empty_kv_rules
 from mobius.rewrite_rules._unpack_qkv import unpack_qkv_rules