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

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

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

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

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