feat: emit GroupQueryAttention directly in Attention when EP supports it by justinchuby · Pull Request #134 · onnxruntime/mobius

justinchuby · 2026-04-09T05:23:58Z

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:

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.

Add GQAContext NamedTuple to _attention.py and a _forward_gqa() method to the Attention class. When attention_bias is a GQAContext, Attention emits com.microsoft::GroupQueryAttention directly (do_rotary=1) instead of the generic ONNX Attention + external RotaryEmbedding pair. TextModel.forward() detects EP support at build time via ep_capabilities() and get_build_dtype(). When the EP supports GQA for the build dtype (e.g. cuda/f16, cpu/f32) and the model uses a standard BaseRope, it builds a GQAContext from the rotary_emb.cos_cache / sin_cache parameters and seqlens_k / total_seq_len derived from attention_mask, then passes the GQAContext as attention_bias to all decoder layers. The RotaryAttentionToGQA rewrite rule is kept as a fallback for models with non-standard RoPE (Qwen3.5 mRoPE, ChatGLM, etc.) and for DML EP (supports_fused_rope=False). Key implementation details: - self.rotary_emb(op, position_ids) is still called in GQA mode to realize cos_cache / sin_cache as ONNX graph initializers; the returned gathered embeddings are discarded (dead nodes removed by optimization). - GQAContext exported from components.__init__ for use by external code. - softcap support: reads attn_logit_softcapping from config (Gemma2). - rotary_embedding_dim support: partial RoPE models respected. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

github-actions · 2026-04-09T05:25:48Z

Performance Comparison

Comparing c87787e → c5d20f6

Model	Metric	Baseline	Current	Delta
bert (feature-extraction)	model_size_bytes	359 KB	359 KB	+0.0%	⚪
bert (feature-extraction)	num_nodes	61	61	+0.0%	⚪
falcon	model_size_bytes	364 KB	364 KB	+0.0%	⚪
falcon	num_nodes	66	66	+0.0%	⚪
gemma2	model_size_bytes	428 KB	428 KB	+0.0%	⚪
gemma2	num_nodes	107	107	+0.0%	⚪
gpt2	model_size_bytes	388 KB	388 KB	+0.0%	⚪
gpt2	num_nodes	53	53	+0.0%	⚪
llama	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama	num_nodes	61	61	+0.0%	⚪
llama (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama (static-cache)	num_nodes	58	58	+0.0%	⚪
mamba (ssm-text-generation)	model_size_bytes	360 KB	360 KB	+0.0%	⚪
mamba (ssm-text-generation)	num_nodes	103	103	+0.0%	⚪
phi3	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3	num_nodes	61	61	+0.0%	⚪
phi3 (static-cache)	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3 (static-cache)	num_nodes	58	58	+0.0%	⚪
qwen2	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2	num_nodes	61	61	+0.0%	⚪
qwen2 (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2 (static-cache)	num_nodes	58	58	+0.0%	⚪
qwen3_5_moe (hybrid-text-generation)	model_size_bytes	506 KB	506 KB	+0.0%	⚪
qwen3_5_moe (hybrid-text-generation)	num_nodes	275	275	+0.0%	⚪
qwen3_5_text (hybrid-text-generation)	model_size_bytes	458 KB	458 KB	+0.0%	⚪
qwen3_5_text (hybrid-text-generation)	num_nodes	129	129	+0.0%	⚪
qwen3_5_vl (hybrid-qwen-vl)	model_size_bytes	977 KB	977 KB	+0.0%	⚪
qwen3_5_vl (hybrid-qwen-vl)	num_nodes	409	409	+0.0%	⚪
t5 (seq2seq)	model_size_bytes	836 KB	836 KB	+0.0%	⚪
t5 (seq2seq)	num_nodes	174	174	+0.0%	⚪
whisper (speech-to-text)	model_size_bytes	1008 KB	1008 KB	+0.0%	⚪
whisper (speech-to-text)	num_nodes	128	128	+0.0%	⚪

No performance regressions.

github-actions · 2026-04-09T05:27:44Z

🏗️ Architecture Diff

Comparing c87787e → c5d20f6

Model	Sub-model	Changes	Status
bert (feature-extraction)	model	0	⚪
falcon	model	2	🔵
gemma2	model	0	⚪
gpt2	model	2	🔵
llama	model	2	🔵
llama (static-cache)	model	2	🔵
mamba (ssm-text-generation)	model	0	⚪
phi3	model	2	🔵
phi3 (static-cache)	model	2	🔵
qwen	model	2	🔵
qwen (static-cache)	model	2	🔵
qwen2	model	2	🔵
qwen2 (static-cache)	model	2	🔵
qwen2_moe	model	2	🔵
qwen2_moe (static-cache)	model	2	🔵
qwen3	model	2	🔵
qwen3 (static-cache)	model	2	🔵
qwen3_5_moe (hybrid-text-generation)	model	1	🔵
qwen3_5_text (hybrid-text-generation)	model	1	🔵
qwen3_5_vl (hybrid-qwen-vl)	decoder	1	🔵
qwen3_5_vl (hybrid-qwen-vl)	embedding	0	⚪
qwen3_5_vl (hybrid-qwen-vl)	vision	0	⚪
qwen3_moe	model	2	🔵
qwen3_moe (static-cache)	model	2	🔵
qwen3_next (hybrid-text-generation)	model	1	🔵
t5 (seq2seq)	decoder	0	⚪
t5 (seq2seq)	encoder	0	⚪
whisper (speech-to-text)	decoder	0	⚪
whisper (speech-to-text)	encoder	0	⚪

falcon / model — 2 change(s)

Op summary: 66 → 66 nodes

No op-sequence changes.

Modified attributes:

node[29] Attention: softcap: None → 0.0
node[52] Attention: softcap: None → 0.0

gpt2 / model — 2 change(s)

Op summary: 53 → 53 nodes

No op-sequence changes.

Modified attributes:

node[17] Attention: softcap: None → 0.0
node[37] Attention: softcap: None → 0.0

llama / model — 2 change(s)

Op summary: 61 → 61 nodes

No op-sequence changes.

Modified attributes:

node[19] Attention: softcap: None → 0.0
node[43] Attention: softcap: None → 0.0

llama (static-cache) / model — 2 change(s)

Op summary: 58 → 58 nodes

No op-sequence changes.

Modified attributes:

node[14] Attention: softcap: None → 0.0
node[40] Attention: softcap: None → 0.0

phi3 / model — 2 change(s)

Op summary: 61 → 61 nodes

No op-sequence changes.

Modified attributes:

node[19] Attention: softcap: None → 0.0
node[43] Attention: softcap: None → 0.0

phi3 (static-cache) / model — 2 change(s)

Op summary: 58 → 58 nodes

No op-sequence changes.

Modified attributes:

node[14] Attention: softcap: None → 0.0
node[40] Attention: softcap: None → 0.0

qwen / model — 2 change(s)

Op summary: 61 → 61 nodes

No op-sequence changes.

Modified attributes:

node[19] Attention: softcap: None → 0.0
node[43] Attention: softcap: None → 0.0

qwen (static-cache) / model — 2 change(s)

Op summary: 58 → 58 nodes

No op-sequence changes.

Modified attributes:

node[14] Attention: softcap: None → 0.0
node[40] Attention: softcap: None → 0.0

qwen2 / model — 2 change(s)

Op summary: 61 → 61 nodes

No op-sequence changes.

Modified attributes:

node[19] Attention: softcap: None → 0.0
node[43] Attention: softcap: None → 0.0

qwen2 (static-cache) / model — 2 change(s)

Op summary: 58 → 58 nodes

No op-sequence changes.

Modified attributes:

node[14] Attention: softcap: None → 0.0
node[40] Attention: softcap: None → 0.0

qwen2_moe / model — 2 change(s)

Op summary: 224 → 224 nodes

No op-sequence changes.

Modified attributes:

node[29] Attention: softcap: None → 0.0
node[131] Attention: softcap: None → 0.0

qwen2_moe (static-cache) / model — 2 change(s)

Op summary: 214 → 214 nodes

No op-sequence changes.

Modified attributes:

node[17] Attention: softcap: None → 0.0
node[121] Attention: softcap: None → 0.0

qwen3 / model — 2 change(s)

Op summary: 73 → 73 nodes

No op-sequence changes.

Modified attributes:

node[25] Attention: softcap: None → 0.0
node[55] Attention: softcap: None → 0.0

qwen3 (static-cache) / model — 2 change(s)

Op summary: 70 → 70 nodes

No op-sequence changes.

Modified attributes:

node[20] Attention: softcap: None → 0.0
node[52] Attention: softcap: None → 0.0

qwen3_5_moe (hybrid-text-generation) / model — 1 change(s)

Op summary: 275 → 275 nodes

No op-sequence changes.

Modified attributes:

node[180] Attention: softcap: None → 0.0

qwen3_5_text (hybrid-text-generation) / model — 1 change(s)

Op summary: 129 → 129 nodes

No op-sequence changes.

Modified attributes:

node[107] Attention: softcap: None → 0.0

qwen3_5_vl (hybrid-qwen-vl) / decoder — 1 change(s)

Op summary: 149 → 149 nodes

No op-sequence changes.

Modified attributes:

node[127] Attention: softcap: None → 0.0

qwen3_moe / model — 2 change(s)

Op summary: 202 → 202 nodes

No op-sequence changes.

Modified attributes:

node[32] Attention: softcap: None → 0.0
node[123] Attention: softcap: None → 0.0

qwen3_moe (static-cache) / model — 2 change(s)

Op summary: 192 → 192 nodes

No op-sequence changes.

Modified attributes:

node[20] Attention: softcap: None → 0.0
node[113] Attention: softcap: None → 0.0

qwen3_next (hybrid-text-generation) / model — 1 change(s)

Op summary: 585 → 585 nodes

No op-sequence changes.

Modified attributes:

node[484] Attention: softcap: None → 0.0

Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed)

Copilot

Pull request overview

This PR updates the core transformer graph-construction path to emit com.microsoft::GroupQueryAttention directly (when the active execution provider supports it), reducing reliance on the post-hoc RotaryAttentionToGQA rewrite rule while keeping it as a fallback.

Changes:

Add a GQAContext bundle and dispatch logic so Attention.forward() can emit GroupQueryAttention directly when given that context.
Add EP-capability-based gating in TextModel.forward() to construct and thread GQAContext (including RoPE cache initializers and sequence-length scalars).
Add component- and build-level tests covering the new dispatch behavior and EP-dependent selection.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
`src/mobius/models/base.py`	Adds EP-driven dispatch to build and pass `GQAContext` through the decoder stack for direct GQA emission.
`src/mobius/components/_attention.py`	Introduces `GQAContext` and implements the direct `GroupQueryAttention` emission path in `Attention`.
`src/mobius/components/_attention_test.py`	Adds tests validating GQAContext dispatch and EP-dependent graph construction results.
`src/mobius/components/__init__.py`	Exports `GQAContext` in the public components API.

1. MRoPE guard: add 'not isinstance(self.rotary_emb, _MRopeBase)' to the use_gqa condition in TextModel.forward(). ChunkedMRope (Qwen2.5-VL) and InterleavedMRope (Qwen3-VL) both inherit from _MRopeBase and use 3D position_ids; GroupQueryAttention do_rotary=1 only implements 1D RoPE, so those models would silently produce wrong outputs on CUDA+float16. They fall through to the RotaryAttentionToGQA rewrite rule instead. 2. softcap forwarded: add softcap parameter to _apply_attention() and pass it to both the static-cache and dynamic-cache op.Attention() calls. Gemma2 on DML/WebGPU/default EPs was silently losing attn_logit_softcap. 3. Test cos_cache shape: fix [32,16] → [32,8]. rotary_dim = head_dim/2 = 8 since inv_freq = arange(0, head_dim, 2) / head_dim has head_dim/2 entries. 4. base.py comments: expand use_gqa condition comment to explain each guard (attention_mask, gqa_dtypes, supports_fused_rope, BaseRope/_MRopeBase). Add comment explaining why position_embeddings=None in GQA mode. 5. New tests: - test_mrope_model_does_not_use_direct_gqa: ChunkedMRope model on CUDA+f16 still goes through the rewrite rule (not direct GQA). - test_direct_gqa_and_rewrite_rule_produce_same_structure: verifies that CPU-EP direct path and default-EP+rewrite-rule path both produce the same GQA node count with no leftover Attention nodes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

justinchuby · 2026-04-09T17:14:25Z

+    :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).


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

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:

Gathering cos/sin for each of the three dimensions separately

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).

@kunal-vaishnavi does this sound reasonable?

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

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:

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).

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.

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.

Let's do it in a separate PR.

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

Document the proven pattern from Gemma 4 implementation (PR #134 GQA pattern applied to MoE): EP capability check, full pre-topk router_probs requirement, CastLike after MoE output to preserve dtype, preprocess_weights stacking, and loop-fallback when EP doesn't support fused MoE. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

justinchuby requested review from a team and Copilot April 9, 2026 05:24

Copilot started reviewing on behalf of justinchuby April 9, 2026 05:24 View session

Copilot AI reviewed Apr 9, 2026

View reviewed changes

Comment thread src/mobius/models/base.py

Comment thread src/mobius/components/_attention.py Outdated

justinchuby commented Apr 9, 2026

View reviewed changes

Copilot started work on behalf of justinchuby April 9, 2026 17:14 View session

Copilot finished work on behalf of justinchuby April 9, 2026 17:17

Copilot started work on behalf of justinchuby April 9, 2026 17:25 View session

Copilot finished work on behalf of justinchuby April 9, 2026 17:29

Update src/mobius/components/_attention.py

c5d20f6

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

justinchuby merged commit c448d4f into main Apr 9, 2026
15 of 17 checks passed

justinchuby deleted the justinchu/gqa-direct branch April 9, 2026 20:34

This was referenced Apr 10, 2026

[DRAFT] Add native GQA attention support for ORT GenAI compatibility #109

Closed

Emit GQA for 3d rope cases #147

Closed

justinchuby mentioned this pull request Apr 25, 2026

Notable changes: April 9–25, 2026 #220

Closed

Navigation Menu

Uh oh!

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

feat: emit GroupQueryAttention directly in Attention when EP supports it#134
justinchuby merged 3 commits into
mainfrom
justinchu/gqa-direct

justinchuby commented Apr 9, 2026

Uh oh!

github-actions Bot commented Apr 9, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Apr 9, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

justinchuby Apr 9, 2026

Uh oh!

Copilot AI Apr 9, 2026

Uh oh!

justinchuby Apr 9, 2026

Uh oh!

justinchuby Apr 9, 2026

Uh oh!

Copilot AI Apr 9, 2026

Uh oh!

justinchuby Apr 9, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

justinchuby commented Apr 9, 2026

Summary

What changed

New: GQAContext NamedTuple (components/_attention.py)

Modified: Attention.forward()

Modified: TextModel.forward() (models/base.py)

What stays unchanged

Tests

Uh oh!

github-actions Bot commented Apr 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Performance Comparison

Uh oh!

github-actions Bot commented Apr 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🏗️ Architecture Diff

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

justinchuby Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

justinchuby Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

justinchuby Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

justinchuby Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

New: `GQAContext` NamedTuple (`components/_attention.py`)

Modified: `Attention.forward()`

Modified: `TextModel.forward()` (`models/base.py`)

github-actions Bot commented Apr 9, 2026 •

edited

Loading

github-actions Bot commented Apr 9, 2026 •

edited

Loading