feat(xla): support Gemma3 VLM additive prefill masks

[inureyes]

Parent and dependency

• Parent epic: epic: OpenXLA multimodal / VLM architecture support (Qwen2-VL, Gemma3n, Phi4MM, Molmo, and more) #566
• Blocked by feat(xla): emit SigLIP/ViT encoders and projectors through IREE #863
• Existing MLX implementation: src/loading/vlm_gemma.rs, src/vision/merge.rs::prepare_inputs_for_multimodal, and src/models/gemma3.rs

Problem

Gemma3 VLM uses the reusable SigLIP/projector shape but not LLaVA's ordinary causal merge. The current MLX path scales projected image features by 1 / sqrt(hidden_size), merges them with text embeddings, then the Gemma3 model scales the complete hidden sequence by sqrt(hidden_size). It also supplies an additive f32 4D prefill mask (0.0 for allowed valid-token pairs, f32::MIN for padded pairs) to every language layer instead of constructing the normal causal/sliding masks.

If the XLA embeddings entry treats the prepared values as post-scale but the preprocessor exports the pre-scale MLX merge result, text/image magnitudes diverge. If it substitutes a causal mask, the model architecture diverges independently of the vision encoder.

Goal

Support a pinned Gemma3 VLM checkpoint end to end on XLA with exact embedding scaling, placeholder/padding semantics, additive 4D prefill attention, SigLIP/projector execution, and ordinary one-token decode.

Required implementation

• Extend the host preprocessor with the Gemma3 prompt wrapping and image-token expansion used by vlm_prompt.rs/vlm_gemma.rs, including configured BOI/image/newline token IDs and exact image-feature cardinality.
• Reuse the IREE SigLIP/projector primitives from feat(xla): emit SigLIP/ViT encoders and projectors through IREE #863 with Gemma3 checkpoint prefixes, shapes, normalization epsilon, image size, and projector mapping validated explicitly.
• Reproduce the net embedding scale at the canonical feat(xla): add a prefill-from-embeddings StableHLO entry #858 boundary: the tensor handed to prefill_embeddings.main must be post-Gemma embedding scale. Document and test that text positions receive sqrt(hidden_size) while the pre-divide/overall-scale combination leaves projected image features at the reference magnitude.
• Construct the exact additive f32 mask currently produced by prepare_inputs_for_multimodal, including padding rows/columns and canonical shape [1, 1, Lp, Lp] for one request. Do not replace it with a boolean or multiplicative mask.
• Pass the external mask to every Gemma3 prefill layer as the existing MLX path does. Decode remains the normal one-token no-mask path over the seeded KV cache.
• Validate, with an independent reference implementation, whether all valid prefill token pairs are bidirectional for the selected checkpoint/config. If upstream semantics differ by variant, encode the variant in configuration rather than guessing from image-token positions.
• Include mask mode, image token IDs, scaling policy, SigLIP/projector config, and Gemma sliding/global attention config in artifact compatibility identity.
• Integrate filtered loading, XLA capabilities, CLI/server batching, cancellation, and no-image text fallback without a full duplicate MLX decoder.

Non-goals

• Gemma3n MobileNet/per-layer inputs or audio.
• Gemma4 unified block masks.
• General arbitrary 4D masks from untrusted API callers.

Validation

• Pin a Gemma3 VLM checkpoint/revision and compare prompt tokens, image-token positions, raw/projected/scaled image features, text and image embedding norms, every mask cell for a small fixture, selected layer outputs, logits/KV, and greedy tokens against MLX and an independent oracle.
• Cover no padding, left/right padding as supported, one and multiple images, placeholder mismatch, all-valid mask, padded mask, no-image text fallback, and context overflow.
• Add negative fixtures for accidental causal masking, 0/1 multiplicative masking, double Gemma scaling, and missing image pre-scale; each must be caught at an intermediate assertion.
• Run CLI and mixed text/image continuous-batch server tests on a production-relevant IREE target.
• Run test(xla): validate a LLaVA reference architecture end to end #862/feat(xla): emit SigLIP/ViT encoders and projectors through IREE #863 and all text-only Gemma3 sliding/global attention regression gates.

Acceptance criteria

• Exported merged embeddings have the exact post-scale magnitudes required by feat(xla): add a prefill-from-embeddings StableHLO entry #858.
• The f32 additive 4D mask matches MLX and the independent checkpoint oracle cell-for-cell.
• External mask routing reaches every prefill layer and does not affect one-token decode.
• SigLIP/projector and final deterministic output meet the documented tolerances/token exactness.
• Padding, multi-image, and no-image behavior are tested through CLI and continuous batching.
• No duplicate full MLX decoder or silent causal-mask fallback remains.
• Existing text-only Gemma3 XLA behavior remains token-exact.