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

scripts/generate_golden.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,

src/mobius/_configs/_base.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/_sub_configs.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/per_model/__init__.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/_gemma4_unified_audio.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_vision.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/_optimizations.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/_registry.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/components/_attention.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

src/mobius/components/_common.py

@@ -161,6 +161,7 @@ def create_attention_bias(
     attention_mask,
     sliding_window: int | None = None,
     dtype: ir.DataType = ir.DataType.FLOAT,
+    block_sequence_ids=None,
 ):
     """Create causal attention bias for use in attention mechanisms.
 
@@ -172,6 +173,18 @@ def create_attention_bias(
         dtype: Data type for the bias tensor. The masked value uses the
             minimum representable value for this dtype (e.g. -65504 for
             float16, -3.4e38 for float32).
+        block_sequence_ids: Optional INT tensor of shape
+            (batch_size, query_length) giving a contiguous vision-block id
+            per current-sequence position (``>= 0`` for vision tokens in the
+            same block, ``-1`` for text). When provided, a bidirectional
+            "blockwise overlay" is OR-ed onto the causal (and sliding) mask:
+            two positions in the same block (same id ``>= 0``) may attend to
+            each other regardless of causal order. This mirrors HuggingFace
+            ``blockwise_overlay`` for Gemma4 ``use_bidirectional_attention``.
+            The returned bias bakes in causal + sliding + padding + blockwise,
+            so the consuming ``Attention`` op MUST be called with
+            ``is_causal=0`` to avoid re-applying the causal constraint and
+            cancelling the bidirectional unmasking.
 
     Returns:
         Attention bias tensor of shape (batch_size, 1, query_length, total_length).
@@ -210,6 +223,34 @@ def create_attention_bias(
         within_window = op.Less(dist, sliding_window)
         full_mask = op.And(full_mask, within_window)
 
+    if block_sequence_ids is not None:
+        # Bidirectional vision-block overlay (OR-ed onto causal/sliding mask,
+        # BEFORE the padding AND, matching HF blockwise_overlay ordering).
+        #
+        # q_group: block id per query position -> (batch, query_length, 1).
+        # block_sequence_ids covers the current input (== the query), so it
+        # aligns 1:1 with the query positions.
+        q_group = op.Unsqueeze(block_sequence_ids, [2])  # (B, q_len, 1)
+        # kv_group: block id per kv position -> (batch, 1, total_length).
+        # The kv axis spans past + current; past positions are text in the
+        # cache, so left-pad with -1 to width total_length.
+        pad_width = op.Sub(total_length, query_length)  # [1], == past length
+        zero_1d = op.Constant(value_ints=[0])
+        # Pad spec for a 2-D tensor [B, q_len]: [b_begin, s_begin, b_end, s_end].
+        pads = op.Concat(zero_1d, pad_width, zero_1d, zero_1d, axis=0)
+        kv_group_2d = op.Pad(
+            block_sequence_ids,
+            pads,
+            op.Constant(value_int=-1),
+        )  # (B, total_length)
+        kv_group = op.Unsqueeze(kv_group_2d, [1])  # (B, 1, total_length)
+        # same_block = (q_group == kv_group) AND (q_group >= 0)
+        same_block = op.And(
+            op.Equal(q_group, kv_group),
+            op.GreaterOrEqual(q_group, op.Constant(value_int=0)),
+        )
+        full_mask = op.Or(full_mask, same_block)
+
     # Combine with attention_mask
     attn_mask_bool = op.Cast(op.Unsqueeze(attention_mask, [1]), to=ir.DataType.BOOL)
     full_mask = op.And(attn_mask_bool, full_mask)