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

src/mobius/_configs/_base.py

@@ -1267,6 +1267,17 @@ 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 vision-token blocks
+      (image/audio placeholders) 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.
+    - ``"all"``: every token attends bidirectionally (no causal mask).
+    """
 
     @classmethod
     def from_transformers(cls, config, parent_config=None) -> Gemma4Config:
@@ -1340,6 +1351,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/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,37 @@
+# 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),
+        )
+    }

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/_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)

src/mobius/components/_common_test.py

@@ -5,9 +5,11 @@
 
 from __future__ import annotations
 
+import numpy as np
 import onnx_ir as ir
 
 from mobius._testing import count_op_type, create_test_builder, create_test_input
+from mobius._testing.ort_inference import OnnxModelSession
 from mobius.components._common import (
     Embedding,
     Linear,
@@ -117,6 +119,100 @@ def test_query_length_from_input_ids_not_attention_mask(self):
         )
 
 
+class TestBlockwiseAttentionBias:
+    """Numerically verify the Gemma4 vision-block bidirectional overlay.
+
+    ``create_attention_bias(block_sequence_ids=...)`` must bake the FULL mask
+    (causal [+ sliding] OR same-block, AND padding) so the Attention op can be
+    called with ``is_causal=0``. We build the graph, run it via ORT, and
+    compare the attended pattern (bias == 0) against a numpy reference.
+    """
+
+    @staticmethod
+    def _build(sliding):
+        b, op, g = create_test_builder()
+        input_ids = create_test_input(b, "input_ids", [1, "S"], dtype=ir.DataType.INT64)
+        attn = create_test_input(b, "attention_mask", [1, "T"], dtype=ir.DataType.INT64)
+        bsid = create_test_input(b, "block_sequence_ids", [1, "S"], dtype=ir.DataType.INT64)
+        bias = create_attention_bias(
+            op,
+            input_ids,
+            attn,
+            sliding_window=sliding,
+            dtype=ir.DataType.FLOAT,
+            block_sequence_ids=bsid,
+        )
+        bias.name = "bias"
+        g.outputs.append(bias)
+        return ir.Model(g, ir_version=10)
+
+    @staticmethod
+    def _ref(block_ids, attn, sliding):
+        cumsum = np.cumsum(attn)
+        qi = cumsum[:, None]
+        ki = cumsum[None, :]
+        m = qi >= ki
+        if sliding is not None:
+            m = m & ((qi - ki) < sliding)
+        qg = np.array(block_ids)[:, None]
+        kg = np.array(block_ids)[None, :]
+        m = m | ((qg == kg) & (qg >= 0))
+        return m & (np.array(attn)[None, :].astype(bool))
+
+    def _run(self, sliding, block_ids, attn):
+        sess = OnnxModelSession(self._build(sliding), device="cpu")
+        block_ids = np.array([block_ids], dtype=np.int64)
+        attn = np.array([attn], dtype=np.int64)
+        out = sess.run(
+            {
+                "input_ids": np.zeros_like(block_ids),
+                "attention_mask": attn,
+                "block_sequence_ids": block_ids,
+            }
+        )["bias"]
+        attended = out[0, 0] > -1.0
+        expected = self._ref(block_ids[0], attn[0], sliding)
+        return attended, expected
+
+    def test_multi_block_full_attention(self):
+        # Two vision blocks (pos 1-2 and 4-5) separated by text.
+        attended, expected = self._run(None, [-1, 0, 0, -1, 1, 1, -1, -1], [1] * 8)
+        assert np.array_equal(attended, expected)
+        # A vision token attends to a LATER token in the same block (bidirectional).
+        assert attended[1, 2]
+        # But text stays causal: position 3 cannot see position 4.
+        assert not attended[3, 4]
+
+    def test_block_wider_than_sliding_window(self):
+        # Single block spanning positions 1..4 with a window of 2: the block
+        # must escape the sliding window (same-block OR overrides the window).
+        attended, expected = self._run(2, [-1, 0, 0, 0, 0, -1, -1, -1], [1] * 8)
+        assert np.array_equal(attended, expected)
+        assert attended[4, 1]  # distance 3 >= window, allowed via same block
+
+    def test_padding_still_masked(self):
+        # Last two positions are padding (attention_mask == 0).
+        attended, expected = self._run(
+            2, [-1, 0, 0, -1, -1, -1, -1, -1], [1, 1, 1, 1, 1, 1, 0, 0]
+        )
+        assert np.array_equal(attended, expected)
+        assert not attended[:, 6:].any()  # nothing attends to padding
+
+    def test_decode_single_query_is_causal(self):
+        # Decode step: q_len=1 (new text token, group -1), kv total length 8.
+        sess = OnnxModelSession(self._build(None), device="cpu")
+        out = sess.run(
+            {
+                "input_ids": np.zeros((1, 1), dtype=np.int64),
+                "attention_mask": np.ones((1, 8), dtype=np.int64),
+                "block_sequence_ids": np.array([[-1]], dtype=np.int64),
+            }
+        )["bias"]
+        assert out.shape == (1, 1, 1, 8)
+        # Text decode token attends to all past positions (pure causal row).
+        assert bool((out[0, 0, 0] > -1.0).all())
+
+
 class TestCreatePaddingMask:
     def test_creates_bool_mask_with_2d_input_ids(self):
         """Standard path: input_ids is 2D [batch, q_len]."""