feat: emit GroupQueryAttention directly in Attention when EP supports it

src/mobius/components/__init__.py

@@ -34,6 +34,7 @@
     "GatedRMSNorm",
     "Gemma3MultiModalProjector",
     "GroupNorm",
+    "GQAContext",
     "INT64_MAX",
     "InputMixer",
     "JambaSelectiveScan",
@@ -109,6 +110,7 @@
 from mobius.components._activations import SiLU, get_activation
 from mobius.components._attention import (
     Attention,
+    GQAContext,
     Qwen35Attention,
     StaticCacheState,
 )

src/mobius/components/_attention.py

@@ -4,8 +4,9 @@
 from __future__ import annotations
 
 import math
-from typing import TYPE_CHECKING, NamedTuple
+from typing import NamedTuple
 
+import onnx_ir as ir
 from onnxscript import nn
 from onnxscript._internal import builder
 
@@ -14,8 +15,38 @@
 from mobius.components._rms_norm import OffsetRMSNorm, RMSNorm
 from mobius.components._rotary_embedding import apply_rotary_pos_emb
 
-if TYPE_CHECKING:
-    import onnx_ir as ir
+
+class GQAContext(NamedTuple):
+    """Context for direct ``com.microsoft::GroupQueryAttention`` emission.
+
+    Created once per graph by :class:`~mobius.models.base.TextModel` when the
+    active EP (from :func:`~mobius._build_context.ep_capabilities`) supports
+    GQA for the current build dtype.  Passed through DecoderLayer as the
+    ``attention_bias`` argument so that :class:`Attention` can detect it and
+    emit ``GroupQueryAttention`` directly instead of the generic
+    ``Attention + RotaryEmbedding`` sequence.
+
+    Using this context skips the post-hoc
+    :class:`~mobius.rewrite_rules._group_query_attention.RotaryAttentionToGQA`
+    rewrite rule for models that use the standard :class:`TextModel` backbone.
+    The rewrite rule remains as a fallback for models with non-standard RoPE
+    (e.g. Qwen3.5 with 3D mRoPE).
+
+    Fields:
+        seqlens_k: Per-batch actual KV sequence length ``[batch]`` INT32.
+            Computed as ``ReduceSum(attention_mask, axis=1) - 1``.
+        total_seq_len: Scalar total sequence length INT32.
+            Computed as ``Shape(attention_mask)[1]``.
+        cos_cache: Full cosine RoPE table ``[max_seq_len, rotary_dim]`` FLOAT.
+            Taken directly from the model's ``rotary_emb.cos_cache`` parameter.
+        sin_cache: Full sine RoPE table ``[max_seq_len, rotary_dim]`` FLOAT.
+            Taken directly from the model's ``rotary_emb.sin_cache`` parameter.
+    """
+
+    seqlens_k: ir.Value
+    total_seq_len: ir.Value
+    cos_cache: ir.Value
+    sin_cache: ir.Value
 
 
 class StaticCacheState(NamedTuple):
@@ -187,6 +218,8 @@ def __init__(
             else int(self.head_dim * config.partial_rotary_factor)
         )
         self._rope_interleave = config.rope_interleave
+        # Gemma2-style logit soft-capping; 0.0 means disabled.
+        self._softcap = getattr(config, "attn_logit_softcapping", 0.0) or 0.0
 
         self.q_proj = linear_class(
             self.hidden_size,
@@ -231,7 +264,7 @@ def forward(
         self,
         op: builder.OpBuilder,
         hidden_states: ir.Value,
-        attention_bias: ir.Value | None,
+        attention_bias: ir.Value | GQAContext | None,
         position_embeddings: tuple | None = None,
         past_key_value: tuple | None = None,
         static_cache: StaticCacheState | None = None,
@@ -254,6 +287,17 @@ def forward(
                 query_states = op.Reshape(query_states, [0, 0, -1])
                 key_states = op.Reshape(key_states, [0, 0, -1])
 
+        # Direct GroupQueryAttention path: skip external RoPE, fuse everything.
+        if isinstance(attention_bias, GQAContext):
+            return self._forward_gqa(
+                op,
+                query_states,
+                key_states,
+                value_states,
+                attention_bias,
+                past_key_value,
+            )
+
         # Apply rotary position embeddings (skip when not provided)
         if position_embeddings is not None:
             query_states = apply_rotary_pos_emb(
@@ -290,6 +334,63 @@ def forward(
         attn_output = self.o_proj(op, attn_output)
         return attn_output, (present_key, present_value)
 
+    def _forward_gqa(
+        self,
+        op: builder.OpBuilder,
+        query_states: ir.Value,
+        key_states: ir.Value,
+        value_states: ir.Value,
+        gqa_ctx: GQAContext,
+        past_key_value: tuple | None,
+    ):
+        """Emit ``com.microsoft::GroupQueryAttention`` directly.
+
+        Called from :meth:`forward` when ``attention_bias`` is a
+        :class:`GQAContext`.  Bypasses the external
+        :class:`~mobius.components._rotary_embedding.RotaryEmbeddingBase`
+        forward pass and the post-hoc
+        :class:`~mobius.rewrite_rules._group_query_attention.RotaryAttentionToGQA`
+        rewrite rule; RoPE is handled by the ``do_rotary=1`` attribute instead.
+
+        Returns ``(attn_output, (present_key, present_value))`` in the same
+        shape as the standard :meth:`forward` path.
+        """
+        past_key = past_key_value[0] if past_key_value is not None else None
+        past_value = past_key_value[1] if past_key_value is not None else None
+
+        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:
+            # Partial RoPE: only rotate the first rotary_embedding_dim elements.
+            gqa_attrs["rotary_embedding_dim"] = self.rotary_embedding_dim
+
+        # Emit GroupQueryAttention: RoPE + attention + KV cache in one op.
+        # Outputs: (attn_output [B, S, hidden], present_key, present_value)
+        attn_out, present_key, present_value = op.GroupQueryAttention(
+            query_states,  # [B, S, num_heads * head_dim]
+            key_states,  # [B, S, kv_heads * head_dim]
+            value_states,  # [B, S, kv_heads * head_dim]
+            past_key,  # [B, kv_heads, past_S, head_dim] or None
+            past_value,  # [B, kv_heads, past_S, head_dim] or None
+            gqa_ctx.seqlens_k,  # [B] INT32
+            gqa_ctx.total_seq_len,  # scalar INT32
+            gqa_ctx.cos_cache,  # [max_seq, rotary_dim]
+            gqa_ctx.sin_cache,  # [max_seq, rotary_dim]
+            _domain="com.microsoft",
+            _outputs=3,
+            **gqa_attrs,
+        )
+
+        attn_out = self.o_proj(op, attn_out)
+        return attn_out, (present_key, present_value)
+
 
 class Qwen35Attention(nn.Module):
     """Multi-head attention with output gating for Qwen3.5.

src/mobius/components/_attention_test.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import onnx_ir as ir
 import pytest
 
 from mobius._testing import (
@@ -186,3 +187,129 @@ def test_forward_builds_graph(self):
         # Should have Attention op + Sigmoid (for gate) + Mul (output gating)
         assert count_op_type(graph, "Attention") >= 1
         assert count_op_type(graph, "Sigmoid") >= 1
+
+
+class TestGQAContextDispatch:
+    """Tests for the GQAContext direct GroupQueryAttention emission path."""
+
+    def test_gqa_context_emits_group_query_attention(self):
+        """When attention_bias is a GQAContext, Attention emits GroupQueryAttention directly."""
+        from mobius.components._attention import GQAContext
+
+        config = make_config()
+        attn = Attention(config)
+        builder, op, graph = create_test_builder()
+
+        hidden = create_test_input(builder, "hidden", [1, 8, 64])
+        past_key = create_test_input(builder, "past_key", [1, 2, 4, 16])
+        past_value = create_test_input(builder, "past_value", [1, 2, 4, 16])
+        seqlens_k = create_test_input(builder, "seqlens_k", [1], dtype=ir.DataType.INT32)
+        total_seq_len = create_test_input(
+            builder, "total_seq_len", [], dtype=ir.DataType.INT32
+        )
+        cos_cache = create_test_input(builder, "cos_cache", [32, 16])
+        sin_cache = create_test_input(builder, "sin_cache", [32, 16])
+
+        gqa_ctx = GQAContext(
+            seqlens_k=seqlens_k,
+            total_seq_len=total_seq_len,
+            cos_cache=cos_cache,
+            sin_cache=sin_cache,
+        )
+
+        output, (pk, pv) = attn(
+            op, hidden, attention_bias=gqa_ctx, past_key_value=(past_key, past_value)
+        )
+        builder._adapt_outputs([output, pk, pv])
+
+        # Direct path: GroupQueryAttention instead of ONNX Attention
+        assert count_op_type(graph, "GroupQueryAttention") >= 1
+        assert count_op_type(graph, "Attention") == 0
+
+    def test_gqa_context_respects_rotary_interleaved(self):
+        """rotary_interleaved attribute is set from config.rope_interleave."""
+        from mobius.components._attention import GQAContext
+
+        config = make_config(rope_interleave=True)
+        attn = Attention(config)
+        builder, op, graph = create_test_builder()
+
+        hidden = create_test_input(builder, "hidden", [1, 8, 64])
+        past_key = create_test_input(builder, "past_key", [1, 2, 4, 16])
+        past_value = create_test_input(builder, "past_value", [1, 2, 4, 16])
+        seqlens_k = create_test_input(builder, "seqlens_k", [1], dtype=ir.DataType.INT32)
+        total_seq_len = create_test_input(
+            builder, "total_seq_len", [], dtype=ir.DataType.INT32
+        )
+        cos_cache = create_test_input(builder, "cos_cache", [32, 16])
+        sin_cache = create_test_input(builder, "sin_cache", [32, 16])
+
+        gqa_ctx = GQAContext(seqlens_k, total_seq_len, cos_cache, sin_cache)
+
+        output, _ = attn(
+            op, hidden, attention_bias=gqa_ctx, past_key_value=(past_key, past_value)
+        )
+        builder._adapt_outputs([output])
+
+        gqa_node = next(n for n in graph if n.op_type == "GroupQueryAttention")
+        assert gqa_node.attributes["rotary_interleaved"].value == 1
+
+    def test_standard_attention_when_no_gqa_context(self):
+        """Without GQAContext, standard ONNX Attention is emitted."""
+        config = make_config()
+        attn = Attention(config)
+        builder, op, graph = create_test_builder()
+
+        hidden = create_test_input(builder, "hidden", [1, 8, 64])
+        bias = create_test_input(builder, "bias", [1, 4, 8, 8])
+
+        output, _ = attn(op, hidden, attention_bias=bias)
+        builder._adapt_outputs([output])
+
+        assert count_op_type(graph, "Attention") >= 1
+        assert count_op_type(graph, "GroupQueryAttention") == 0
+
+    def test_build_with_cuda_ep_emits_gqa_directly(self):
+        """build_from_module with CUDA EP and float16 config emits GroupQueryAttention directly."""
+        import onnx_ir as ir
+
+        from mobius._builder import build_from_module
+        from mobius._registry import registry
+        from mobius.rewrite_rules._testing_utils import count_ops
+
+        config = make_config(
+            dtype=ir.DataType.FLOAT16,
+            max_position_embeddings=128,
+            rope_type="default",
+            rope_theta=10000.0,
+        )
+        pkg = build_from_module(
+            registry.get("llama")(config),
+            config,
+            execution_provider="cuda",
+        )
+        ops = count_ops(pkg["model"])
+        # Direct generation: each layer should have a GroupQueryAttention node
+        assert ops.get("GroupQueryAttention", 0) == config.num_hidden_layers
+        # Standard ONNX Attention should not appear
+        assert ops.get("Attention", 0) == 0
+
+    def test_build_with_default_ep_uses_standard_attention(self):
+        """build_from_module with default EP keeps standard ONNX Attention (no GQA)."""
+        from mobius._builder import build_from_module
+        from mobius._registry import registry
+        from mobius.rewrite_rules._testing_utils import count_ops
+
+        config = make_config(
+            max_position_embeddings=128,
+            rope_type="default",
+            rope_theta=10000.0,
+        )
+        pkg = build_from_module(
+            registry.get("llama")(config),
+            config,
+            execution_provider="default",
+        )
+        ops = count_ops(pkg["model"])
+        assert ops.get("GroupQueryAttention", 0) == 0
+        assert ops.get("Attention", 0) == config.num_hidden_layers