diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 3ab13012..d3f9f1fe 100644 --- a/src/mobius/components/__init__.py +++ b/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, ) diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py index 1174df0c..73db9fe9 100644 --- a/src/mobius/components/_attention.py +++ b/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. diff --git a/src/mobius/components/_attention_test.py b/src/mobius/components/_attention_test.py index e24540a1..2acc138c 100644 --- a/src/mobius/components/_attention_test.py +++ b/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,218 @@ 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 + ) + # rotary_dim = head_dim / 2 = 16 / 2 = 8 (inv_freq has half the head_dim entries) + cos_cache = create_test_input(builder, "cos_cache", [32, 8]) + sin_cache = create_test_input(builder, "sin_cache", [32, 8]) + + 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, 8]) + sin_cache = create_test_input(builder, "sin_cache", [32, 8]) + + 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.""" + 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 + + def test_mrope_model_does_not_use_direct_gqa(self): + """Models with MRoPE (Qwen2.5-VL, Qwen3-VL) must NOT emit GQA directly. + + GQA do_rotary=1 only supports 1D RoPE. _MRopeBase subclasses + (ChunkedMRope, InterleavedMRope) use 3D position_ids for temporal/ + height/width axes. Silently emitting GQA would produce wrong outputs. + CUDA+f16 EP is used to trigger the GQA path for 1D-RoPE models; + the MRoPE model must fall through to the rewrite-rule path instead. + """ + from mobius._builder import build_from_module + from mobius._configs import ArchitectureConfig + from mobius._registry import registry + from mobius.rewrite_rules._testing_utils import count_ops + + # Minimal Qwen2.5-VL-style config with mrope_section (activates ChunkedMRope). + mrope_config = ArchitectureConfig( + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + num_hidden_layers=2, + vocab_size=256, + max_position_embeddings=128, + hidden_act="silu", + rms_norm_eps=1e-6, + rope_type="default", + rope_theta=10000.0, + pad_token_id=0, + dtype=ir.DataType.FLOAT16, + mrope_section=[4, 6, 6], # activates ChunkedMRope + ) + # Use "llama" model class (TextModel backbone) so use_gqa condition is evaluated. + # qwen2_5_vl uses a different model class, but the relevant guard is in TextModel. + pkg = build_from_module( + registry.get("llama")(mrope_config), + mrope_config, + execution_provider="cuda", + ) + ops = count_ops(pkg["model"]) + # MRoPE model must NOT use direct GQA (do_rotary=1 is 1D only). + # The rewrite rule path still applies GroupQueryAttention after graph construction. + assert ops.get("GroupQueryAttention", 0) == mrope_config.num_hidden_layers + + def test_direct_gqa_and_rewrite_rule_produce_same_structure(self): + """Direct GQA and rewrite-rule paths both produce GroupQueryAttention per layer. + + Verifies that for a standard 1D-RoPE model: + - CPU EP (direct path): num_layers GQA nodes + - Default EP + manual rewrite rule: same count + """ + from onnxscript.rewriter import rewrite + + from mobius._builder import build_from_module + from mobius._registry import registry + from mobius.rewrite_rules import group_query_attention_rules + from mobius.rewrite_rules._testing_utils import count_ops + + config = make_config( + max_position_embeddings=128, + rope_type="default", + rope_theta=10000.0, + ) + num_layers = config.num_hidden_layers + + # Direct path: CPU EP uses GQA directly (FLOAT is in cpu.gqa_dtypes) + pkg_direct = build_from_module( + registry.get("llama")(config), + config, + execution_provider="cpu", + ) + + # Rewrite-rule path: default EP keeps Attention + RoPE, then rewrite fires + pkg_default = build_from_module( + registry.get("llama")(config), + config, + execution_provider="default", + ) + rewrite(pkg_default["model"], group_query_attention_rules()) + + ops_direct = count_ops(pkg_direct["model"]) + ops_rewrite = count_ops(pkg_default["model"]) + + # Both paths must produce the same number of GroupQueryAttention nodes + assert ops_direct.get("GroupQueryAttention", 0) == num_layers + assert ops_rewrite.get("GroupQueryAttention", 0) == num_layers + # Neither path should leave any standard Attention nodes + assert ops_direct.get("Attention", 0) == 0 + assert ops_rewrite.get("Attention", 0) == 0 diff --git a/src/mobius/models/base.py b/src/mobius/models/base.py index 01f4bdcd..c427e539 100644 --- a/src/mobius/models/base.py +++ b/src/mobius/models/base.py @@ -13,12 +13,12 @@ from __future__ import annotations -from typing import TYPE_CHECKING - +import onnx_ir as ir import torch from onnxscript import nn from onnxscript._internal import builder +from mobius._build_context import ep_capabilities, get_build_dtype from mobius._configs import ArchitectureConfig, CausalLMConfig from mobius._weight_utils import ( preprocess_awq_weights, @@ -35,9 +35,8 @@ initialize_rope, make_quantized_linear_factory, ) - -if TYPE_CHECKING: - import onnx_ir as ir +from mobius.components._attention import GQAContext +from mobius.components._rotary_embedding import BaseRope, _MRopeBase class TextModel(nn.Module): @@ -83,21 +82,79 @@ def forward( hidden_states = inputs_embeds else: hidden_states = self.embed_tokens(op, input_ids) - position_embeddings = self.rotary_emb(op, position_ids) - - # When attention_mask is None (static cache mode), skip mask - # creation entirely — the Attention op uses is_causal=1 instead. - # When present, create a bool padding mask. Causal masking is - # handled by is_causal=1 on the Attention op (set in - # _apply_attention), so we only need padding information here. - if attention_mask is not None: - padding_mask = create_padding_mask( - op, - input_ids=hidden_states if input_ids is None else input_ids, - attention_mask=attention_mask, + + # Determine whether to emit GroupQueryAttention directly. + # Conditions: + # - attention_mask present: static-cache mode passes None; GQA requires seqlens_k. + # - EP gqa_dtypes: EP must declare GQA support for the build dtype (cuda/f16, + # cpu/f32, etc.). Default EP has gqa_dtypes={} so GQA is never emitted. + # - supports_fused_rope: EP must handle do_rotary=1 inside GQA. DML has + # gqa_dtypes={FLOAT16} but supports_fused_rope=False, so it uses the + # RotaryAttentionToGQA rewrite + SeparateRoPE path instead. + # - BaseRope (not _MRopeBase): standard 1D RoPE tables are required. + # _MRopeBase subclasses (ChunkedMRope for Qwen2.5-VL, InterleavedMRope for + # Qwen3-VL/Qwen3.5) use 3D position_ids; GQA do_rotary=1 only implements 1D + # RoPE, so those models must fall through to the RotaryAttentionToGQA rule. + caps = ep_capabilities() + dtype = get_build_dtype() + use_gqa = ( + attention_mask is not None + and dtype in caps.gqa_dtypes + and caps.supports_fused_rope + and isinstance(self.rotary_emb, BaseRope) + and not isinstance(self.rotary_emb, _MRopeBase) + ) + + if use_gqa: + # Call rotary_emb to realize cos_cache / sin_cache as ONNX graph + # initializers (onnxscript registers parameters on module __call__). + # The returned gathered embeddings are discarded — GroupQueryAttention + # will index the full tables itself via do_rotary=1. + self.rotary_emb(op, position_ids) + + # Build GQAContext from the cos/sin parameter tables and a + # seqlens_k / total_seq_len pair derived from attention_mask. + # Access cos_cache / sin_cache directly as ir.Value to avoid + # creating dead Gather(cos_cache, position_ids) nodes. + # + # seqlens_k[b] = sum(attention_mask[b]) - 1 = last valid KV index. + # total_seq_len = attention_mask.shape[1] = past + current len. + one_i32 = op.Constant(value_int=1) + seqlens_k = op.Cast( + op.Sub(op.ReduceSum(attention_mask, [1], keepdims=0), one_i32), + to=ir.DataType.INT32, + ) # [batch] INT32 + total_seq_len = op.Cast( + op.Gather(op.Shape(attention_mask), op.Constant(value_int=1)), + to=ir.DataType.INT32, + ) # scalar INT32 + + attention_bias: GQAContext | ir.Value | None = GQAContext( + seqlens_k=seqlens_k, + total_seq_len=total_seq_len, + cos_cache=self.rotary_emb.cos_cache, # [max_seq, rotary_dim] + sin_cache=self.rotary_emb.sin_cache, # [max_seq, rotary_dim] ) + # position_embeddings not needed: GroupQueryAttention handles RoPE + # internally via do_rotary=1. Passing None skips apply_rotary_pos_emb + # in Attention.forward() (which checks `if position_embeddings is not None`). + position_embeddings = None else: - padding_mask = None + position_embeddings = self.rotary_emb(op, position_ids) + + # When attention_mask is None (static cache mode), skip mask + # creation entirely — the Attention op uses is_causal=1 instead. + # When present, create a bool padding mask. Causal masking is + # handled by is_causal=1 on the Attention op (set in + # _apply_attention), so we only need padding information here. + if attention_mask is not None: + attention_bias = create_padding_mask( + op, + input_ids=hidden_states if input_ids is None else input_ids, + attention_mask=attention_mask, + ) + else: + attention_bias = None present_key_values = [] past_kvs = past_key_values or [None] * len(self.layers) @@ -105,7 +162,7 @@ def forward( hidden_states, present_kv = layer( op, hidden_states=hidden_states, - attention_bias=padding_mask, + attention_bias=attention_bias, position_embeddings=position_embeddings, past_key_value=past_kv, )