Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
justinchuby merged 3 commits into from
Apr 9, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6220588
feat: emit GroupQueryAttention directly in Attention when EP supports it
justinchuby Apr 9, 2026
aa2b0af
fix: guard GQA direct path against MRoPE, forward softcap, add tests
justinchuby Apr 9, 2026
c5d20f6
Update src/mobius/components/_attention.py
justinchuby Apr 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/mobius/components/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
"GatedRMSNorm",
"Gemma3MultiModalProjector",
"GroupNorm",
"GQAContext",
"INT64_MAX",
"InputMixer",
"JambaSelectiveScan",
Expand Down Expand Up @@ -109,6 +110,7 @@
from mobius.components._activations import SiLU, get_activation
from mobius.components._attention import (
Attention,
GQAContext,
Qwen35Attention,
StaticCacheState,
)
Expand Down
114 changes: 110 additions & 4 deletions src/mobius/components/_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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).

@justinchuby justinchuby Apr 9, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need to rewrite for 3D mRoPE? What's the difference? @copilot

Copilot AI Apr 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The GroupQueryAttention kernel's do_rotary=1 mode expects a simple 1D cos/sin lookup table [max_seq, rotary_dim] and a scalar position per token — it gathers one vector per position from those tables internally.

3D mRoPE (_MRopeBase subclasses: ChunkedMRope for Qwen2.5-VL, InterleavedMRope for Qwen3-VL) uses position_ids of shape (3, batch, seq_len) with three separate coordinates per token (temporal T, height H, width W). The RoPE is computed by:

  1. Gathering cos/sin for each of the three dimensions separately
  2. Blending them per-channel using Where nodes controlled by h_mask/w_mask arrays

The resulting per-token, per-channel cos/sin embeddings can't be expressed as a single lookup table that GQA's do_rotary=1 understands — the "table" is dynamically mixed rather than a simple positional index.

So for 3D mRoPE the direct GQA path is skipped, and the RotaryAttentionToGQA / AttentionToGQA rewrite rules handle these models instead, using do_rotary=0 (Q/K already have RoPE pre-applied before entering GQA).

@justinchuby justinchuby Apr 9, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@kunal-vaishnavi does this sound reasonable?

@justinchuby justinchuby Apr 9, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@copilot would it be a good idea to just set do_rotary=0 here? What's missing?

Copilot AI Apr 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, do_rotary=0 could work for 3D mRoPE at build time — the AttentionToGQA rewrite rule already does exactly this today. Making it work directly in the build path is feasible but requires a few pieces that are currently missing:

  1. Different payload in context: GQAContext currently carries cos_cache/sin_cache (full tables [max_seq, rotary_dim]) specifically for do_rotary=1 — GQA indexes them internally. For do_rotary=0, we need the already-gathered, per-token (cos, sin) tensors (batch, seq, rotary_dim) produced by _MRopeBase.forward(), not raw tables. These are a completely different shape/purpose and would need to be plumbed through (either by extending GQAContext or creating a separate context type).

  2. Apply RoPE inside _forward_gqa(): When do_rotary=0, apply_rotary_pos_emb() must be called on Q/K before emitting the GroupQueryAttention node (no cos/sin inputs to GQA in this mode). This branch doesn't exist in _forward_gqa() yet.

  3. Remove the _MRopeBase guard in TextModel.forward(), or replace it with a condition that chooses do_rotary=0 for _MRopeBase vs do_rotary=1 for standard BaseRope.

Worth noting: the supports_fused_rope guard would also need revisiting — do_rotary=0 doesn't need fused RoPE inside the GQA kernel, so it could be applied even on EPs where supports_fused_rope=False (as long as dtype in caps.gqa_dtypes).

Happy to add this if you'd like to extend the PR. It would fully eliminate AttentionToGQA as a fallback for mRoPE models on GQA-capable EPs.

@justinchuby justinchuby Apr 9, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's do it in a separate PR.


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):
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
)
Expand All @@ -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,
)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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.
Expand Down
Loading
Loading