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

src/mobius/integrations/ort_genai/genai_config.py

@@ -60,7 +60,7 @@ def _default_search_params(*, ep: str, context_length: int) -> dict[str, Any]:
     from mobius._execution_providers import ep_registry
 
     caps = ep_registry.get(ep)
-    share_buffer = caps.supports_past_present_share_buffer if caps is not None else False
+    share_buffer = True
     if share_buffer:
         # EPs that pre-allocate KV-cache for the full max_length at load time
         # (e.g. WebGPU) need a capped default to avoid pre-allocating huge

src/mobius/models/gemma4.py

@@ -727,37 +727,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 layers should not receive GQAContext")
-            # 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,
@@ -1464,11 +1506,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)
@@ -1512,60 +1553,27 @@ def forward(
                 "full_attention": self.rotary_emb_global(op, position_ids),
             }
 
-        # Fallback attention bias for non-GQA layers (KV-shared layers always
-        # use this, 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] = {}
-        need_fallback = not use_gqa or any(
-            layer.self_attn.is_kv_shared_layer for layer in self.layers
-        )
+        need_fallback = not use_gqa
         if need_fallback:
-            if use_gqa:
-                # GQA is active for non-shared layers. KV-shared layers use
-                # the standard Attention op with is_causal=1, so we only need
-                # bool masks (not additive float bias). This avoids the
-                # CumSum/GreaterOrEqual chain used by create_attention_bias.
-                # Full-attention: simple padding mask (causality handled by op)
-                # Sliding-window: still needs CumSum for window constraint
-                fallback_bias_dict = {
-                    "sliding_attention": create_sliding_window_mask(
-                        op,
-                        input_ids=query_input,
-                        attention_mask=attention_mask,
-                        window_size=self.sliding_window or 512,
-                    ),
-                    "full_attention": create_padding_mask(
-                        op,
-                        input_ids=query_input,
-                        attention_mask=attention_mask,
-                    ),
-                }
-            else:
-                fallback_bias_dict = {
-                    "sliding_attention": create_attention_bias(
-                        op,
-                        input_ids=query_input,
-                        attention_mask=attention_mask,
-                        sliding_window=self.sliding_window,
-                        dtype=self._dtype,
-                    ),
-                    "full_attention": create_attention_bias(
-                        op,
-                        input_ids=query_input,
-                        attention_mask=attention_mask,
-                        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_bias_dict = {
+                "sliding_attention": create_attention_bias(
+                    op,
+                    input_ids=query_input,
+                    attention_mask=attention_mask,
+                    sliding_window=self.sliding_window,
+                    dtype=self._dtype,
+                ),
+                "full_attention": create_attention_bias(
+                    op,
+                    input_ids=query_input,
+                    attention_mask=attention_mask,
+                    dtype=self._dtype,
+                ),
+            }
+            fallback_pos_dict = position_embeddings_dict
         else:
             fallback_pos_dict = {}
 
@@ -1591,10 +1599,10 @@ 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.
+            # Per-layer decision: use GQA when available. KV-shared layers
+            # also use GQA (with empty K/V and shared past buffer).
             is_shared = layer.self_attn.is_kv_shared_layer
-            if use_gqa and not is_shared:
+            if use_gqa:
                 attn_bias = gqa_ctx_dict[layer_type]
                 pos_emb = None
             else:

src/mobius/rewrite_rules/_skip_norm.py

@@ -64,6 +64,19 @@ 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]
+        rank_a = input_a.shape.rank() if input_a.shape is not None else None
+        rank_b = input_b.shape.rank() if input_b.shape 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

@@ -300,7 +300,7 @@ def _build_vision(
 
         pixel_values = builder.input(
             "pixel_values",
-            dtype=config.dtype,
+            dtype=ir.DataType.FLOAT,
             shape=[batch, num_patches, pixel_dim],
         )
         pixel_position_ids = builder.input(
@@ -309,6 +309,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,