From dc7c56b71a07957a50e74dc3a639911a0d7b177f Mon Sep 17 00:00:00 2001
From: dengcchi <dengcchi@meta.com>
Date: Fri, 7 Aug 2026 20:07:24 -0700
Subject: [PATCH 1/3] Gemma-4 vision: add per-projection clipped-linears
 (E2B/E4B parity prerequisite)

The Gemma-4 E2B/E4B reference checkpoint ships per-projection activation clip
bounds for the vision tower: each of the 7 vision projections
(self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj) in every encoder block
carries a scalar {input,output}_{min,max} (16 blocks x 7 x 4 = 448 bounds), and
the reference forward clamps each projection's input and output by those bounds.
MaxText does not model them today, so the values are silently dropped on
conversion. They are one of the pieces required before E2B/E4B image inputs can
match the reference.

This adds the clip bounds as an opt-in, checkpoint-resident, non-trainable
feature gated on use_clipped_linears_for_vit (exact no-op when False):

- gemma4_vision.py: clip-bound helpers (_clip_in/_clip_out, _ClipBounds,
  NaN-sentinel validate_clip_bounds, path-based clip_optimizer_freeze_mask).
  Gemma4Attention overrides its q/k/v/o projection methods to clamp-in ->
  DenseGeneral -> clamp-out; Gemma4ClippedMlpBlock does the same for
  gate/up/down. The shared attentions.Attention / linears.MlpBlock are untouched,
  and the underlying DenseGeneral weights and checkpoint key paths are unchanged
  (the bounds are separate scalar leaves).
- param_mapping.py: maps the 448 clip-bound scalars from the HF checkpoint when
  the flag is set.
- types.py / base.yml: adds the use_clipped_linears_for_vit flag (default False).

The bounds are plain nnx.Param leaves so they round-trip through the
nnx->linen->orbax checkpoint path and map to the canonical params collection;
clip_optimizer_freeze_mask keeps them out of optimizer updates and weight decay
via a leaf-path mask. A NaN sentinel + validate_clip_bounds hard-fails on a
missing/non-finite bound rather than silently degrading to an identity clamp.

Scope / status (deliberately conservative):
- This does NOT unblock E2B/E4B multimodal. The existing validator that gates
  E2B/E4B image inputs is left in place, because clip-bounds are necessary but
  NOT sufficient for image parity on their own.
- Verified: the 448 bounds convert and load with real finite values; the clamp
  math is bit-exact vs jnp.clip, dtype-preserving, and an exact no-op when
  disabled (checked at the helper level and on a real Gemma4EncoderBlock forward,
  where wide bounds reproduce the disabled-path output exactly and tight bounds
  bound the projected activations).
- End-to-end teacher-forced image parity is NOT yet achieved: with clip-bounds
  loaded, the text path is exact but the image span still diverges from the HF
  reference, because the current Gemma-4 vision forward is missing other pieces
  of the reference contract (e.g. pad-patch attention masking and external image
  position threading). Those are out of scope for this change and tracked
  separately; this PR lands the clip-bounds building block on its own.
---
 .../utils/param_mapping.py                    |  21 +-
 src/maxtext/configs/base.yml                  |   6 +
 src/maxtext/configs/types.py                  |   9 +
 src/maxtext/models/gemma4_vision.py           | 234 +++++++++++++++++-
 4 files changed, 267 insertions(+), 3 deletions(-)

diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py
index 26359cdddc..22f80d9e09 100644
--- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py
+++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py
@@ -3107,7 +3107,9 @@ def GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers
           }
       )
 
-  # TODO: gemma4-small multimodal not yet supported — vision-encoder mappings below are dead.
+  # Gemma-4 E2B/E4B vision-encoder param mapping. Active when use_multimodal is set;
+  # the clipped-linears activation clip bounds are additionally mapped when
+  # use_clipped_linears_for_vit is enabled (required for image parity on E2B/E4B).
   if maxtext_config.use_multimodal and vcfg:
     nvis = vcfg.get("num_hidden_layers", 0)
     mapping.update(
@@ -3163,6 +3165,23 @@ def GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers
               f"{prefix}-mlp-wo-kernel": f"{hf_prefix}.mlp.down_proj.linear.weight",
           }
       )
+      # Gemma-4 vision clipped-linears: per-projection activation clip bounds
+      # (scalar {input,output}_{min,max}) carried in the reference checkpoint.
+      # Only mapped when the clipped-linears path is enabled; the nnx leaves live
+      # at <proj>_clip.{input,output}_{min,max} under attention/mlp.
+      if getattr(maxtext_config, "use_clipped_linears_for_vit", False):
+        _clip_proj = {
+            "attention-q_clip": f"{hf_prefix}.self_attn.q_proj",
+            "attention-k_clip": f"{hf_prefix}.self_attn.k_proj",
+            "attention-v_clip": f"{hf_prefix}.self_attn.v_proj",
+            "attention-o_clip": f"{hf_prefix}.self_attn.o_proj",
+            "mlp-gate_clip": f"{hf_prefix}.mlp.gate_proj",
+            "mlp-up_clip": f"{hf_prefix}.mlp.up_proj",
+            "mlp-down_clip": f"{hf_prefix}.mlp.down_proj",
+        }
+        for mt_sub, hf_proj in _clip_proj.items():
+          for bound in ("input_min", "input_max", "output_min", "output_max"):
+            mapping[f"{prefix}-{mt_sub}-{bound}"] = f"{hf_proj}.{bound}"
 
   return {k: v for k, v in mapping.items() if v is not None}
 
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 7a725fc4ca..54b019aee8 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -1188,6 +1188,12 @@ freeze_vision_encoder_params: true
 freeze_audio_encoder_params: true
 dtype_mm: "float32"  # Data type for multimodal model's vision encoder
 remat_policy_for_vit: "minimal"  # Remat policy for multimodal model's vision encoder. Check `remat_policy` for options.
+# Gemma-4 vision only: apply the per-projection activation clip bounds carried in the
+# reference checkpoint (self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj each have
+# scalar {input,output}_{min,max}). A prerequisite for Gemma-4 E2B/E4B image parity
+# (necessary but not on its own sufficient); no-op for other vision encoders. Bounds are
+# checkpoint-resident, non-trainable scalars.
+use_clipped_linears_for_vit: false
 image_size_for_vit: 896 # Default for Gemma3, and should be overwritten by model's config
 image_path: "" # Local image path used for decoding, can be multiple paths separated by comma, exp "/path/image1.jpg,/path/image2.jpg"
 video_path: "" # Local video path used for decoding, can be multiple paths separated by comma, exp "/path/video1.mp4,/path/video2.mp4"
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 769a548422..8f33482afb 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -2137,6 +2137,15 @@ class MultimodalGeneral(BaseModel):
 
   use_multimodal: bool = Field(False, description="Enable multimodal capabilities.")
   attention_for_vit: str = Field("dot_product", description="The attention algorithm to use for vision encoder.")
+  use_clipped_linears_for_vit: bool = Field(
+      False,
+      description=(
+          "Gemma-4 vision only: apply the per-projection activation clip bounds carried in the reference "
+          "checkpoint (self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj each have scalar "
+          "{input,output}_{min,max}). A prerequisite for Gemma-4 E2B/E4B image parity (necessary but not "
+          "on its own sufficient); no-op for other encoders."
+      ),
+  )
   vision_encoder_block: VisionEncoderBlockType = Field(
       VisionEncoderBlockType.NONE,
       description="The style of VisionEncoderBlock to use (e.g., 'gemma3', 'llama4').",
diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py
index 72cd841f81..d6fcbb5ad4 100644
--- a/src/maxtext/models/gemma4_vision.py
+++ b/src/maxtext/models/gemma4_vision.py
@@ -16,6 +16,8 @@
 """Vision transformer implementation for Gemma4."""
 
 from typing import cast
+import functools
+import operator
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
@@ -30,6 +32,108 @@
 from maxtext.layers import normalizations
 
 
+# =============================================================================
+# Gemma-4 vision clipped-linears (Navi upstream contribution)
+# -----------------------------------------------------------------------------
+# The Gemma-4 E2B/E4B vision tower ships per-projection activation clip bounds in
+# the reference (HF) checkpoint: for each of the 7 vision projections
+# (self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj) in each of the 16 encoder
+# blocks, a scalar {input_min,input_max,output_min,output_max} = 16*7*4 = 448
+# checkpoint tensors. The reference forward clamps each projection's input by
+# [input_min,input_max] and its output by [output_min,output_max]. Omitting the
+# clamps produces large activation drift in the image span (empirically KL 4-17
+# on a 340-token teacher-forced parity harness), because a handful of vision
+# activations blow up without the trained saturation. Upstream MaxText marks
+# E2B/E4B multimodal "not yet supported" and does not model these bounds.
+#
+# This module adds them as OPT-IN, checkpoint-resident, NON-TRAINABLE scalars,
+# gated on ``config.use_clipped_linears_for_vit`` (exact no-op when False).
+# Design: plain ``nnx.Param`` bounds (so they map to the canonical ``params``
+# collection and round-trip through the nnx->linen->orbax checkpoint path), plus
+# a leaf-PATH optimizer-freeze mask so the 448 bounds are excluded from optimizer
+# updates and weight decay without a custom nnx.Variable subclass (subclasses are
+# renamed by the linen bridge and silently dropped from the saved checkpoint).
+# A NaN sentinel + ``validate_clip_bounds`` hard-fails on a missing/non-finite
+# bound rather than silently degrading to an identity clamp.
+
+# Leaf-name tokens that identify a clip-bound scalar in a flattened params tree.
+_CLIP_LEAF_TOKENS = ("q_clip", "k_clip", "v_clip", "o_clip", "gate_clip", "up_clip", "down_clip")
+_CLIP_BOUND_NAMES = ("input_min", "input_max", "output_min", "output_max")
+
+
+def _mk_clip_bound(init_val=jnp.nan):
+  """A checkpoint-resident scalar clip bound as a plain ``nnx.Param`` (maps to the
+  canonical ``params`` collection; NaN sentinel marks an unloaded bound)."""
+  return nnx.Param(jnp.asarray(init_val, dtype=jnp.float32))
+
+
+def _is_clip_bound_path(path) -> bool:
+  """True iff a flattened-params key path addresses a clip-bound scalar."""
+  s = "/".join(str(getattr(p, "key", p)) for p in path) if not isinstance(path, str) else path
+  return any(tok in s for tok in _CLIP_LEAF_TOKENS) and any(b in s for b in _CLIP_BOUND_NAMES)
+
+
+def clip_optimizer_freeze_mask(params_tree):
+  """Bool pytree (same structure as ``params_tree``): True for TRAINABLE leaves,
+  False for the immutable clip bounds. Feed to ``optax.masked``/``multi_transform``
+  so the bounds get ``set_to_zero()`` updates. Path-based, so it survives the
+  nnx->linen->orbax round-trip regardless of leaf type erasure."""
+  flat = jax.tree_util.tree_flatten_with_path(params_tree)[0]
+  leaves_mask = [not _is_clip_bound_path(path) for path, _ in flat]
+  treedef = jax.tree_util.tree_structure(params_tree)
+  return jax.tree_util.tree_unflatten(treedef, leaves_mask)
+
+
+def _clip_in(x, cb):
+  """clamp(x, input_min, input_max) in x's dtype; no-op if ``cb`` is None."""
+  if cb is None:
+    return x
+  xd = x.dtype
+  return jnp.clip(x, cb.input_min.value.astype(xd), cb.input_max.value.astype(xd))
+
+
+def _clip_out(y, cb):
+  """clamp(y, output_min, output_max) in y's dtype; no-op if ``cb`` is None."""
+  if cb is None:
+    return y
+  yd = y.dtype
+  return jnp.clip(y, cb.output_min.value.astype(yd), cb.output_max.value.astype(yd))
+
+
+class _ClipBounds(nnx.Module):
+  """Holds the four checkpoint-resident scalar clip bounds as ``nnx.Param`` leaves."""
+
+  def __init__(self):
+    self.input_min = _mk_clip_bound()
+    self.input_max = _mk_clip_bound()
+    self.output_min = _mk_clip_bound()
+    self.output_max = _mk_clip_bound()
+
+
+def _make_clip_state():
+  """Four NaN-sentinel scalar bounds (checkpoint-resident, non-trainable)."""
+  return _ClipBounds()
+
+
+def validate_clip_bounds(cb, where=""):
+  """Hard-fail: every bound finite + scalar. Raises ValueError otherwise. No-op if ``cb`` is None."""
+  if cb is None:
+    return
+  for nm in ("input_min", "input_max", "output_min", "output_max"):
+    v = getattr(cb, nm).value
+    if getattr(v, "shape", ()) not in ((), (1,)):
+      raise ValueError(f"Gemma4 vision clip bound '{nm}'{(' @ '+where) if where else ''} has non-scalar "
+                       f"shape {v.shape}; expected scalar.")
+    fv = float(jnp.reshape(v, (-1,))[0])
+    if not bool(jnp.isfinite(jnp.asarray(fv))):
+      raise ValueError(f"Gemma4 vision clip bound '{nm}'{(' @ '+where) if where else ''} = {fv} is non-finite "
+                       f"(missing/NaN/Inf). use_clipped_linears_for_vit=True declares a FINITE clipped model; "
+                       f"refusing to fall back to an identity clamp.")
+
+
+
+
+
 def factorized_posemb(posemb: jax.Array, positions_xy: jax.Array, precision) -> jax.Array:
   """Computes factorized position embedding from (x, y) coordinates.
 
@@ -409,7 +513,15 @@ def __call__(self, inputs: jax.Array, positions: jax.Array) -> jax.Array:
 
 
 class Gemma4Attention(attentions.Attention):
-  """Gemma 4 specific Attention module."""
+  """Gemma 4 specific Attention module.
+
+  When ``use_clipped_linears`` is enabled, the q/k/v/o projections apply the
+  per-projection activation clip bounds carried in the Gemma-4 vision checkpoint
+  (input clamp before the matmul, output clamp after). The clamps are wired by
+  overriding the base ``Attention`` projection methods, so the underlying
+  ``DenseGeneral`` weights and their checkpoint key paths are unchanged. When the
+  flag is off, every override is an exact delegate to the base implementation.
+  """
 
   def init_rotary_embedding(self) -> Gemma4VisionRotaryEmbedding:
     """Initializes the rotary position embedding module for Gemma 4 vision."""
@@ -418,6 +530,116 @@ def init_rotary_embedding(self) -> Gemma4VisionRotaryEmbedding:
         rotary_fraction=None,  # Or assume it from config if available
     )
 
+  def enable_vision_clip_bounds(self):
+    """Attach the four checkpoint-resident clip-bound scalars for each of q/k/v/o.
+
+    Called once by ``Gemma4EncoderBlock`` after construction when
+    ``config.use_clipped_linears_for_vit`` is set. Idempotent.
+    """
+    if getattr(self, "_use_clipped_linears", False):
+      return
+    self._use_clipped_linears = True
+    self.q_clip = _make_clip_state()
+    self.k_clip = _make_clip_state()
+    self.v_clip = _make_clip_state()
+    self.o_clip = _make_clip_state()
+
+  def validate_clip_bounds(self):
+    if not getattr(self, "_use_clipped_linears", False):
+      return
+    validate_clip_bounds(self.q_clip, "q_proj")
+    validate_clip_bounds(self.k_clip, "k_proj")
+    validate_clip_bounds(self.v_clip, "v_proj")
+    validate_clip_bounds(self.o_clip, "o_proj")
+
+  # --- projection overrides: clamp(input) -> DenseGeneral -> clamp(output) ---
+  def query_projection(self, inputs_q, out_sharding=None):
+    if not getattr(self, "_use_clipped_linears", False):
+      return super().query_projection(inputs_q, out_sharding=out_sharding)
+    x = _clip_in(inputs_q, self.q_clip)
+    y = self.query(x, out_sharding=out_sharding)
+    return _clip_out(y, self.q_clip)
+
+  def kv_projection(self, inputs_kv, proj_name, out_sharding=None):
+    if not getattr(self, "_use_clipped_linears", False):
+      return super().kv_projection(inputs_kv, proj_name=proj_name, out_sharding=out_sharding)
+    if proj_name == "key":
+      cb, module = self.k_clip, self.key
+    elif proj_name == "value":
+      cb, module = self.v_clip, self.value
+    else:
+      raise ValueError(f"proj_name must be 'key' or 'value', but got {proj_name}")
+    x = _clip_in(inputs_kv, cb)
+    y = module(x, out_sharding=out_sharding)
+    return _clip_out(y, cb)
+
+  def out_projection(self, out, out_sharding=None):
+    if not getattr(self, "_use_clipped_linears", False):
+      return super().out_projection(out, out_sharding=out_sharding)
+    x = _clip_in(out, self.o_clip)
+    y = self.out(x, out_sharding=out_sharding)
+    return _clip_out(y, self.o_clip)
+
+
+class Gemma4ClippedMlpBlock(linears.MlpBlock):
+  """MlpBlock that applies the Gemma-4 vision per-projection activation clip bounds
+  to the gate (wi_0), up (wi_1) and down (wo) projections.
+
+  Only the non-fused activation path is supported (E2B/E4B use
+  ``activations=("gelu", "linear")`` with ``fused_mlp=False``), because gate and up
+  carry distinct clip bounds. When ``use_clipped_linears`` is off, this delegates to
+  the base ``MlpBlock`` unchanged.
+  """
+
+  def __init__(self, *args, use_clipped_linears=False, **kwargs):
+    super().__init__(*args, **kwargs)
+    self._use_clipped_linears = bool(use_clipped_linears)
+    if self._use_clipped_linears:
+      self.gate_clip = _make_clip_state()  # wi_0
+      self.up_clip = _make_clip_state()    # wi_1
+      self.down_clip = _make_clip_state()  # wo
+
+  def validate_clip_bounds(self):
+    if not self._use_clipped_linears:
+      return
+    validate_clip_bounds(self.gate_clip, "gate_proj")
+    validate_clip_bounds(self.up_clip, "up_proj")
+    validate_clip_bounds(self.down_clip, "down_proj")
+
+  def __call__(self, inputs, decode=False, deterministic=False,
+               intermediate_sharding=None, out_sharding=None):
+    if not self._use_clipped_linears:
+      return super().__call__(inputs, decode=decode, deterministic=deterministic,
+                              intermediate_sharding=intermediate_sharding, out_sharding=out_sharding)
+    cfg = self.config
+    if getattr(cfg, "fused_mlp", False):
+      # Clipped vision MLP requires the unfused path so gate/up get their own output clamps.
+      raise ValueError("Gemma4ClippedMlpBlock requires fused_mlp=False (per-projection clip bounds).")
+    if self.mlp_layer_norm is not None:
+      inputs = self.mlp_layer_norm(inputs)
+    clips = [self.gate_clip, self.up_clip]  # order matches activations ("gelu", "linear") == (gate, up)
+    activations = []
+    for idx, act_fn in enumerate(self.activations):
+      dense_name = "wi" if len(self.activations) == 1 else f"wi_{idx}"
+      module = getattr(self, dense_name)
+      x = _clip_in(inputs, clips[idx])
+      x = module(x, out_sharding=intermediate_sharding)
+      x = _clip_out(x, clips[idx])
+      x = linears.checkpoint_name(x, "mlp" + dense_name)
+      if cfg.activations_in_float32:
+        x = x.astype(jnp.float32)
+      x = linears._convert_to_activation_function(act_fn)(x)
+      activations.append(x)
+    x = functools.reduce(operator.mul, activations).astype(self.dtype)
+    x = self.dropout(x, deterministic=deterministic)
+    x = self._maybe_shard_with_logical(x, self.intermediate_logical)
+    x = _clip_in(x, self.down_clip)
+    output = self.wo(x, out_sharding=out_sharding)
+    output = _clip_out(output, self.down_clip)
+    output = linears.checkpoint_name(output, "mlpwo")
+    return output
+
+
 
 class Gemma4EncoderBlock(nnx.Module):
   """Single transformer encoder block (MHSA + MLP)."""
@@ -487,6 +709,11 @@ def __init__(
         is_vision=True,
         rngs=self.rngs,
     )
+    # Opt-in Gemma-4 vision clipped-linears: attach the q/k/v/o checkpoint-resident
+    # clip bounds and route the projections through the clamp-in/clamp-out overrides.
+    self._use_clipped_linears = bool(getattr(config, "use_clipped_linears_for_vit", False))
+    if self._use_clipped_linears:
+      self.attention.enable_vision_clip_bounds()
 
     self.pre_ffw_norm = normalizations.RMSNorm(
         num_features=config.hidden_size_for_vit,
@@ -506,7 +733,9 @@ def __init__(
         rngs=self.rngs,
     )
 
-    self.mlp = linears.MlpBlock(
+    mlp_cls = Gemma4ClippedMlpBlock if self._use_clipped_linears else linears.MlpBlock
+    mlp_kwargs = {"use_clipped_linears": True} if self._use_clipped_linears else {}
+    self.mlp = mlp_cls(
         config=config,
         mesh=mesh,
         in_features=config.hidden_size_for_vit,
@@ -516,6 +745,7 @@ def __init__(
         weight_dtype=config.weight_dtype,
         intermediate_dropout_rate=config.dropout_rate,
         rngs=self.rngs,
+        **mlp_kwargs,
     )
 
   def __call__(self, x: jax.Array, positions: jax.Array | None = None, deterministic: bool = False) -> jax.Array:

From 084ef28294743695a787bb2db906ad9344d6b527 Mon Sep 17 00:00:00 2001
From: dengcchi <dengcchi@meta.com>
Date: Sat, 8 Aug 2026 03:12:57 -0700
Subject: [PATCH 2/3] Gemma-4 E2B/E4B: close image parity (padded-patch masking
 + PLE image-row substitution)

Builds on the vision clipped-linears to make Gemma-4 E2B/E4B image inputs match the HF
reference end to end. The clipped-linears alone are necessary but not sufficient; the
reference contract also needs the vision padded-patch handling and a decoder-side per-layer
embedding (PLE) fix.

Vision (models/gemma4_vision.py):
- Gemma4EncoderBlock threads decoder_segment_ids into attention (valid=1 / pad=2) so the
  phantom padded patches are masked out of vision self-attention.
- Gemma4VisionEncoderLayer gains a padded-patch path (image_position_ids != None): consume
  pre-patchified patches + real per-patch positions (-1 = pad), build the segment ids, pool
  by the real positions, and return the pooled-token validity mask. image_position_ids is
  None -> byte-identical legacy path.

Threading (layers/encoders.py, models/models.py):
- VisionEncoder / the model forward thread encoder_image_position_ids to the Gemma-4 vision
  encoder and route the returned validity mask into MultimodalInput.image_masks, so exactly
  the valid pooled tokens land in the image placeholders (merge_mm_embeddings.token_masks).

Decoder (layers/nnx_decoders.py):
- ple_pad_substitute_image_rows: Gemma-4 E2B/E4B build the per-layer inputs from llm_input_ids
  with image placeholder tokens mapped to pad_token_id (HF modeling_gemma4), rather than
  feeding the placeholder id into the PLE path. Without this the per-layer embeddings at the
  image positions diverge, corrupting the image-span and post-image logits.
- use_bidirectional_image_attn: E2B/E4B image spans are causal; suppress the bidirectional
  attention carve-out unless explicitly enabled (bidirectional-image models set it True).

Config (configs/types.py, configs/base.yml):
- Adds use_bidirectional_image_attn, ple_pad_substitute_image_rows, ple_pad_mode,
  image_placeholder_token_id, ple_pad_token_id (defaults preserve behavior for other models).
- Allows E2B/E4B multimodal when use_clipped_linears_for_vit is set.

Validation (CPU, teacher-forced 340-token forward vs HF reference logits, converted E2B
checkpoint): with the full fix and clip-bounds enabled,
  pre_image  max_KL 2.7e-5, argmax 1.0
  image_span max_KL 9.0e-5, argmax 1.0
  post_image max_KL 4.3e-5, argmax 1.0   (frozen gate: <= 1.26e-3 and argmax >= 0.995)
matching the reference. With clip-bounds disabled the image span diverges, confirming the
clip-bounds are required. Defaults keep all changes off for non-Gemma-4 models.
---
 src/maxtext/configs/base.yml        |   6 ++
 src/maxtext/configs/types.py        |  42 +++++++++-
 src/maxtext/layers/encoders.py      |  12 ++-
 src/maxtext/layers/nnx_decoders.py  |  23 ++++-
 src/maxtext/models/gemma4_vision.py | 126 ++++++++++++++++++++++++----
 src/maxtext/models/models.py        |  22 +++--
 6 files changed, 198 insertions(+), 33 deletions(-)

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 54b019aee8..3b61b967f4 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -1194,6 +1194,12 @@ remat_policy_for_vit: "minimal"  # Remat policy for multimodal model's vision en
 # (necessary but not on its own sufficient); no-op for other vision encoders. Bounds are
 # checkpoint-resident, non-trainable scalars.
 use_clipped_linears_for_vit: false
+# Gemma-4 E2B/E4B decoder image-handling (defaults preserve behavior for other models):
+use_bidirectional_image_attn: false  # E2B/E4B image spans are causal
+ple_pad_substitute_image_rows: false  # substitute pad id for image rows in the per-layer-embedding path (HF gemma4)
+ple_pad_mode: "identity"  # 'identity' (token-id path) or 'both' (also context/embedding path)
+image_placeholder_token_id: 258880  # GEMMA4_TOKEN_PLACEHOLDER
+ple_pad_token_id: 0  # E2B text_config.pad_token_id
 image_size_for_vit: 896 # Default for Gemma3, and should be overwritten by model's config
 image_path: "" # Local image path used for decoding, can be multiple paths separated by comma, exp "/path/image1.jpg,/path/image2.jpg"
 video_path: "" # Local video path used for decoding, can be multiple paths separated by comma, exp "/path/video1.mp4,/path/video2.mp4"
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 8f33482afb..712c042c85 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -2146,6 +2146,34 @@ class MultimodalGeneral(BaseModel):
           "on its own sufficient); no-op for other encoders."
       ),
   )
+  use_bidirectional_image_attn: bool = Field(
+      False,
+      description=(
+          "Whether image placeholder tokens attend bidirectionally in the text decoder. Gemma-4 E2B/E4B "
+          "use causal image spans (False); bidirectional-image models (Gemma-3, gemma4-26b/31b) use True."
+      ),
+  )
+  ple_pad_substitute_image_rows: bool = Field(
+      False,
+      description=(
+          "Gemma-4 E2B/E4B per-layer-embedding (PLE) path: substitute ple_pad_token_id for image placeholder "
+          "rows before the per-layer embedder, matching HF modeling_gemma4 (llm_input_ids pad substitution). "
+          "Default False preserves the native PLE for other models."
+      ),
+  )
+  ple_pad_mode: str = Field(
+      "identity",
+      description=(
+          "PLE pad-substitution scope when ple_pad_substitute_image_rows=True: 'identity' (token-id path only) "
+          "or 'both' (also substitute the pad embedding in the context path)."
+      ),
+  )
+  image_placeholder_token_id: int = Field(
+      258880, description="Gemma-4 image placeholder token id (GEMMA4_TOKEN_PLACEHOLDER)."
+  )
+  ple_pad_token_id: int = Field(
+      0, description="Pad token id used for PLE image-row substitution (Gemma-4 E2B text_config.pad_token_id=0)."
+  )
   vision_encoder_block: VisionEncoderBlockType = Field(
       VisionEncoderBlockType.NONE,
       description="The style of VisionEncoderBlock to use (e.g., 'gemma3', 'llama4').",
@@ -3581,16 +3609,22 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           f"{self.model_name} requires scan_layers=False (per-layer KV sharing is incompatible with nn.scan)."
       )
     if self.use_multimodal:
-      # Gemma 4 small (E2B / E4B) only supports text for now; multimodal
-      # support is pending clipped-linears in the vision encoder.
-      if self.model_name in ("gemma4-e2b", "gemma4-e4b"):
-        raise ValueError(f"Multimodal is not yet supported for {self.model_name}; only text inputs are supported.")
+      # Gemma 4 small (E2B / E4B) multimodal requires the vision-encoder clipped-linears AND the
+      # padded-patch masking / position-threading path; gate on the clipped-linears flag.
+      if self.model_name in ("gemma4-e2b", "gemma4-e4b") and not self.use_clipped_linears_for_vit:
+        raise ValueError(
+            f"Multimodal for {self.model_name} requires use_clipped_linears_for_vit=True "
+            "(the vision encoder ships per-projection activation clip bounds; without them the "
+            "image span diverges). Set use_clipped_linears_for_vit=True to enable image inputs."
+        )
       valid_mm_models = (
           "gemma3-4b",
           "gemma3-12b",
           "gemma3-27b",
           "gemma4-26b",
           "gemma4-31b",
+          "gemma4-e2b",
+          "gemma4-e4b",
           "llama4-17b-16e",
           "llama4-17b-128e",
           "qwen3-omni-30b-a3b",
diff --git a/src/maxtext/layers/encoders.py b/src/maxtext/layers/encoders.py
index e5e32d044a..a5a0c38ff8 100644
--- a/src/maxtext/layers/encoders.py
+++ b/src/maxtext/layers/encoders.py
@@ -106,13 +106,21 @@ def _setup_vision_encoder_layers(self):
 
     return encoder_name, projector_name
 
-  def __call__(self, input_images, input_masks=None, video_grid_thw=None, deterministic=False):
+  def __call__(self, input_images, input_masks=None, video_grid_thw=None, deterministic=False,
+               image_position_ids=None):
     # vision encoder output, frozen params in many cases
     encoder = getattr(self, self.encoder_name)
+    vision_image_masks = None
     if self.vision_encoder_block.value.startswith("qwen3") and input_masks is not None:
       encoder_output = encoder(
           input_images, video_mask=input_masks, video_grid_thw=video_grid_thw, deterministic=deterministic
       )
+    elif self.vision_encoder_block == VisionEncoderBlockType.GEMMA4 and image_position_ids is not None:
+      # Gemma-4 padded-patch path: pre-patchified patches + per-patch positions (-1 = pad). The
+      # encoder returns (embeddings, image_masks); the mask marks the valid pooled tokens.
+      encoder_output = encoder(input_images, deterministic=deterministic, image_position_ids=image_position_ids)
+      embeddings, vision_image_masks = encoder_output
+      encoder_output = embeddings
     else:
       encoder_output = encoder(input_images, deterministic=deterministic)
     deep_feats = None
@@ -131,7 +139,7 @@ def __call__(self, input_images, input_masks=None, video_grid_thw=None, determin
     projector = getattr(self, self.projector_name)
     embeddings = projector(embeddings)
 
-    return embeddings, deep_feats
+    return embeddings, deep_feats, vision_image_masks
 
 
 class MultimodalMLPProjector(nnx.Module):
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 895ea27c14..51b32ea112 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -2231,10 +2231,31 @@ def _apply_gemma4_small_layers(
     """Apply Gemma 4 small (E2B/E4B) decoder layers (pure-NNX)."""
     cfg = self.config
     bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None
+    # Gemma-4 E2B/E4B image spans are causal (text_config.use_bidirectional_attention is unset), unlike
+    # the bidirectional-image Gemma-3 / 26B / 31B models. Suppress the bidirectional attention carve-out
+    # unless explicitly enabled via config (gate on a flag, not the model name).
+    if not bool(getattr(cfg, "use_bidirectional_image_attn", False)):
+      bidirectional_mask_value = None
 
     per_layer_inputs = None
     if cfg.hidden_size_per_layer_input > 0 and cfg.vocab_size_per_layer_input > 0:
-      per_layer_inputs = self.per_layer_embedder(decoder_input_tokens, y)
+      ple_tokens = decoder_input_tokens
+      ple_context = y
+      # Gemma-4 E2B/E4B build the per-layer inputs from llm_input_ids with the image placeholder
+      # tokens mapped to pad_token_id (HF modeling_gemma4.py), rather than feeding the image
+      # placeholder id / merged image features into the PLE path. Without this substitution the
+      # per-layer embeddings at the image placeholder positions diverge from the reference, which
+      # corrupts the image-span and post-image logits. Gated on ple_pad_substitute_image_rows
+      # (default False preserves the native PLE for other models).
+      if bool(getattr(cfg, "ple_pad_substitute_image_rows", False)) and multimodal_input is not None:
+        _img_id = int(getattr(cfg, "image_placeholder_token_id", 258880))
+        _pad_id = int(getattr(cfg, "ple_pad_token_id", 0))
+        _img_row = decoder_input_tokens.astype(jnp.int32) == _img_id
+        ple_tokens = jnp.where(_img_row, _pad_id, decoder_input_tokens.astype(jnp.int32))
+        if str(getattr(cfg, "ple_pad_mode", "identity")) == "both" and hasattr(self, "shared_embedding"):
+          _pad_vec = self.shared_embedding(jnp.full_like(decoder_input_tokens, _pad_id).astype(jnp.int32))
+          ple_context = jnp.where(_img_row[..., None], _pad_vec, y)
+      per_layer_inputs = self.per_layer_embedder(ple_tokens, ple_context)
 
     layer_types = gemma4_small.build_layer_types(cfg.num_decoder_layers, cfg.model_name)
     num_kv_shared = cfg.num_kv_shared_layers
diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py
index d6fcbb5ad4..60ab571f85 100644
--- a/src/maxtext/models/gemma4_vision.py
+++ b/src/maxtext/models/gemma4_vision.py
@@ -748,11 +748,29 @@ def __init__(
         **mlp_kwargs,
     )
 
-  def __call__(self, x: jax.Array, positions: jax.Array | None = None, deterministic: bool = False) -> jax.Array:
-    """Applies the encoder block (MHSA + MLP) to the inputs."""
+  def __call__(
+      self,
+      x: jax.Array,
+      positions: jax.Array | None = None,
+      deterministic: bool = False,
+      decoder_segment_ids: jax.Array | None = None,
+  ) -> jax.Array:
+    """Applies the encoder block (MHSA + MLP) to the inputs.
+
+    When ``decoder_segment_ids`` is provided, patches carrying distinct segment
+    ids cannot attend to each other. This is used by the padded-patch path
+    (valid patches = segment 1, padded/sentinel patches = segment 2) so that the
+    phantom pad patches are masked out of vision self-attention.
+    """
     x_normed = self.pre_attention_norm(x)
-    # Pass positions to attention for RoPE
-    x_attn, _ = self.attention(x_normed, x_normed, inputs_positions=positions, deterministic=deterministic)
+    # Pass positions to attention for RoPE (+ optional segment mask for padded patches).
+    x_attn, _ = self.attention(
+        x_normed,
+        x_normed,
+        inputs_positions=positions,
+        decoder_segment_ids=decoder_segment_ids,
+        deterministic=deterministic,
+    )
     x_attn = self.post_attention_norm(x_attn)
     x_after_attn = x_attn + x
 
@@ -803,34 +821,106 @@ def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs):
         nnx.initializers.ones(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype), sharding=(None,)
     )
 
-  def __call__(self, inputs: jax.Array, deterministic: bool = False) -> jax.Array:
-    """Applies the vision encoder layer."""
+  def __call__(
+      self,
+      inputs: jax.Array,
+      deterministic: bool = False,
+      image_position_ids: jax.Array | None = None,
+  ):
+    """Applies the vision encoder layer.
+
+    Two contracts:
+
+    (A) Legacy all-valid (``image_position_ids is None``): ``inputs`` are raw images
+        [B, N, H, W, C] (or [B, H, W, C]); patchify -> full unmasked attention -> pool by the
+        derived positions -> return embeddings only (4D array [B, N, K, D]).
+
+    (B) Padded-patch dynamic-N (``image_position_ids is not None``): ``inputs`` are ALREADY
+        patchified pixel_values with shape [B, L, P*P*C] (or [B, N, L, P*P*C]) and
+        ``image_position_ids`` is [B, L, 2] (or [B, N, L, 2]) with -1 sentinel rows marking padded
+        patches. The pre-patchified patches + REAL positions are fed to VisionEntry, per-patch
+        ``decoder_segment_ids`` (valid=1, pad=2) mask the phantom pad patches out of self-attention,
+        pooling uses the real positions (``avg_pool_by_positions`` maps a -1 patch to a zero-weight
+        bucket), and the VisionExit validity mask is returned. Returns a 2-tuple
+        ``(embeddings[B, N, K, D], image_masks[B*N, K])`` where ``image_masks.sum()`` is the number
+        of valid pooled tokens, threaded to ``merge_mm_embeddings.token_masks`` so exactly the valid
+        pooled tokens land in the image placeholders.
+    """
+    if image_position_ids is None:
+      # ---- Legacy path: raw images -> patchify -> full (unmasked) attention ----
+      if inputs.ndim == 4:
+        inputs = jnp.expand_dims(inputs, 1)
+      b, n, h, w, c = inputs.shape
+      inputs_flat = jnp.reshape(inputs, (b * n, h, w, c))
+
+      x, positions_xy = self.vision_entry(inputs_flat)
+
+      for i in range(self.config.num_hidden_layers_for_vit):
+        layer = getattr(self, f"layer_{i}")
+        x = layer(x, positions=positions_xy, deterministic=deterministic)
+
+      vision_exit_results = self.vision_exit(x, positions_xy=positions_xy)
+      (embeddings, _) = vision_exit_results[0]
+
+      embeddings = (embeddings - self.std_bias.value.astype(embeddings.dtype)) * self.std_scale.value.astype(
+          embeddings.dtype
+      )
+
+      # Unflatten batch and num_images
+      final_x = jnp.reshape(embeddings, (b, n, embeddings.shape[1], embeddings.shape[2]))
+      return final_x
+
+    # ---- Padded-patch dynamic-N path: pre-patchified patches + sentinel positions ----
+    # inputs: [B, L, P*P*C] pre-patchified pixel_values. Support an optional per-image N dim
+    # [B, N, L, F] -> flatten to [B*N, L, F] to match positions.
     if inputs.ndim == 4:
-      inputs = jnp.expand_dims(inputs, 1)
-    b, n, h, w, c = inputs.shape
-    inputs_flat = jnp.reshape(inputs, (b * n, h, w, c))
+      b, n, l, f = inputs.shape
+      patches = jnp.reshape(inputs, (b * n, l, f))
+      pos = jnp.reshape(image_position_ids, (b * n, l, 2))
+    else:
+      assert inputs.ndim == 3, f"padded-patch path expects pre-patchified [B, L, F] patches, got {inputs.shape}"
+      b, l, f = inputs.shape
+      n = 1
+      patches = inputs
+      pos = image_position_ids
+      if pos.ndim == 2:
+        pos = jnp.broadcast_to(pos, (b, l, 2))
 
-    x, positions_xy = self.vision_entry(inputs_flat)
+    pos = pos.astype(jnp.int32)
+
+    # VisionEntry consumes pre-patchified patches + REAL positions (incl -1 sentinels).
+    x, positions_xy = self.vision_entry(patches, positions_xy=pos)
+
+    # Segment ids: valid patch (any coord != -1) -> 1, padded sentinel patch -> 2. Distinct segments
+    # cannot attend to each other, masking the phantom pad patches out of self-attention.
+    is_pad = (positions_xy == -1).all(axis=-1)  # [B*N, L]
+    decoder_segment_ids = jnp.where(is_pad, 2, 1).astype(jnp.int32)
 
     for i in range(self.config.num_hidden_layers_for_vit):
       layer = getattr(self, f"layer_{i}")
-      x = layer(x, positions=positions_xy, deterministic=deterministic)
+      x = layer(
+          x,
+          positions=positions_xy,
+          deterministic=deterministic,
+          decoder_segment_ids=decoder_segment_ids,
+      )
 
+    # Pool with REAL positions; avg_pool_by_positions returns (embeddings, validity_mask).
     vision_exit_results = self.vision_exit(x, positions_xy=positions_xy)
-
-    # Return embeddings from VisionExit tuple
-    # vision_exit_results is a tuple of (embeddings, mask) tuples, one for each output length.
-    # We take the first result.
-    (embeddings, _) = vision_exit_results[0]
+    (embeddings, image_masks) = vision_exit_results[0]  # embeddings [B*N, K, D], mask [B*N, K]
 
     embeddings = (embeddings - self.std_bias.value.astype(embeddings.dtype)) * self.std_scale.value.astype(
         embeddings.dtype
     )
 
-    # Unflatten batch and num_images
     final_x = jnp.reshape(embeddings, (b, n, embeddings.shape[1], embeddings.shape[2]))
+    if image_masks is None:
+      image_masks = jnp.ones((b * n, embeddings.shape[1]), dtype=jnp.int32)
+    else:
+      # merge_mm_embeddings does argsort(-token_mask); use int32 so negation/sort is well-defined.
+      image_masks = image_masks.astype(jnp.int32)
 
-    return final_x
+    return final_x, image_masks
 
 
 class Gemma4VisionProjector(nnx.Module):
diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py
index a70ad780a3..aab8f1bdab 100644
--- a/src/maxtext/models/models.py
+++ b/src/maxtext/models/models.py
@@ -129,6 +129,7 @@ def __call__(
       decoder_segment_ids=None,
       encoder_images: None | jnp.ndarray = None,
       encoder_image_masks: None | jnp.ndarray = None,
+      encoder_image_position_ids: None | jnp.ndarray = None,
       encoder_videos: None | jnp.ndarray = None,
       encoder_video_masks: None | jnp.ndarray = None,
       encoder_video_grid_thw: None | jnp.ndarray = None,
@@ -164,17 +165,19 @@ def __call__(
     video_embeddings = None
     audio_embeddings = None
     deepstack_visual_embeds = None
+    vision_image_masks = None
 
     if getattr(self.config, "use_multimodal", False) and encoder_images is not None:
-      image_embeddings, deepstack_visual_embeds = self.vision_encoder(  # pyrefly: ignore[not-callable]
-          input_images=encoder_images, deterministic=not enable_dropout
+      image_embeddings, deepstack_visual_embeds, vision_image_masks = self.vision_encoder(  # pyrefly: ignore[not-callable]
+          input_images=encoder_images, deterministic=not enable_dropout,
+          image_position_ids=encoder_image_position_ids,
       )
       bidirectional_mask_image = mm_processor.get_bidirectional_mask_vision(
           self.config, decoder_input_tokens, is_video=False
       )
 
     if getattr(self.config, "use_multimodal", False) and encoder_videos is not None:
-      video_embeddings, deepstack_visual_embeds = self.vision_encoder(  # pyrefly: ignore[not-callable]
+      video_embeddings, deepstack_visual_embeds, _ = self.vision_encoder(  # pyrefly: ignore[not-callable]
           input_images=encoder_videos,
           input_masks=encoder_video_masks,
           video_grid_thw=encoder_video_grid_thw,
@@ -197,7 +200,7 @@ def __call__(
     if image_embeddings is not None or video_embeddings is not None or audio_embeddings is not None:
       multimodal_input = MultimodalInput(
           image_embeddings=image_embeddings,
-          image_masks=encoder_image_masks,
+          image_masks=vision_image_masks if vision_image_masks is not None else encoder_image_masks,
           video_embeddings=video_embeddings,
           video_masks=encoder_video_masks,
           audio_embeddings=audio_embeddings,
@@ -449,6 +452,7 @@ def __call__(
       cache=None,
       encoder_images: jax.Array | None = None,
       encoder_image_masks: jax.Array | None = None,
+      encoder_image_position_ids: jax.Array | None = None,
       encoder_videos: jax.Array | None = None,
       encoder_video_masks: jax.Array | None = None,
       encoder_video_grid_thw: jax.Array | None = None,
@@ -499,16 +503,18 @@ def __call__(
     video_embeddings = None
     audio_embeddings = None
     deepstack_visual_embeds = None
+    vision_image_masks = None
     if getattr(self.config, "use_multimodal", False) and encoder_images is not None:
-      image_embeddings, deepstack_visual_embeds = self.vision_encoder(  # pyrefly: ignore[not-callable]
-          input_images=encoder_images, deterministic=not enable_dropout
+      image_embeddings, deepstack_visual_embeds, vision_image_masks = self.vision_encoder(  # pyrefly: ignore[not-callable]
+          input_images=encoder_images, deterministic=not enable_dropout,
+          image_position_ids=encoder_image_position_ids,
       )
       bidirectional_mask_image = mm_processor.get_bidirectional_mask_vision(
           self.config, decoder_input_tokens, is_video=False
       )
 
     if getattr(self.config, "use_multimodal", False) and encoder_videos is not None:
-      video_embeddings, deepstack_visual_embeds = self.vision_encoder(  # pyrefly: ignore[not-callable]
+      video_embeddings, deepstack_visual_embeds, _ = self.vision_encoder(  # pyrefly: ignore[not-callable]
           input_images=encoder_videos,
           input_masks=encoder_video_masks,
           video_grid_thw=encoder_video_grid_thw,
@@ -531,7 +537,7 @@ def __call__(
     if image_embeddings is not None or video_embeddings is not None or audio_embeddings is not None:
       multimodal_input = MultimodalInput(
           image_embeddings=image_embeddings,
-          image_masks=encoder_image_masks,
+          image_masks=vision_image_masks if vision_image_masks is not None else encoder_image_masks,
           video_embeddings=video_embeddings,
           video_masks=encoder_video_masks,
           audio_embeddings=audio_embeddings,

From 753b58f0d48371d399f72206b2d739d2583db235 Mon Sep 17 00:00:00 2001
From: dengcchi <dengcchi@meta.com>
Date: Sat, 8 Aug 2026 03:17:01 -0700
Subject: [PATCH 3/3] Gemma-4 E2B/E4B: set image-contract defaults in the model
 configs

So enabling image inputs only requires use_multimodal=true + use_clipped_linears_for_vit=true;
the E2B/E4B model configs supply the decoder image-contract flags (causal image spans, PLE
image-row pad substitution, placeholder/pad token ids). Also drops the stale 'multimodal not
yet supported' comment.
---
 src/maxtext/configs/models/gemma4-e2b.yml | 10 +++++++++-
 src/maxtext/configs/models/gemma4-e4b.yml | 10 +++++++++-
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/maxtext/configs/models/gemma4-e2b.yml b/src/maxtext/configs/models/gemma4-e2b.yml
index 81c8d4ea66..c4a4b26c26 100644
--- a/src/maxtext/configs/models/gemma4-e2b.yml
+++ b/src/maxtext/configs/models/gemma4-e2b.yml
@@ -44,7 +44,8 @@ global_rope_proportion: 0.25
 local_rope_proportion: 1.0
 final_logits_soft_cap: 30.0
 
-# Vision encoder flags — multimodal not yet supported for E2B / E4B.
+# Vision encoder flags. Image (multimodal) parity requires use_clipped_linears_for_vit=true
+# (set at runtime alongside use_multimodal=true); the flags below configure the E2B image contract.
 vision_encoder_block: "gemma4"
 rope_theta_for_vit: 100
 image_size_for_vit: [672, 960]
@@ -58,3 +59,10 @@ num_attention_heads_for_vit: 12
 image_placeholder: "<|image|>"
 vision_output_length: 280
 num_position_embeddings_for_vit: 10240
+# E2B image contract (decoder-side): image spans are causal and the per-layer-embedding path
+# substitutes the pad token for image placeholder rows (matches HF modeling_gemma4).
+use_bidirectional_image_attn: False
+ple_pad_substitute_image_rows: True
+ple_pad_mode: "identity"
+image_placeholder_token_id: 258880
+ple_pad_token_id: 0
diff --git a/src/maxtext/configs/models/gemma4-e4b.yml b/src/maxtext/configs/models/gemma4-e4b.yml
index b8da0d1d46..cd752ae22b 100644
--- a/src/maxtext/configs/models/gemma4-e4b.yml
+++ b/src/maxtext/configs/models/gemma4-e4b.yml
@@ -45,7 +45,8 @@ global_rope_proportion: 0.25
 local_rope_proportion: 1.0
 final_logits_soft_cap: 30.0
 
-# Vision encoder flags — multimodal not yet supported for E2B / E4B.
+# Vision encoder flags. Image (multimodal) parity requires use_clipped_linears_for_vit=true
+# (set at runtime alongside use_multimodal=true); the flags below configure the E4B image contract.
 vision_encoder_block: "gemma4"
 rope_theta_for_vit: 100
 image_size_for_vit: [672, 960]
@@ -59,3 +60,10 @@ num_attention_heads_for_vit: 12
 image_placeholder: "<|image|>"
 vision_output_length: 280
 num_position_embeddings_for_vit: 10240
+# E4B image contract (decoder-side): image spans are causal and the per-layer-embedding path
+# substitutes the pad token for image placeholder rows (matches HF modeling_gemma4).
+use_bidirectional_image_attn: False
+ple_pad_substitute_image_rows: True
+ple_pad_mode: "identity"
+image_placeholder_token_id: 258880
+ple_pad_token_id: 0