fix(gemma4): bidirectional vision attention for unified soft tokens — fixes color corruption (#127)

[Pushkinist]

Closes #127.

Problem

The encoder-free gemma4_unified (12B) vision path (added in #126/#120)
corrupted color: on solid-color PNGs the unified 12B misnamed green→black,
yellow→purple, white→black, while the standard gemma4 SigLIP path (e4b) named
all 5 correctly on the same images. (The original #120 proof used a synthetic
4-quadrant image and a "describe it" prompt, which was too soft to catch this.)

Root cause — the issue's preprocessing diagnosis was falsified

The filed diagnosis (channel order / mean-std / RGBA stride) is wrong. The
patchify + embedder were verified numerically bit-exact to the HF
Gemma4UnifiedVisionEmbedder reference (NumPy port, all digits). The real defect:
the text decoder ran the image soft tokens causally, but Gemma4 unified sets
use_bidirectional_attention: "vision" — every soft token in an image block must
attend to every other soft token of the same block. The SigLIP path (e4b/26b/31b)
hides this because its ViT pre-integrates the image before projection; the
encoder-free path projects raw patches with no pre-integrated context, so causal
masking scrambles them (green/yellow → black/purple, spatial layout hallucinated).

Fix (general)

• build_vision_bidi_overlay (gemma4/model.rs): builds a [1,1,seq,seq]
  additive overlay (0.0 = allowed) opening intra-image-block cells, keyed off the
  <start_of_image>/<end_of_image> markers, merged via element-wise maximum
  (combine_bidi_overlay) with each layer's causal/SWA prefill mask. Threaded as
  Option<&Array> through DecoderLayer/Attention; None (no-op, byte-identical)
  on text / audio / decode / speculative paths, and gated on has_image so a
  text prefill never even syncs ids to host.
• LayerNorm eps corrected to the PyTorch nn.LayerNorm default 1e-5.
• BOI/EOI marker ids promoted to a single pub const (was duplicated across crates).

Proof (real 12B-mxfp8, temp=0, identical PNGs)

Unified solid-color battery, BEFORE → AFTER: green Purple→Green, yellow
purple→Yellow, red Red→Red, blue blue→Blue = chromatic 4/4 correct. The
white-card image (white bg, red border, dark text) now reads "light gray
background, thin red border, black center" (structure correct) vs BEFORE
"corrupted pinks/mauves/grays". e4b SigLIP control unchanged (5/5). Unified text
("Paris") + unified audio paths un-regressed.

Documented limitation (not an rMLX bug)

A 100%-uniform solid white fill still reads as black, and this is inherent
to the architecture: patch_ln1 is a 6912-dim LayerNorm applied to the raw
patch as the first op, so any achromatic uniform patch has zero variance and
collapses to the bias — white ≡ gray ≡ black by construction (confirmed by the
weight shapes and the rust-reviewer). The real Gemma4 12B has the same property;
the SigLIP path escapes it only via its ViT. Real images (with inter-patch
variance) are unaffected. Documented in docs/MODELS.md.

Tests + guards

Model-free: patchify channel-preservation + interior-layout; bidi overlay opens
intra-block / does not cross blocks / absent without markers; SWA+overlay
override. make lint (-D warnings) + make test green. Reviewed by the
rust-reviewer agent (no CRITICAL/HIGH; all findings fixed in the second commit).
No new deps, no new env vars.

🤖 Generated with Claude Code