Fix Gemma4 bidirectional attention + add gemma-4-12B (gemma4_unified) (#338)

## Summary

Two related changes for the Gemma 4 family:

### 1. Fix bidirectional vision-block attention (existing gemma4 models)
The existing gemma4 (32B) and gemma4 MoE multimodal models always built
a
purely **causal** decoder, ignoring
`use_bidirectional_attention="vision"`.
HuggingFace gemma4 makes contiguous vision-token spans attend
**bidirectionally** (via a blockwise overlay OR-ed onto the
causal/sliding
mask, applied to both full- and sliding-attention layers). This adds:

- `use_bidirectional_attention` to `Gemma4Config` + extraction.
- `block_sequence_ids` plumbed embedding → decoder;
`create_attention_bias`
gains a `block_sequence_ids` blockwise-overlay path (forces the
float-bias,
  non-GQA attention path with `is_causal=0`).

### 2. Add gemma-4-12B (`gemma4_unified`)
The full encoder-free multimodal
`Gemma4UnifiedForConditionalGeneration`:

- **Text backbone**: 48 layers, hidden 3840, dual head_dim (local 256 /
global 512), `attention_k_eq_v` with a single global KV head, dual RoPE,
  final-logit softcapping. Reuses `Gemma4CausalLMModel`.
- **Vision**: encoder-free embedder — raw patches (6912) → LN → Dense →
LN
  → factorized 2D posemb → LN → scale-free RMSNorm → Linear(→3840).
- **Audio**: encoder-free embedder — raw frames (640) → scale-free
RMSNorm
  → Linear(→3840).
- Both strip padding inside the ONNX graph (Compress).
- `Gemma4UnifiedModel` + `Gemma4UnifiedTask` build a 3/4-model package
  (decoder, vision_encoder, embedding, optional audio_encoder).

## Testing

Fast suite: 2783 passed. L1 graph-build tests for text and 4-model
multimodal packages.

Real `google/gemma-4-12B` checkpoint, float32 on CUDA (H200):
- **Text prefill** parity vs HF: max_abs_diff ~0.01.
- **Bidirectional mask** parity test.
- **Full multimodal prefill** parity vs HF: vision cos_sim=1.000000,
  decoder last-token cos_sim=1.000000, max_abs_diff=1.10, argmax match,
  no NaN. (Bidirectional reference requires passing `mm_token_type_ids`
  to HF.)

---------

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

Author: justinchuby

examples/gemma4_unified_ort_genai.py

@@ -403,10 +403,14 @@ def _generate_vision_language(case: TestCase, json_path: Path, device: str) -> N
     # Load images from testdata/
     images = [Image.open(Path("testdata") / img_path) for img_path in case.images]
 
-    # Build chat-formatted prompt with image placeholders if the
-    # processor supports apply_chat_template (Qwen-VL, Gemma-3, etc.)
+    # Build chat-formatted prompt with image placeholders if the processor has a
+    # usable chat template (Qwen-VL, Gemma-3, etc.). Base checkpoints (e.g.
+    # google/gemma-4-12B) ship no chat template, so fall back to manually
+    # prepending one image placeholder token per image — the processor then
+    # expands each into the correct number of soft tokens (mirrors how
+    # examples/gemma4_unified_ort_genai.py formats image prompts).
     prompt_text = case.prompts[0]
-    if hasattr(processor, "apply_chat_template"):
+    if getattr(processor, "chat_template", None):
         content: list[dict[str, str]] = []
         for img_path in case.images:
             content.append({"type": "image", "image": str(Path("testdata") / img_path)})
@@ -415,6 +419,8 @@ def _generate_vision_language(case: TestCase, json_path: Path, device: str) -> N
         prompt_text = processor.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True
         )
+    elif getattr(processor, "image_token", None):
+        prompt_text = processor.image_token * len(case.images) + prompt_text
 
     # Process multimodal inputs through the HF processor
     processed = processor(
@@ -739,7 +745,7 @@ def _prepare_speech_language_inputs(
     else:
         # Gemma4-style: text prompt + audio
         prompt_text = case.prompts[0]
-        if hasattr(processor, "apply_chat_template"):
+        if getattr(processor, "chat_template", None):
             content: list[dict[str, str]] = [
                 {"type": "audio", "audio": str(audio_path)},
                 {"type": "text", "text": prompt_text},
@@ -748,6 +754,10 @@ def _prepare_speech_language_inputs(
             prompt_text = processor.apply_chat_template(
                 messages, tokenize=False, add_generation_prompt=True
             )
+        elif getattr(processor, "audio_token", None):
+            # Base checkpoint (no chat template): manually prepend the audio
+            # placeholder; the processor expands it to the right token count.
+            prompt_text = processor.audio_token + prompt_text
         model_device = _get_model_device(model, device)
         processed = processor(
             text=prompt_text,

scripts/generate_golden.py

@@ -1267,6 +1267,22 @@ class Gemma4Config(VisionLanguageConfig):
     enable_moe_block: bool = False
     attention_k_eq_v: bool = False
     boa_token_id: int | None = None
+    use_bidirectional_attention: str | None = None
+    """Bidirectional attention mode for the text decoder.
+
+    Mirrors HF ``Gemma4TextConfig.use_bidirectional_attention``:
+    - ``None``: fully causal (smaller Gemma4 models, e.g. E2B).
+    - ``"vision"``: text stays causal, but contiguous image-token blocks
+      attend bidirectionally within each block (larger models, e.g.
+      12B/26B/32B). Implemented via a per-position ``block_sequence_ids``
+      overlay added onto the causal mask. Audio placeholders are *not*
+      included (HF marks audio as token-type 3, excluded from the vision
+      block mask), so audio tokens keep causal attention.
+    - ``"all"``: HF mode where every token attends bidirectionally. Not used
+      by any currently supported Gemma4 model and not implemented here; the
+      decoder raises ``NotImplementedError`` rather than silently degrading to
+      causal attention (only ``None`` and ``"vision"`` are accepted).
+    """
 
     @classmethod
     def from_transformers(cls, config, parent_config=None) -> Gemma4Config:
@@ -1340,6 +1356,7 @@ def from_transformers(cls, config, parent_config=None) -> Gemma4Config:
             enable_moe_block=getattr(config, "enable_moe_block", False),
             attention_k_eq_v=getattr(config, "attention_k_eq_v", False),
             boa_token_id=getattr(parent_config, "boa_token_id", None),
+            use_bidirectional_attention=getattr(config, "use_bidirectional_attention", None),
         )
 
 

src/mobius/_configs/_base.py

@@ -201,6 +201,9 @@ class AudioConfig:
     audio_start_token_id: int | None = None
     audio_end_token_id: int | None = None
     classify_num: int | None = None
+    # RMSNorm epsilon for the audio encoder/embedder (may differ from the text
+    # decoder's rms_norm_eps). Falls back to the text value when unset.
+    rms_norm_eps: float | None = None
     # Qwen3-ASR chunked conv parameters. ``n_window`` is half the
     # number of mel frames per conv chunk (so chunk_size = 2 *
     # n_window). ``n_window_infer`` is the attention window in mel

src/mobius/_configs/_sub_configs.py

@@ -24,6 +24,8 @@
 # may freely re-sort this block.
 from mobius._configs.per_model import (  # noqa: F401
     _gemma4_audio,
+    _gemma4_unified_audio,
+    _gemma4_unified_vision,
     _hunyuan_vl_mot_vision,
     _internvl_vision,
     _phi4mm_audio,

src/mobius/_configs/per_model/__init__.py

@@ -0,0 +1,38 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Gemma4-unified (gemma-4-12B) audio extractor hook.
+
+The ``gemma4_unified`` audio config describes an *encoder-free* embedder (no
+Conformer tower).  It exposes only ``audio_embed_dim`` (input feature size for
+the projection) and ``rms_norm_eps``.  This hook maps those onto
+:class:`Gemma4AudioConfig` so
+:class:`~mobius.models.gemma4._Gemma4UnifiedAudioEmbedderModel` can read them.
+"""
+
+from __future__ import annotations
+
+from mobius._configs._extractors import register_audio_hook
+from mobius._configs._sub_configs import Gemma4AudioConfig
+
+_UNIFIED_TYPES = ("gemma4_unified", "gemma4_unified_text", "gemma4_unified_audio")
+
+
+@register_audio_hook
+def _gemma4_unified_audio(config, parent_config, model_type: str, fields: dict):
+    composite = parent_config or config
+    parent_model_type = getattr(composite, "model_type", "")
+    if model_type not in _UNIFIED_TYPES and parent_model_type != "gemma4_unified":
+        return None
+    hf_audio = getattr(composite, "audio_config", None)
+    if hf_audio is None:
+        return None
+    audio_embed_dim = getattr(hf_audio, "audio_embed_dim", 640)
+    return {
+        "audio": Gemma4AudioConfig(
+            hidden_size=audio_embed_dim,
+            output_proj_dims=audio_embed_dim,
+            audio_token_id=getattr(composite, "audio_token_id", None),
+            rms_norm_eps=getattr(hf_audio, "rms_norm_eps", None),
+        )
+    }

src/mobius/_configs/per_model/_gemma4_unified_audio.py

@@ -0,0 +1,49 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Gemma4-unified (gemma-4-12B) vision extractor hook.
+
+The ``gemma4_unified`` vision config describes an *encoder-free* embedder, not
+a SigLIP tower.  Its fields differ from the generic ``vision_config``:
+
+- ``patch_size`` / ``pooling_kernel_size`` → merged ``model_patch_size``
+- ``mm_embed_dim`` → embedder hidden size (``VisionConfig.hidden_size``)
+- ``mm_posemb_size`` → factorized positional-embedding table size
+  (``VisionConfig.position_embedding_size``)
+- ``output_proj_dims`` → projection input dim (``VisionConfig.out_hidden_size``)
+
+This hook maps those onto :class:`VisionConfig` so
+:class:`~mobius.models.gemma4._Gemma4UnifiedVisionEmbedderModel` can read them.
+"""
+
+from __future__ import annotations
+
+from mobius._configs._extractors import register_vision_hook
+
+_UNIFIED_TYPES = ("gemma4_unified", "gemma4_unified_text", "gemma4_unified_vision")
+
+
+@register_vision_hook
+def _gemma4_unified_vision(config, parent_config, model_type: str, fields: dict):
+    composite = parent_config or config
+    parent_model_type = getattr(composite, "model_type", "")
+    if model_type not in _UNIFIED_TYPES and parent_model_type != "gemma4_unified":
+        return None
+    hf_vision = getattr(composite, "vision_config", None)
+    if hf_vision is None:
+        return None
+
+    def _get(name, default=None):
+        return getattr(hf_vision, name, default)
+
+    fields.update(
+        model_type="gemma4_unified_vision",
+        hidden_size=_get("mm_embed_dim", 3840),
+        patch_size=_get("patch_size", 16),
+        pooling_kernel_size=_get("pooling_kernel_size", 3),
+        position_embedding_size=_get("mm_posemb_size", 1120),
+        out_hidden_size=_get("output_proj_dims", _get("mm_embed_dim", 3840)),
+        norm_eps=_get("rms_norm_eps", 1e-6),
+    )
+    fields["image_token_id"] = getattr(composite, "image_token_id", None)
+    return None

src/mobius/_configs/per_model/_gemma4_unified_vision.py

@@ -284,7 +284,12 @@ def _get_optimization_passes(
 
     # --- Attention fusion (decoder only) ---
     if model_role == "decoder" and dtype in caps.gqa_dtypes:
-        fuse.append(("GQAFusion", list(group_query_attention_rules())))
+        fuse.append(
+            (
+                "GQAFusion",
+                list(group_query_attention_rules()),
+            )
+        )
 
     # --- QKV packing (decoder only, gated by qkv_pack_dtypes) ---
     if model_role == "decoder" and dtype in caps.qkv_pack_dtypes:

src/mobius/_optimizations.py

@@ -45,6 +45,7 @@
     Gemma3MultiModalModel,
     Gemma4CausalLMModel,
     Gemma4Model,
+    Gemma4UnifiedModel,
     GemmaCausalLMModel,
     Glm4CausalLMModel,
     Glm4MoECausalLMModel,
@@ -398,6 +399,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "gemma3n": ModelRegistration(Gemma3nCausalLMModel),
     "gemma3n_text": ModelRegistration(Gemma3nCausalLMModel),
     "gemma4_text": ModelRegistration(Gemma4CausalLMModel, config_class=Gemma4Config),
+    "gemma4_unified_text": ModelRegistration(Gemma4CausalLMModel, config_class=Gemma4Config),
     "glm": ModelRegistration(GlmCausalLMModel),
     "glm4": ModelRegistration(Glm4CausalLMModel),
     "gpt_neox": ModelRegistration(GPTNeoXCausalLMModel),
@@ -472,6 +474,9 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "florence2": ModelRegistration(LLaVAModel, task="vision-language"),
     "fuyu": ModelRegistration(LLaVAModel, task="vision-language"),
     "gemma4": ModelRegistration(Gemma4Model, task="gemma4", config_class=Gemma4Config),
+    "gemma4_unified": ModelRegistration(
+        Gemma4UnifiedModel, task="gemma4-unified", config_class=Gemma4Config
+    ),
     "glm4v": ModelRegistration(LLaVAModel, task="vision-language"),
     "glm4v_moe": ModelRegistration(LLaVAModel, task="vision-language"),
     "glm4v_moe_text": ModelRegistration(Glm4MoECausalLMModel),
@@ -819,6 +824,8 @@ def _create_default_registry() -> ModelRegistry:
     "llava_next": "llava-hf/llava-v1.6-mistral-7b-hf",
     "mllama": "meta-llama/Llama-3.2-11B-Vision-Instruct",
     "gemma4": "google/gemma-4-E2B-it",
+    "gemma4_unified": "google/gemma-4-12B",
+    "gemma4_unified_text": "google/gemma-4-12B",
     "internvl2": "OpenGVLab/InternVL2-1B",
     "phi4mm": "microsoft/Phi-4-multimodal-instruct",
     "phi4_multimodal": "microsoft/Phi-4-multimodal-instruct",

src/mobius/_registry.py

@@ -88,6 +88,7 @@ def _apply_attention(
     scale: float,
     softcap: float = 0.0,
     static_cache: StaticCacheState | None = None,
+    is_causal: int = 1,
 ) -> tuple[ir.Value, ir.Value, ir.Value]:
     """Apply the ONNX Attention op with internal or static KV cache.
 
@@ -103,10 +104,20 @@ def _apply_attention(
         Also uses ``is_causal=1``.
         Returns ``(attn_output, updated_key_cache, updated_value_cache)``.
 
+    Args:
+        is_causal: Whether the Attention op applies its built-in causal
+            mask (default ``1``). Set to ``0`` when ``attn_mask`` already
+            bakes the FULL mask (causal + sliding + padding, and any
+            bidirectional unmasking such as Gemma4's vision-block overlay)
+            into a float additive bias. Leaving ``is_causal=1`` in that
+            case would re-apply causality and cancel any future-position
+            unmasking encoded in the bias.
+
     Note:
-        Both paths set ``is_causal=1`` on the Attention op, which enables
-        built-in causal masking. This means ``attn_mask`` should encode
-        only padding information (as a bool mask), not causality.
+        Both paths default to ``is_causal=1`` on the Attention op, which
+        enables built-in causal masking. This means ``attn_mask`` should
+        encode only padding information (as a bool mask), not causality,
+        unless ``is_causal=0`` is passed explicitly.
 
     Note:
         ``nonpad_kv_seqlen`` (input #6) is only valid in static cache mode
@@ -167,7 +178,7 @@ def _apply_attention(
             kv_num_heads=num_key_value_heads,
             scale=scale,
             softcap=softcap,
-            is_causal=1,
+            is_causal=is_causal,
             _outputs=3,
         )
         return attn_output, updated_k, updated_v
@@ -191,7 +202,7 @@ def _apply_attention(
         kv_num_heads=num_key_value_heads,
         scale=scale,
         softcap=softcap,
-        is_causal=1,
+        is_causal=is_causal,
         _outputs=3,
     )
     return attn_output, present_key, present_value