epic: OpenXLA multimodal / VLM architecture support (Qwen2-VL, Gemma3n, Phi4MM, Molmo, and more)
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem / current state
Parent architecture epic: #493. Scoping design: #503 and `spike/openxla/MULTIMODAL_VLM_DESIGN.md`.
The current OpenXLA/IREE path is still text-only at the request boundary, but the implementation has evolved beyond the original 2025 design assumptions:
- `src/lib/mlxcel-xla/src/emitter/model.rs` has a mature shared text layer stack, but token prefill still gathers `embed[tokens]`; there is no embeddings-prefill entry.
- `src/lib/mlxcel-xla/src/emitter/model.rs` and `src/lib/mlxcel-xla/src/iree.rs` hard-code a 256-token context/prefill shape. A normal LLaVA image alone contributes 576 or 729 patch tokens at common resolutions, before text and generation headroom.
- `src/lib/mlxcel-xla/src/iree.rs` and `src/lib/mlxcel-xla/csrc/xla_iree.c` load only token-prefill/decode modules, and `XlaBatchEngine::Pending` owns only token IDs.
- `src/backend/session.rs` has an XLA path that can ignore supplied input embeddings and run token-only generation. Multimodal support must fail closed until the real path exists.
- `src/server/batch/xla_worker.rs` rejects images, audio, and video. `spawn_xla_model_worker` owns no VLM/audio preprocessor.
- The mature MLX stack returns `MlxArray`-backed `InputEmbeddings`; its VLM loaders commonly construct a full text decoder. XLA needs an owned backend-neutral payload and filtered processor/encoder/embedding loading, not a duplicate decoder.
- Qwen M-RoPE, Qwen3 DeepStack, Gemma3 additive masks, Molmo sparse/indexed pooling, and Youtu windowed vision have different contracts and cannot be represented by one generic “merge image tokens” task.
- Phi4MM and Gemma3n audio are not currently executable reference paths in mlxcel. Phi4MM explicitly rejects audio placeholders and skips audio/speech-LoRA weights; Gemma3n loads text/vision but no audio tower. Their MLX reference implementations must be qualified before XLA ports.
## Goal
Accept supported multimodal requests on the OpenXLA/IREE backend and produce reference-correct output through CLI and continuous-batch OpenAI-compatible serving, while preserving text-only behavior and using explicit, bounds-checked contracts for:
1. compiled context capacity and admission;
2. prefill from owned embeddings, positions, masks, and optional family side inputs;
3. host preprocessing without a duplicate full decoder;
4. IREE encoder/projector execution where required;
5. per-request/slot multimodal state, cancellation, metrics, and capability truth;
6. independent intermediate and token-exact architecture validation.
## Implementation DAG
GitHub native sub-issue and `blocked-by` relationships are authoritative. The checklist below is a readable index.
### A. Independent foundation / oracle work
- [x] #857 — parameterize XLA context and multimodal prefill capacity; enforce expanded prompt plus generation bounds.
- [x] #858 — emit a distinct `prefill_embeddings.main` with explicit additive attention bias and token-prefill parity.
- [x] #859 — extract an owned backend-neutral host VLM preprocessor and filtered LLaVA reference loader.
- [x] #874 — implement and qualify the missing MLX Phi4MM audio reference, including speech LoRA.
- [x] #875 — implement and qualify the missing MLX Gemma3n audio reference.
- [ ] #932 — define versioned operator-level MLX/IREE numeric contracts and deterministic micro-oracles.
### B. Core runtime and first end-to-end image path
- [x] #860 — load token-prefill/embeddings-prefill/decode as one compatible IREE bundle and seed single/batched slots from embeddings. Blocked by #857, #858, #859.
- [x] #861 — wire image preprocessing/admission into CLI and continuous-batch serving with truthful capability reporting. Blocked by #860.
- [x] #862 — qualify a pinned LLaVA-style checkpoint through intermediate oracles, CLI, and streaming server. Blocked by #861.
- [x] #863 — move the reference SigLIP/ViT encoder and projector to IREE, retaining host-first only as an explicit oracle/fallback. Blocked by #859 and #862.
### C. Family-side language contracts
- [x] #864 — add explicit 3-axis M-RoPE positions and per-slot decode deltas. Blocked by #860.
- [x] #867 — add sparse per-layer multimodal feature injection for Qwen3 DeepStack. Blocked by #858 and #860.
- [x] #876 — add the Gemma3n AltUp/Laurel text backbone and dense PLE embeddings-prefill contract. Blocked by #857.
### D. Vision-family breadth
- [x] #865 — support Qwen2-VL packed vision plus M-RoPE end to end. Blocked by #862 and #864.
- [ ] #866 — support Qwen2.5-VL windowed/full vision attention and restoration ordering. Blocked by #865 and #932.
- [ ] #868 — support Qwen3-VL main plus DeepStack vision outputs, dense and qualified MoE text variants. Blocked by #865, #867, and #932.
- [ ] #869 — support Gemma3 VLM post-scale embeddings and additive f32 4D prefill masks. Blocked by #863 and #932.
- [ ] #870 — support Molmo processor-supplied sparse additive visual embeddings. Blocked by #862 and #932.
- [ ] #871 — support Molmo2 indexed attention pooling and additive `image_patch_id` merge. Blocked by #862 and #932.
- [ ] #872 — support Youtu-VL flattened patches, windowed vision/merger, and dense MLA text inference. Blocked by #862 and #932.
### E. Audio serving and XLA execution
- [x] #873 — add bounded, cancellable, backend-neutral audio preprocessing and serving plumbing. Blocked by #861.
- [x] #877 — execute Phi4MM Conformer/projection and per-slot speech/vision LoRA modes through IREE. Blocked by #873 and #874.
- [ ] #878 — execute Gemma3n audio encoder/projection through IREE and feed exact dense PLE prefill. Blocked by #873, #875, #876, and #932.
## Architectural invariants
- Text prefill, embeddings prefill, family-specific prefill, and decode are separate versioned module contracts; no path accepts a tensor and then ignores it.
- Compiled context capacity, dtype, architecture config, position/mask mode, KV layout, encoder shape buckets, and side-input schema participate in artifact compatibility.
- Logical token IDs remain available for penalties, detokenization, stop matching, and public usage accounting; expanded multimodal length owns cache admission/positioning.
- Host preprocessors return owned contiguous tensor descriptors and may retain only required processor/encoder/projector/embedding weights, never a duplicate full decoder.
- Media acquisition/decode/resampling follows existing security and size limits in a bounded cancellable stage outside the scheduler loop.
- All request-specific M-RoPE deltas, DeepStack/PLE tensors, audio modes/adapters, and large host buffers are slot-scoped and cleared on prefill completion, cancellation, error, and slot reuse.
- Host-first execution is an explicit qualified fallback/oracle. An IREE-capable family may not silently execute its encoder on MLX.
- A family is advertised in `SessionCapabilities` only when model weights/config, processor, runtime modules, and execution dispatch all support it through the same predicate.
## Validation and merge gates
- Every emitter task needs structural/golden IR tests and a real IREE execution; emitter-string success alone is insufficient.
- Every family pins checkpoint/source revisions, deterministic media/prompts, dtype-specific intermediate tolerances, and token-exact greedy output.
- Compare processor outputs, encoder/projector intermediates, merged embeddings/positions/masks/side inputs, prefill logits, selected all-layer KV values, and final tokens. Report the first divergent stage.
- Exercise non-streaming CLI and streaming server paths, including mixed text/multimodal batching, cancellation, slot reuse, capacity overflow, malformed media, and per-request failure isolation.
- Run the existing text-only XLA architecture oracle and continuous-batch suites for every merge.
- Record target/device, commands, compile time, host/device memory, transfers, prefill latency, and decode throughput for real IREE qualification.
## Epic acceptance criteria
- [ ] Capacity and module compatibility prevent shape/config drift before native execution.
- [ ] Token-derived embeddings match token prefill logits and all-layer KV through the public runtime/C shim.
- [ ] Supported media reaches real embeddings/side-input prefill in CLI and server; unsupported media fails explicitly and is never dropped.
- [ ] The pinned LLaVA reference is token-exact through a production-relevant IREE target and streaming server.
- [ ] Each advertised Qwen/Gemma/Molmo/Youtu family passes its issue's intermediate and end-to-end oracle gates.
- [ ] Phi4MM and Gemma3n audio have qualified MLX references before their IREE ports and pass deterministic CLI/server ASR fixtures.
- [ ] Text-only requests remain token-exact and mixed continuous batching has no state leakage.
- [ ] All native sub-issues are completed, or any deliberate de-scope is recorded in this epic with replacement tracking and acceptance-impact rationale.
## References
- Parent architecture epic: #493
- Scoping design: #503
- Design document: `spike/openxla/MULTIMODAL_VLM_DESIGN.md`
- Shared attention and mask foundations: #494 and #495
Contributor guide
Assessment
This issue has not been assessed yet.