Enable GQA for KV-shared layers and fix dtype issues in vision and audio

src/mobius/models/gemma4.py

@@ -35,7 +35,6 @@
 
 from mobius._build_context import ep_capabilities
 from mobius._configs import ArchitectureConfig, Gemma4Config
-from mobius._flags import flags
 from mobius._weight_utils import vlm_decoder_weights, vlm_embedding_weights
 from mobius.components import (
     MLP,
@@ -736,41 +735,79 @@ def forward(
             )
 
         if self.is_kv_shared_layer:
-            # KV-shared layers always use standard Attention path because
-            # they borrow K,V from a source layer (no own KV cache).
-            if use_gqa:
-                raise ValueError(
-                    "KV-shared GQA path not yet implemented. "
-                    "Set MOBIUS_USE_GQA_FOR_KV_SHARED=0 or wait for "
-                    "ORT GQA new_kv_length=0 support."
-                )
-            # Borrow full-history K,V from source layer.
-            # present_key/value from the ONNX Attention op is 4D:
-            #   [batch, kv_heads, total_seq, head_dim]
-            # The Attention op expects key/value as 3D:
-            #   [batch, total_seq, kv_heads * head_dim]
-            # Transpose and reshape to match.
+            # KV-shared layers borrow K,V from a source layer (no own KV cache).
             src_key, src_value = shared_kv_states[self.kv_shared_layer_index]
 
-            # [B, kv_heads, total_seq, head_dim] → [B, total_seq, kv_heads*head_dim]
-            src_key = op.Transpose(src_key, perm=[0, 2, 1, 3])
-            src_key = op.Reshape(src_key, [0, 0, -1])
-            src_value = op.Transpose(src_value, perm=[0, 2, 1, 3])
-            src_value = op.Reshape(src_value, [0, 0, -1])
-
-            attn_output, present_key, present_value = _apply_attention(
-                op,
-                query_states,
-                src_key,
-                src_value,
-                attention_bias,
-                past_key=None,
-                past_value=None,
-                num_attention_heads=self.num_attention_heads,
-                num_key_value_heads=self.num_key_value_heads,
-                scale=self.scaling,
-                softcap=self.softcap,
-            )
+            if use_gqa:
+                # GQA path for shared KV: pass empty K/V tensors and wire the
+                # source layer's present_key/value as past_key/past_value.
+                # The shared buffer is already in BNSH format, so GQA reads it
+                # directly — no Transpose/Reshape needed.
+                gqa_ctx = attention_bias
+
+                # Create empty K/V tensors with kv_sequence_length=0.
+                # Shape: [batch, 0, kv_heads * head_dim]
+                batch_dim = op.Shape(query_states, start=0, end=1)
+                kv_hidden = self.num_key_value_heads * self.head_dim
+                empty_shape = op.Concat(
+                    batch_dim,
+                    op.Constant(value_ints=[0, kv_hidden]),
+                    axis=0,
+                )
+                empty_kv = op.CastLike(op.ConstantOfShape(empty_shape), query_states)
+
+                gqa_attrs: dict = {
+                    "num_heads": self.num_attention_heads,
+                    "kv_num_heads": self.num_key_value_heads,
+                    "scale": self.scaling,
+                    "do_rotary": 1,
+                    "rotary_interleaved": int(self._rope_interleave),
+                }
+                if self.softcap:
+                    gqa_attrs["softcap"] = self.softcap
+                if self.rotary_embedding_dim:
+                    gqa_attrs["rotary_embedding_dim"] = self.rotary_embedding_dim
+                if gqa_ctx.local_window_size > 0:
+                    gqa_attrs["local_window_size"] = gqa_ctx.local_window_size
+
+                attn_output, present_key, present_value = op.GroupQueryAttention(
+                    query_states,
+                    empty_kv,   # key: empty (kv_sequence_length=0)
+                    empty_kv,   # value: empty (kv_sequence_length=0)
+                    src_key,    # past_key: shared KV in BNSH
+                    src_value,  # past_value: shared KV in BNSH
+                    gqa_ctx.seqlens_k,
+                    gqa_ctx.total_seq_len,
+                    gqa_ctx.cos_cache,
+                    gqa_ctx.sin_cache,
+                    _domain="com.microsoft",
+                    _outputs=3,
+                    **gqa_attrs,
+                )
+            else:
+                # Fallback Attention path: transpose shared KV from BNSH to 3D.
+                # present_key/value from the ONNX Attention op is 4D:
+                #   [batch, kv_heads, total_seq, head_dim]
+                # The Attention op expects key/value as 3D:
+                #   [batch, total_seq, kv_heads * head_dim]
+                src_key = op.Transpose(src_key, perm=[0, 2, 1, 3])
+                src_key = op.Reshape(src_key, [0, 0, -1])
+                src_value = op.Transpose(src_value, perm=[0, 2, 1, 3])
+                src_value = op.Reshape(src_value, [0, 0, -1])
+
+                attn_output, present_key, present_value = _apply_attention(
+                    op,
+                    query_states,
+                    src_key,
+                    src_value,
+                    attention_bias,
+                    past_key=None,
+                    past_value=None,
+                    num_attention_heads=self.num_attention_heads,
+                    num_key_value_heads=self.num_key_value_heads,
+                    scale=self.scaling,
+                    softcap=self.softcap,
+                )
         elif use_gqa:
             # GQA path: emit com.microsoft.GroupQueryAttention directly.
             # The op fuses RoPE + attention + KV cache into a single op,
@@ -1477,11 +1514,10 @@ def forward(
         )
 
         if use_gqa:
-            # Realize cos/sin caches as ONNX graph initializers.
-            # The returned gathered embeddings are saved for potential reuse
-            # by KV-shared layers that fall back to standard Attention.
-            local_pos_emb = self.rotary_emb_local(op, position_ids)
-            global_pos_emb = self.rotary_emb_global(op, position_ids)
+            # Realize cos/sin caches as ONNX graph initializers so that
+            # the GQA op can reference them.
+            self.rotary_emb_local(op, position_ids)
+            self.rotary_emb_global(op, position_ids)
 
             # seqlens_k[b] = sum(attention_mask[b]) - 1  (last valid KV idx)
             # total_seq_len = attention_mask.shape[1]     (past + current)
@@ -1525,18 +1561,11 @@ def forward(
                 "full_attention": self.rotary_emb_global(op, position_ids),
             }
 
-        # Fallback attention bias for non-GQA layers (KV-shared layers use
-        # this when use_gqa_for_kv_shared is False, plus all layers when
-        # use_gqa is False).
+        # Fallback attention bias for non-GQA layers (used when use_gqa is False).
         query_input = input_ids if input_ids is not None else hidden_states
         fallback_bias_dict: dict[str, ir.Value | None] = {}
-        has_kv_shared = any(layer.self_attn.is_kv_shared_layer for layer in self.layers)
-        need_fallback = not use_gqa or (has_kv_shared and not flags.use_gqa_for_kv_shared)
+        need_fallback = not use_gqa
         if need_fallback:
-            # All fallback layers use float additive bias masks encoding
-            # causal + sliding window + padding constraints. Float bias
-            # works with both unfused and MEA kernel paths on CUDA EP.
-            # Padding mask is required for batch > 1 correctness.
             fallback_bias_dict = {
                 "sliding_attention": create_attention_bias(
                     op,
@@ -1552,16 +1581,7 @@ def forward(
                     dtype=self._dtype,
                 ),
             }
-            # KV-shared layers also need position embeddings for the
-            # standard Attention path (manual RoPE). Reuse the embeddings
-            # already gathered when realizing cos/sin caches above.
-            if use_gqa:
-                fallback_pos_dict = {
-                    "sliding_attention": local_pos_emb,
-                    "full_attention": global_pos_emb,
-                }
-            else:
-                fallback_pos_dict = position_embeddings_dict
+            fallback_pos_dict = position_embeddings_dict
         else:
             fallback_pos_dict = {}
 
@@ -1587,12 +1607,9 @@ def forward(
         ):
             per_layer_input = per_layer_inputs[i] if per_layer_inputs is not None else None
 
-            # Per-layer decision: use GQA for non-shared layers when
-            # available, fall back to standard Attention for KV-shared
-            # layers (unless the use_gqa_for_kv_shared flag is set).
-            is_shared = layer.self_attn.is_kv_shared_layer
-            use_gqa_this_layer = use_gqa and (not is_shared or flags.use_gqa_for_kv_shared)
-            if use_gqa_this_layer:
+            # Per-layer decision: use GQA when available. KV-shared layers
+            # also use GQA (with empty K/V and shared past buffer).
+            if use_gqa:
                 attn_bias = gqa_ctx_dict[layer_type]
                 pos_emb = None
             else:

src/mobius/rewrite_rules/_skip_norm.py

@@ -64,6 +64,21 @@ def check(self, context, add_out, norm_out, **_):
         if producer is None or producer.op_type != "Add":
             return result.fail("Input to RMSNorm is not from an Add node")
 
+        # Both Add inputs must have the same rank (SkipSimplifiedLayerNormalization
+        # requires input and skip to have the same shape). A rank mismatch
+        # indicates broadcasting (e.g. Add(MatMul, bias) where bias is 1D).
+        input_a = producer.inputs[0]
+        input_b = producer.inputs[1]
+        shape_a = input_a.shape
+        shape_b = input_b.shape
+        rank_a = len(shape_a) if shape_a is not None else None
+        rank_b = len(shape_b) if shape_b is not None else None
+        if rank_a is not None and rank_b is not None and rank_a != rank_b:
+            return result.fail(
+                f"Add inputs have different ranks ({rank_a} vs {rank_b}); "
+                "SkipSimplifiedLayerNormalization requires same-shape inputs"
+            )
+
         # Don't fuse if add_out is itself a graph output — that indicates we're inside
         # an ONNX function body (e.g. SkipSimplifiedLayerNormalization_body) where
         # replace_all_uses_with would fail, or produce nested fusion.

src/mobius/tasks/_gemma4.py

@@ -311,6 +311,10 @@ def _build_vision(
             shape=[batch, num_patches, 2],
         )
 
+        # Cast float32 pixel_values from image processor to model dtype (e.g. fp16)
+        if config.dtype != ir.DataType.FLOAT:
+            pixel_values = op.Cast(pixel_values, to=config.dtype)
+
         image_features = vision(
             op,
             pixel_values=pixel_values,