feat(xla): support Gemma3 VLM additive prefill masks
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Parent and dependency
- Parent epic: #566
- Blocked by #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 #863 with Gemma3 checkpoint prefixes, shapes, normalization epsilon, image size, and projector mapping validated explicitly.
- Reproduce the net embedding scale at the canonical #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 #862/#863 and all text-only Gemma3 sliding/global attention regression gates.
## Acceptance criteria
- [ ] Exported merged embeddings have the exact post-scale magnitudes required by #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.
Contributor guide
Research direction
Start with src/loading/vlm_gemma.rs, src/vision/merge.rs::prepare_inputs_for_multimodal, and src/models/gemma3.rs, then review the XLA embedding boundary in #858 and the SigLIP/projector work in #863. Compare a pinned checkpoint against MLX and an independent oracle, including embeddings, every mask cell, layer outputs, logits, KV, and greedy tokens. Done means the listed padding, multi-image, fallback, batching, and text-only regression tests pass without causal-mask or duplicate-decoder behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, machine-learning, testing
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 25/100