diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index f5f07eab..b3eb4bb5 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -139,10 +139,26 @@ def _cmd_build(args: argparse.Namespace) -> None: "Remove --task to use --static-cache." ) + # Validate --gqa + --static-cache are mutually exclusive + if args.gqa and args.static_cache: + raise SystemExit("Error: --gqa and --static-cache are mutually exclusive.") + + # Validate --gqa + --task compatibility + if args.gqa and args.task is not None: + raise SystemExit( + "Error: --gqa cannot be combined with --task. Remove --task to use --gqa." + ) + load_weights = not args.no_weights task: str | ModelTask | None = args.task + module_kwargs: dict | None = None if args.static_cache: task = CausalLMTask(static_cache=True, max_seq_len=args.max_seq_len) + elif args.gqa: + from mobius.components import GQAAttention + + task = CausalLMTask(gqa=True) + module_kwargs = {"attention_class": GQAAttention} trust_remote_code = args.trust_remote_code output_dir = args.output_dir os.makedirs(output_dir, exist_ok=True) @@ -183,7 +199,7 @@ def _cmd_build(args: argparse.Namespace) -> None: if task is None: task = _default_task_for_model(model_type) module_class = registry.get(model_type) - model_module = module_class(config) + model_module = module_class(config, **(module_kwargs or {})) pkg = build_from_module(model_module, config, task=task) for name, model in pkg.items(): model.graph.name = f"{config_path}/{name}" @@ -199,6 +215,7 @@ def _cmd_build(args: argparse.Namespace) -> None: dtype=dtype_override, load_weights=load_weights, trust_remote_code=trust_remote_code, + module_kwargs=module_kwargs, ) _save_package(pkg, output_dir, args, optimize, component_filter) @@ -457,6 +474,12 @@ def main(argv: list[str] | None = None) -> None: help="Maximum sequence length for static cache buffers. " "Only used with --static-cache. Defaults to max_position_embeddings from config.", ) + build_parser.add_argument( + "--gqa", + action="store_true", + help="Use com.microsoft::GroupQueryAttention with fused RoPE and " + "in-place KV cache support. Compatible with onnxruntime-genai.", + ) build_parser.set_defaults(func=_cmd_build) # --- build-gguf --- diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index bde20e8b..970490cb 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -246,6 +246,7 @@ def build( dtype: str | ir.DataType | None = None, load_weights: bool = True, trust_remote_code: bool = False, + module_kwargs: dict | None = None, ) -> ModelPackage: """Build an ONNX :class:`ModelPackage` from a HuggingFace model ID. @@ -278,6 +279,8 @@ def build( load_weights: Whether to download and apply weights from HuggingFace. trust_remote_code: Whether to trust remote code when loading the HuggingFace config. + module_kwargs: Extra keyword arguments passed to the module class + constructor (e.g. ``{"attention_class": GQAAttention}``). Returns: A :class:`ModelPackage` containing the built model(s). @@ -378,7 +381,7 @@ def build( if task is None: task = _default_task_for_model(model_type) - model_module = module_class(config) + model_module = module_class(config, **(module_kwargs or {})) pkg = build_from_module(model_module, config, task) # Set graph names diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 07483f5f..d8e359fc 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -29,6 +29,8 @@ "EncoderDecoderAttention", "EncoderLayer", "FCMLP", + "GQAAttention", + "GQAContext", "GatedDeltaNet", "GatedRMSNorm", "Gemma3MultiModalProjector", @@ -171,6 +173,7 @@ EncoderDecoderAttention, ) from mobius.components._gated_deltanet import GatedDeltaNet +from mobius.components._gqa_attention import GQAAttention, GQAContext from mobius.components._lightning_attention import LightningAttention from mobius.components._lora import LoRALinear from mobius.components._mamba_block import Mamba2Block, MambaBlock diff --git a/src/mobius/components/_decoder.py b/src/mobius/components/_decoder.py index 9132f7f1..6fac56ce 100644 --- a/src/mobius/components/_decoder.py +++ b/src/mobius/components/_decoder.py @@ -11,6 +11,7 @@ from mobius._configs import ArchitectureConfig from mobius.components._attention import Attention, StaticCacheState +from mobius.components._gqa_attention import GQAContext from mobius.components._mlp import MLP from mobius.components._rms_norm import RMSNorm @@ -52,15 +53,18 @@ def __init__( attention_scale: float | None = None, post_norm: bool = False, linear_class: type | None = None, + attention_class: type[nn.Module] | None = None, ): super().__init__() if norm_class is None: norm_class = RMSNorm + if attention_class is None: + attention_class = Attention self._post_norm = post_norm self._residual_multiplier = residual_multiplier - self.self_attn = Attention( + self.self_attn = attention_class( config, rms_norm_class=norm_class, scale=attention_scale, @@ -97,6 +101,17 @@ def forward( else: static_cache = None + # GQA mode: when attention_bias is a GQAContext, the attention + # component is GQAAttention and handles masking + RoPE internally. + if isinstance(attention_bias, GQAContext): + gqa_context = attention_bias + return self._forward_gqa( + op, + hidden_states, + gqa_context, + past_key_value, + ) + if self._post_norm: return self._forward_post_norm( op, @@ -150,6 +165,42 @@ def _forward_pre_norm( return hidden_states, present_key_value + def _forward_gqa( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + gqa_context: GQAContext, + past_key_value: tuple | None, + ): + """Forward pass for GQA mode (pre-norm only). + + GQAAttention handles RoPE and causal masking internally, so + position_embeddings and attention_bias are not needed. + """ + residual = hidden_states + hidden_states = self.input_layernorm(op, hidden_states) + + attn_output, present_key_value = self.self_attn( + op, + hidden_states=hidden_states, + gqa_context=gqa_context, + past_key_value=past_key_value, + ) + + if not math.isclose(self._residual_multiplier, 1.0): + attn_output = op.Mul(attn_output, self._residual_multiplier) + hidden_states = op.Add(residual, attn_output) + + residual = hidden_states + hidden_states = self.post_attention_layernorm(op, hidden_states) + hidden_states = self.mlp(op, hidden_states) + + if not math.isclose(self._residual_multiplier, 1.0): + hidden_states = op.Mul(hidden_states, self._residual_multiplier) + hidden_states = op.Add(residual, hidden_states) + + return hidden_states, present_key_value + def _forward_post_norm( self, op: builder.OpBuilder, @@ -185,6 +236,7 @@ def create_decoder_layer( norm_class: type[nn.Module] | None = None, post_norm: bool = False, linear_class: type | None = None, + attention_class: type[nn.Module] | None = None, ) -> DecoderLayer: """Config-driven factory for creating decoder layers. @@ -200,6 +252,8 @@ def create_decoder_layer( post_norm: If True, use post-norm residual connections (OLMo-2 style). linear_class: Factory callable for projection layers. Pass a LoRA factory for LoRA-adapted layers. + attention_class: Attention module class override (default: Attention). + Pass GQAAttention for ORT GenAI-compatible models. Returns: A configured DecoderLayer instance. @@ -214,6 +268,7 @@ def create_decoder_layer( attention_scale=attention_scale, post_norm=post_norm, linear_class=linear_class, + attention_class=attention_class, ) diff --git a/src/mobius/components/_gqa_attention.py b/src/mobius/components/_gqa_attention.py new file mode 100644 index 00000000..b0dc56e9 --- /dev/null +++ b/src/mobius/components/_gqa_attention.py @@ -0,0 +1,196 @@ +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""GroupQueryAttention component for ORT GenAI-compatible model generation. + +Emits the ``com.microsoft::GroupQueryAttention`` contrib op directly, +instead of the standard ONNX ``Attention`` op. This produces models +compatible with the onnxruntime-genai runtime, which uses in-place KV +cache updates (``past_present_share_buffer``) for efficient generation. + +When ``do_rotary=True``, RoPE is fused inside the GQA op using +``cos_cache`` / ``sin_cache`` graph initializers — no external +``RotaryEmbedding`` nodes are needed and ``position_ids`` is not a +graph input. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple + +from onnxscript import nn +from onnxscript._internal import builder + +from mobius._configs import ArchitectureConfig +from mobius.components._common import Linear +from mobius.components._rms_norm import RMSNorm + +if TYPE_CHECKING: + import onnx_ir as ir + + +class GQAContext(NamedTuple): + """Graph-level context for GroupQueryAttention layers. + + Created once by the task and passed through the model to each + GQAAttention instance. All layers share the same ``seqlens_k`` + and ``total_seq_len`` values (computed from ``attention_mask``). + + Fields: + seqlens_k: Per-batch actual KV 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]``. + sin_cache: Full sine RoPE table ``[max_seq_len, rotary_dim]``. + """ + + seqlens_k: ir.Value + total_seq_len: ir.Value + cos_cache: ir.Value + sin_cache: ir.Value + + +class GQAAttention(nn.Module): + """Multi-head attention using com.microsoft::GroupQueryAttention. + + Drop-in replacement for :class:`Attention` that targets the ORT GenAI + runtime. Supports GQA (grouped query attention), optional QK norm, + and fused rotary embeddings. + + The GQA op handles RoPE internally when ``cos_cache``/``sin_cache`` + are provided via :class:`GQAContext`, eliminating the need for + external ``RotaryEmbedding`` nodes. + + Args: + config: Architecture configuration. + rms_norm_class: Norm class for Q/K normalization (default: RMSNorm). + scale: Custom attention scale factor (default: 1/sqrt(head_dim)). + linear_class: Factory for projection layers (default: Linear). + """ + + def __init__( + self, + config: ArchitectureConfig, + rms_norm_class: type[nn.Module] | None = None, + scale: float | None = None, + linear_class: type | None = None, + ): + super().__init__() + if linear_class is None: + linear_class = Linear + + self.hidden_size = config.hidden_size + self.head_dim = config.head_dim + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.scaling = scale if scale is not None else self.head_dim**-0.5 + self._rope_interleave = config.rope_interleave + self._window_size = getattr(config, "sliding_window", None) or -1 + self._softcap = getattr(config, "attn_logit_softcapping", None) or 0.0 + + self.q_proj = linear_class( + self.hidden_size, + self.num_attention_heads * self.head_dim, + bias=config.attn_qkv_bias, + ) + self.k_proj = linear_class( + self.hidden_size, + self.num_key_value_heads * self.head_dim, + bias=config.attn_qkv_bias, + ) + self.v_proj = linear_class( + self.hidden_size, + self.num_key_value_heads * self.head_dim, + bias=config.attn_qkv_bias, + ) + self.o_proj = linear_class( + self.num_attention_heads * self.head_dim, + self.hidden_size, + bias=config.attn_o_bias, + ) + + if config.attn_qk_norm: + rms_norm_class = RMSNorm if rms_norm_class is None else rms_norm_class + self._qk_norm_full = config.attn_qk_norm_full + if self._qk_norm_full: + self.q_norm = rms_norm_class( + self.num_attention_heads * self.head_dim, eps=config.rms_norm_eps + ) + self.k_norm = rms_norm_class( + self.num_key_value_heads * self.head_dim, eps=config.rms_norm_eps + ) + else: + self.q_norm = rms_norm_class(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = rms_norm_class(self.head_dim, eps=config.rms_norm_eps) + else: + self._qk_norm_full = False + self.q_norm = None + self.k_norm = None + + def forward( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + gqa_context: GQAContext, + past_key_value: tuple | None = None, + ): + """Forward pass emitting com.microsoft::GroupQueryAttention. + + Args: + op: ONNX op builder. + hidden_states: ``[batch, seq_len, hidden_size]``. + gqa_context: Shared GQA context with seqlens_k, total_seq_len, + cos_cache, sin_cache. + past_key_value: Optional ``(past_key, past_value)`` tuple. + + Returns: + ``(attn_output, (present_key, present_value))`` + """ + query_states = self.q_proj(op, hidden_states) + key_states = self.k_proj(op, hidden_states) + value_states = self.v_proj(op, hidden_states) + + # Optional QK norm (applied before GQA, which handles RoPE) + if self.q_norm is not None and self.k_norm is not None: + if self._qk_norm_full: + query_states = self.q_norm(op, query_states) + key_states = self.k_norm(op, key_states) + else: + query_states = op.Reshape(query_states, [0, 0, -1, self.head_dim]) + key_states = op.Reshape(key_states, [0, 0, -1, self.head_dim]) + query_states = self.q_norm(op, query_states) + key_states = self.k_norm(op, key_states) + query_states = op.Reshape(query_states, [0, 0, -1]) + key_states = op.Reshape(key_states, [0, 0, -1]) + + 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 + + # Emit com.microsoft::GroupQueryAttention + # GQA handles RoPE internally via cos_cache/sin_cache and do_rotary=1. + # It also handles KV cache concatenation and causal masking via + # seqlens_k/total_seq_len. + attn_output, present_key, present_value = op.GroupQueryAttention( + query_states, + key_states, + value_states, + past_key, + past_value, + gqa_context.seqlens_k, + gqa_context.total_seq_len, + gqa_context.cos_cache, + gqa_context.sin_cache, + _domain="com.microsoft", + num_heads=self.num_attention_heads, + kv_num_heads=self.num_key_value_heads, + scale=self.scaling, + do_rotary=1, + rotary_interleaved=1 if self._rope_interleave else 0, + local_window_size=self._window_size, + softcap=self._softcap, + _outputs=3, + ) + + attn_output = self.o_proj(op, attn_output) + return attn_output, (present_key, present_value) diff --git a/src/mobius/models/base.py b/src/mobius/models/base.py index 9af62113..837081e3 100644 --- a/src/mobius/models/base.py +++ b/src/mobius/models/base.py @@ -28,6 +28,7 @@ from mobius.components import ( DecoderLayer, Embedding, + GQAContext, LayerNorm, Linear, RMSNorm, @@ -43,7 +44,7 @@ class TextModel(nn.Module): """Base text model with embedding, decoder layers, and final norm.""" - def __init__(self, config: ArchitectureConfig): + def __init__(self, config: ArchitectureConfig, attention_class=None): super().__init__() self._dtype = config.dtype @@ -63,7 +64,11 @@ def __init__(self, config: ArchitectureConfig): ) self.layers = nn.ModuleList( [ - DecoderLayer(config, linear_class=linear_class) + DecoderLayer( + config, + linear_class=linear_class, + attention_class=attention_class, + ) for _ in range(config.num_hidden_layers) ] ) @@ -75,14 +80,36 @@ def forward( op: builder.OpBuilder, input_ids: ir.Value, attention_mask: ir.Value | None, - position_ids: ir.Value, + position_ids: ir.Value | None, past_key_values: list | None = None, inputs_embeds: ir.Value | None = None, + gqa_context: GQAContext | None = None, ): if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.embed_tokens(op, input_ids) + + if gqa_context is not None: + # GQA mode: RoPE is fused inside the GQA op, and masking + # is handled by seqlens_k/total_seq_len. Pass GQAContext + # as the attention_bias to trigger the GQA path in DecoderLayer. + present_key_values = [] + past_kvs = past_key_values or [None] * len(self.layers) + for layer, past_kv in zip(self.layers, past_kvs): + hidden_states, present_kv = layer( + op, + hidden_states=hidden_states, + attention_bias=gqa_context, + position_embeddings=None, + past_key_value=past_kv, + ) + present_key_values.append(present_kv) + + hidden_states = self.norm(op, hidden_states) + return hidden_states, present_key_values + + # Standard path: compute RoPE and padding mask position_embeddings = self.rotary_emb(op, position_ids) # When attention_mask is None (static cache mode), skip mask @@ -132,10 +159,10 @@ class CausalLMModel(nn.Module): category: str = "Text Generation" config_class: type = CausalLMConfig - def __init__(self, config: ArchitectureConfig): + def __init__(self, config: ArchitectureConfig, attention_class=None): super().__init__() self.config = config - self.model = TextModel(config) + self.model = TextModel(config, attention_class=attention_class) self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) def forward( @@ -143,15 +170,20 @@ def forward( op: builder.OpBuilder, input_ids: ir.Value, attention_mask: ir.Value | None, - position_ids: ir.Value, + position_ids: ir.Value | None, past_key_values: list | None = None, + gqa_context: GQAContext | None = None, ): + kwargs = {} + if gqa_context is not None: + kwargs["gqa_context"] = gqa_context hidden_states, present_key_values = self.model( op, input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, + **kwargs, ) logits = self.lm_head(op, hidden_states) return logits, present_key_values @@ -182,8 +214,8 @@ class LayerNormTextModel(TextModel): weight and bias) rather than the bias-free RMS normalisation. """ - def __init__(self, config: ArchitectureConfig): - super().__init__(config) + def __init__(self, config: ArchitectureConfig, attention_class=None): + super().__init__(config, attention_class=attention_class) # Replace per-layer norms: DecoderLayer defaults to RMSNorm; override with LayerNorm. qc = getattr(config, "quantization", None) linear_class = None @@ -195,7 +227,12 @@ def __init__(self, config: ArchitectureConfig): ) self.layers = nn.ModuleList( [ - DecoderLayer(config, linear_class=linear_class, norm_class=LayerNorm) + DecoderLayer( + config, + linear_class=linear_class, + norm_class=LayerNorm, + attention_class=attention_class, + ) for _ in range(config.num_hidden_layers) ] ) @@ -216,7 +253,7 @@ class LayerNormCausalLMModel(CausalLMModel): and ``StableLmForCausalLM``. """ - def __init__(self, config: ArchitectureConfig): - super().__init__(config) + def __init__(self, config: ArchitectureConfig, attention_class=None): + super().__init__(config, attention_class=attention_class) # Replace TextModel with the LayerNorm-based variant. - self.model = LayerNormTextModel(config) + self.model = LayerNormTextModel(config, attention_class=attention_class) diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index b81be8bf..c8f664ad 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -1,7 +1,7 @@ # Copyright (c) ONNX Project Contributors # SPDX-License-Identifier: Apache-2.0 -"""Causal language model tasks with internal and static KV cache.""" +"""Causal language model tasks with internal, static, and GQA KV cache.""" from __future__ import annotations @@ -11,6 +11,7 @@ from mobius._configs import ArchitectureConfig from mobius._model_package import ModelPackage from mobius.components._attention import StaticCacheState +from mobius.components._gqa_attention import GQAContext from mobius.tasks._base import ( ModelTask, _make_graph, @@ -23,10 +24,18 @@ ) +def _find_rotary_emb(module: nn.Module) -> nn.Module | None: + """Find the rotary embedding module by walking the module tree.""" + for name, child in module.named_modules(): + if name.endswith("rotary_emb") and hasattr(child, "cos_cache"): + return child + return None + + class CausalLMTask(ModelTask): """Causal language model with KV cache for text generation. - Supports two cache modes: + Supports three cache modes: **Dynamic cache** (default): Standard KV cache with dynamic sequence lengths. Past keys/values @@ -61,11 +70,30 @@ class CausalLMTask(ModelTask): No ``attention_mask`` input — causal masking uses ``is_causal=1``. + **GQA mode** (``gqa=True``): + Uses ``com.microsoft::GroupQueryAttention`` with fused RoPE and + in-place KV cache support. Compatible with the onnxruntime-genai + runtime (``past_present_share_buffer=true``). + + Inputs: + - input_ids: [batch, sequence_len] INT64 + - attention_mask: [batch, total_seq_len] INT64 + - past_key_values.{i}.key: [batch, num_kv_heads, past_seq_len, head_dim] + - past_key_values.{i}.value: [batch, num_kv_heads, past_seq_len, head_dim] + Outputs: + - logits: FLOAT + - present.{i}.key / present.{i}.value: FLOAT + + No ``position_ids`` input — RoPE is fused inside GroupQueryAttention + using cos_cache/sin_cache initializers. + The module's ``forward()`` must accept ``(op, input_ids, attention_mask, position_ids, past_key_values)`` and return ``(logits, list_of_(key, value)_tuples)``. In static cache mode, ``attention_mask`` will be ``None`` and ``past_key_values`` - entries will be :class:`StaticCacheState` tuples. + entries will be :class:`StaticCacheState` tuples. In GQA mode, + ``position_ids`` will be ``None`` and a ``gqa_context`` kwarg is + passed. Args: static_cache: If ``True``, use pre-allocated static KV cache @@ -73,6 +101,8 @@ class CausalLMTask(ModelTask): max_seq_len: Maximum sequence length for static cache buffers. Only used when ``static_cache=True``. Defaults to ``config.max_position_embeddings``. + gqa: If ``True``, use ``com.microsoft::GroupQueryAttention`` + with fused RoPE and in-place KV cache support. """ def __init__( @@ -80,15 +110,20 @@ def __init__( *, static_cache: bool = False, max_seq_len: int | None = None, + gqa: bool = False, ): self._static_cache = static_cache self._max_seq_len = max_seq_len + self._gqa = gqa def build( self, module: nn.Module, config: ArchitectureConfig, ) -> ModelPackage: + if self._gqa: + return self._build_gqa(module, config) + static = self._static_cache # --- Static-cache pre-validation --- @@ -195,6 +230,107 @@ def build( return ModelPackage({"model": _make_model(graph)}, config=config) + def _build_gqa( + self, + module: nn.Module, + config: ArchitectureConfig, + ) -> ModelPackage: + """Build a model using com.microsoft::GroupQueryAttention. + + Creates the graph with: + - No position_ids input (RoPE is fused inside GQA) + - seqlens_k and total_seq_len computed from attention_mask + - cos_cache and sin_cache as graph initializers + - Standard 4D KV cache I/O (same naming as dynamic cache) + """ + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + past_seq_len = ir.SymbolicDim("past_sequence_len") + + # --- Graph inputs --- + input_ids = ir.Value( + name="input_ids", + shape=ir.Shape([batch, seq_len]), + type=ir.TensorType(ir.DataType.INT64), + ) + attention_mask = ir.Value( + name="attention_mask", + shape=ir.Shape([batch, "past_seq_len + seq_len"]), + type=ir.TensorType(ir.DataType.INT64), + ) + # No position_ids — GQA handles RoPE internally + + graph_inputs = [input_ids, attention_mask] + + # --- KV cache inputs (same 4D shape as dynamic cache) --- + cache_inputs, past_key_values = _make_kv_cache_inputs( + config.num_hidden_layers, + config.num_key_value_heads, + config.head_dim, + config.dtype, + batch, + past_seq_len, + ) + graph_inputs.extend(cache_inputs) + + # --- Build graph --- + graph, builder = _make_graph(graph_inputs) + op = builder.op + + # --- Compute seqlens_k and total_seq_len from attention_mask --- + # seqlens_k = Cast(ReduceSum(attention_mask, axis=1) - 1, INT32) + axis = op.Constant(value_ints=[1]) + reduce_sum = op.ReduceSum(attention_mask, axis) + one = op.Constant(value_ints=[1]) + seqlens_k = op.Cast(op.Sub(reduce_sum, one), to=ir.DataType.INT32) + + # total_seq_len = Cast(Gather(Shape(attention_mask), 1), INT32) + mask_shape = op.Shape(attention_mask) + idx_1 = op.Constant(value_int=1) + total_seq_len = op.Cast(op.Gather(mask_shape, idx_1), to=ir.DataType.INT32) + + # --- Get cos/sin cache from the model's rotary embedding --- + # The model must have a rotary_emb with cos_cache/sin_cache parameters. + rotary_emb = _find_rotary_emb(module) + if rotary_emb is None: + raise ValueError( + "GQA mode requires the model to have a rotary embedding " + "module with cos_cache and sin_cache parameters. " + "Ensure the model uses initialize_rope()." + ) + # Register cos/sin caches as graph initializers so they get + # serialized. The rotary_emb module isn't called in GQA mode + # (GQA handles RoPE internally), so its parameters won't be + # auto-registered by the builder's module call machinery. + cos_cache = builder.initializer(rotary_emb.cos_cache.const_value, name="cos_cache") + sin_cache = builder.initializer(rotary_emb.sin_cache.const_value, name="sin_cache") + + gqa_context = GQAContext( + seqlens_k=seqlens_k, + total_seq_len=total_seq_len, + cos_cache=cos_cache, + sin_cache=sin_cache, + ) + + # --- Invoke module --- + logits, present_key_values = module( + op, + input_ids=input_ids, + attention_mask=None, + position_ids=None, + past_key_values=past_key_values, + gqa_context=gqa_context, + ) + + logits.name = "logits" + graph.outputs.append(logits) + _register_kv_cache_outputs(graph, present_key_values) + + model = _make_model(graph) + # Register the com.microsoft domain opset import for GQA + model.graph.opset_imports["com.microsoft"] = 1 + return ModelPackage({"model": model}, config=config) + class HybridCausalLMTask(ModelTask): """Causal LM with hybrid KV cache + DeltaNet recurrent states. diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 3d25ac8f..fceb2ebb 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -3784,3 +3784,113 @@ def test_outputs_have_shapes_and_dtypes(self, model_type: str, config_overrides: task = get_task(_default_task_for_model(model_type)) pkg = task.build(module, config) _assert_outputs_have_shapes_and_dtypes(pkg, model_type) + + +class TestBuildGQAGraph: + """Verify CausalLMTask(gqa=True) builds a valid GQA graph.""" + + def _build_gqa_model(self, model_type: str = "qwen2", **config_overrides): + """Build a model with GQA mode and return (model, config).""" + from mobius.components import GQAAttention + + config = _base_config(**config_overrides) + model_cls = registry.get(model_type) + module = model_cls(config, attention_class=GQAAttention) + task = CausalLMTask(gqa=True) + pkg = task.build(module, config) + return pkg["model"], config + + def test_gqa_graph_builds(self): + """Build a Qwen2 model with GQA mode.""" + model, _ = self._build_gqa_model() + assert model.graph is not None + assert len(model.graph.inputs) > 0 + assert len(model.graph.outputs) > 0 + + def test_gqa_graph_inputs(self): + """Verify inputs: input_ids, attention_mask, KV caches (no position_ids).""" + model, config = self._build_gqa_model() + input_names = {inp.name for inp in model.graph.inputs} + num_layers = config.num_hidden_layers + + assert "input_ids" in input_names + assert "attention_mask" in input_names + # GQA mode: no position_ids (RoPE is fused inside GQA op) + assert "position_ids" not in input_names + + for i in range(num_layers): + assert f"past_key_values.{i}.key" in input_names + assert f"past_key_values.{i}.value" in input_names + + # Exact count: 2 standard + 2*num_layers caches + expected_count = 2 + 2 * num_layers + assert len(model.graph.inputs) == expected_count, ( + f"Expected {expected_count} inputs, got {len(model.graph.inputs)}" + ) + + def test_gqa_graph_outputs(self): + """Verify outputs: logits + present KV caches.""" + model, config = self._build_gqa_model() + output_names = {out.name for out in model.graph.outputs} + num_layers = config.num_hidden_layers + + assert "logits" in output_names + for i in range(num_layers): + assert f"present.{i}.key" in output_names + assert f"present.{i}.value" in output_names + + expected_count = 1 + 2 * num_layers + assert len(model.graph.outputs) == expected_count + + def test_gqa_has_group_query_attention_ops(self): + """Verify the graph uses GroupQueryAttention, not standard Attention.""" + model, _config = self._build_gqa_model() + op_types = {n.op_type for n in model.graph} + + assert "GroupQueryAttention" in op_types + assert "Attention" not in op_types, "GQA graph should not have standard Attention ops" + + def test_gqa_no_rotary_embedding_nodes(self): + """Verify RoPE is fused inside GQA (no RotaryEmbedding nodes).""" + model, _ = self._build_gqa_model() + op_types = {n.op_type for n in model.graph} + assert "RotaryEmbedding" not in op_types + + def test_gqa_has_cos_sin_initializers(self): + """Verify cos_cache and sin_cache are graph initializers.""" + model, _ = self._build_gqa_model() + init_names = list(model.graph.initializers) + assert any("cos_cache" in n for n in init_names), ( + "GQA graph should have cos_cache initializer" + ) + assert any("sin_cache" in n for n in init_names), ( + "GQA graph should have sin_cache initializer" + ) + + def test_gqa_has_microsoft_opset(self): + """Verify com.microsoft opset is imported.""" + model, _ = self._build_gqa_model() + assert "com.microsoft" in model.graph.opset_imports + + def test_gqa_graph_validates(self): + """Verify the graph survives a serialization round-trip.""" + model, _ = self._build_gqa_model() + proto = ir.serde.serialize_model(model) + assert len(proto.SerializeToString()) > 0 + + def test_gqa_do_rotary_attribute(self): + """Verify GroupQueryAttention nodes have do_rotary=1.""" + model, config = self._build_gqa_model() + gqa_nodes = [n for n in model.graph if n.op_type == "GroupQueryAttention"] + assert len(gqa_nodes) == config.num_hidden_layers + + for node in gqa_nodes: + do_rotary = node.attributes.get("do_rotary") + assert do_rotary is not None + assert do_rotary.value == 1 + + def test_gqa_llama_model(self): + """Verify GQA mode also works with llama model type.""" + model, _ = self._build_gqa_model(model_type="llama") + op_types = {n.op_type for n in model.graph} + assert "GroupQueryAttention" in op_types