feat(xla): support Molmo sparse additive visual embeddings
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Parent and dependency
- Parent epic: #566
- Blocked by #862
- Existing MLX implementation: `src/vision/molmo_vl.rs`, `src/vision/encoders/molmo.rs`, `src/vision/processors/molmo.rs`, and `src/models/molmo.rs`
## Problem
Molmo v1 does not replace placeholder embeddings like LLaVA. Its processor supplies `image_input_idx`; negative entries are skipped, and selected projected feature rows are added to the existing OLMo text embeddings at the remaining target positions. The vision path also includes multi-crop patch masks, selected ViT features, attention pooling, and a SwiGLU projector.
Scanning for an image token or replacing embedding rows would implement a different merge while still producing shape-valid inputs.
## Goal
Support a pinned Molmo v1 checkpoint end to end on XLA with IREE vision execution, processor-supplied sparse additive merge, OLMo language prefill/decode, CLI, and continuous batching.
## Required implementation
- Reuse/share `MolmoProcessor` crop, patch, `image_masks`, and `image_input_idx` generation. Preserve index order and negative-sentinel semantics in an owned prepared payload.
- Define static maxima/buckets for crop count, patches per crop, active feature rows, and expanded text length. Validate allocation arithmetic and reject overflow before tensor construction.
- Emit the Molmo CLIP-style ViT, configured intermediate-layer selection/combination, patch-mask behavior, attention pooling, and SwiGLU projector through StableHLO/IREE with explicit weight-layout checks.
- Keep mask semantics distinct: image patch coverage controls the vision encoder/pooler, while the language model uses its normal causal mask.
- Implement sparse additive merge from `(feature_row, target_position)` pairs. Require target positions to be in range and unique unless the reference explicitly defines repeated accumulation; require feature rows to be in range and preserve their mapping after sentinel removal.
- Cast projected features once to the post-scale text-embedding dtype, then add rather than replace. Document any OLMo embedding scaling at the #858 boundary.
- Include crop/patch config, selected vision layers, pool/projector config, merge mode, and OLMo text architecture in artifact compatibility identity.
- Integrate filtered loading, capabilities, CLI/server batching, cancellation, and no-media behavior without retaining a duplicate full MLX decoder.
## Non-goals
- Molmo2 token-scanned merge or pooled-patch index contract.
- Molmo Point coordinate output.
- General arbitrary sparse writes supplied by API clients.
## Validation
- Pin one Molmo checkpoint/revision and compare processor crops/patches/masks/indices, selected vision states, pooled/projected features, sparse active pairs, merged embeddings, prefill logits/KV, and greedy tokens against MLX and an independent oracle.
- Cover negative sentinels, non-contiguous target positions, one and multiple crops, masked patches, empty active features, invalid/duplicate indices, capacity overflow, no-image behavior, cancellation, and slot reuse.
- Add a negative fixture that performs replacement instead of addition and prove the merged-embedding/logit oracle catches it.
- Run CLI and mixed continuous-batch server tests on a production-relevant IREE target, plus #862 and text-only OLMo gates.
## Acceptance criteria
- [ ] Processor-supplied `image_input_idx` is authoritative; no token scan substitutes for it.
- [ ] Negative sentinels and feature-row/target-position correspondence are preserved exactly.
- [ ] Vision/pooling/projector intermediates and sparse additive merge match both oracles.
- [ ] Deterministic output is token-exact on the pinned fixture.
- [ ] Index/capacity errors fail before native execution and do not affect other slots.
- [ ] No full decoder duplication or silent MLX vision fallback remains.
- [ ] Text-only OLMo behavior remains token-exact.
Contributor guide
Research direction
Start with the dependency #862 and the existing MLX paths src/vision/molmo_vl.rs, src/vision/encoders/molmo.rs, src/vision/processors/molmo.rs, and src/models/molmo.rs. Trace the processor payload, vision outputs, merge contract, and OLMo integration before implementing the XLA/IREE path. Done means the pinned fixture matches both oracles token-for-token, including sparse additive merges, validation, batching, cancellation, and text-only behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 20/100