feat: emit GroupQueryAttention directly in Attention when EP supports it (#134)

## Summary

Eliminates the post-hoc `RotaryAttentionToGQA` rewrite rule as the
primary path for GQA-capable EPs, replacing it with direct
`com.microsoft::GroupQueryAttention` emission at graph construction
time.

## What changed

### New: `GQAContext` NamedTuple (`components/_attention.py`)
A typed bundle of per-graph scalars passed through the decoder stack as
`attention_bias`:
- `seqlens_k`: per-batch KV sequence length `[batch]` INT32
- `total_seq_len`: scalar INT32
- `cos_cache` / `sin_cache`: full RoPE tables (not gathered slices)

### Modified: `Attention.forward()`
Checks `isinstance(attention_bias, GQAContext)` at the top and
dispatches to `_forward_gqa()` which emits `GroupQueryAttention`
directly with `do_rotary=1`, bypassing the external
`RotaryEmbeddingBase.forward()` + `apply_rotary_pos_emb()` path.

### Modified: `TextModel.forward()` (`models/base.py`)
Adds an EP-driven dispatch at graph construction time:
```python
caps = ep_capabilities()
use_gqa = (
    attention_mask is not None
    and get_build_dtype() in caps.gqa_dtypes
    and caps.supports_fused_rope
    and isinstance(self.rotary_emb, BaseRope)
)
```
When True: calls `self.rotary_emb(op, position_ids)` to realize
`cos_cache`/`sin_cache` as ONNX initializers (discards result), then
builds `GQAContext` with the raw parameter tensors and computes
`seqlens_k`/`total_seq_len` from `attention_mask`.

## What stays unchanged

- `RotaryAttentionToGQA` rewrite rule — kept as fallback for:
- Qwen3.5 with `_MRopeBase` 3D mRoPE (`isinstance(..., BaseRope)` is
True but `supports_fused_rope=False` on affected EPs excludes it)
  - DML EP (`supports_fused_rope=False`)  
  - Any model not using `TextModel` (VLMs, Qwen3.5, etc.)
- Graph I/O: `position_ids` remains in the graph inputs for consistency
with the rewrite-rule path (also leaves `position_ids` as a dead input
post-optimization in both paths)

## Tests

Five new tests in
`components/_attention_test.py::TestGQAContextDispatch`:
- `test_gqa_context_emits_group_query_attention`: component-level,
verifies GQA is emitted
- `test_gqa_context_respects_rotary_interleaved`: checks
`rotary_interleaved` attribute
- `test_standard_attention_when_no_gqa_context`: standard path
unaffected
- `test_build_with_cuda_ep_emits_gqa_directly`: CUDA EP + f16 → GQA
present, Attention absent
- `test_build_with_default_ep_uses_standard_attention`: default EP →
standard Attention

All 2327 tests pass.

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

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,39 @@
 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 last valid KV index ``[batch]`` INT32.
+            Computed as ``ReduceSum(attention_mask, axis=1) - 1``; this is a
+            0-based index into the valid KV tokens, not the KV length itself.
+        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):
@@ -51,6 +83,7 @@ def _apply_attention(
     num_attention_heads: int,
     num_key_value_heads: int,
     scale: float,
+    softcap: float = 0.0,
     static_cache: StaticCacheState | None = None,
 ) -> tuple[ir.Value, ir.Value, ir.Value]:
     """Apply the ONNX Attention op with internal or static KV cache.
@@ -125,6 +158,7 @@ def _apply_attention(
             q_num_heads=num_attention_heads,
             kv_num_heads=num_key_value_heads,
             scale=scale,
+            softcap=softcap,
             is_causal=1,
             _outputs=3,
         )
@@ -144,6 +178,7 @@ def _apply_attention(
         q_num_heads=num_attention_heads,
         kv_num_heads=num_key_value_heads,
         scale=scale,
+        softcap=softcap,
         is_causal=1,
         _outputs=3,
     )
@@ -187,6 +222,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 +268,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 +291,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(
@@ -284,12 +332,70 @@ def forward(
             num_attention_heads=self.num_attention_heads,
             num_key_value_heads=self.num_key_value_heads,
             scale=self.scaling,
+            softcap=self._softcap,
             static_cache=static_cache,
         )
 
         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.