epic: OpenXLA multimodal / VLM architecture support (Qwen2-VL, Gemma3n, Phi4MM, Molmo, and more)

[inureyes]

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

• feat(xla): parameterize context and multimodal prefill capacity #857 — parameterize XLA context and multimodal prefill capacity; enforce expanded prompt plus generation bounds.
• feat(xla): add a prefill-from-embeddings StableHLO entry #858 — emit a distinct prefill_embeddings.main with explicit additive attention bias and token-prefill parity.
• refactor(vlm): extract a host multimodal preprocessor for XLA #859 — extract an owned backend-neutral host VLM preprocessor and filtered LLaVA reference loader.
• feat(vlm): add a token-exact Phi4MM audio reference path #874 — implement and qualify the missing MLX Phi4MM audio reference, including speech LoRA.
• feat(vlm): add a token-exact Gemma3n audio reference path #875 — implement and qualify the missing MLX Gemma3n audio reference.
• feat(xla): define operator-level MLX/IREE numeric contracts #932 — define versioned operator-level MLX/IREE numeric contracts and deterministic micro-oracles.

B. Core runtime and first end-to-end image path

• feat(xla): seed single and batched IREE sessions from embeddings #860 — load token-prefill/embeddings-prefill/decode as one compatible IREE bundle and seed single/batched slots from embeddings. Blocked by feat(xla): parameterize context and multimodal prefill capacity #857, feat(xla): add a prefill-from-embeddings StableHLO entry #858, refactor(vlm): extract a host multimodal preprocessor for XLA #859.
• feat(xla): admit image requests in CLI and continuous-batch serving #861 — wire image preprocessing/admission into CLI and continuous-batch serving with truthful capability reporting. Blocked by feat(xla): seed single and batched IREE sessions from embeddings #860.
• test(xla): validate a LLaVA reference architecture end to end #862 — qualify a pinned LLaVA-style checkpoint through intermediate oracles, CLI, and streaming server. Blocked by feat(xla): admit image requests in CLI and continuous-batch serving #861.
• feat(xla): emit SigLIP/ViT encoders and projectors through IREE #863 — move the reference SigLIP/ViT encoder and projector to IREE, retaining host-first only as an explicit oracle/fallback. Blocked by refactor(vlm): extract a host multimodal preprocessor for XLA #859 and test(xla): validate a LLaVA reference architecture end to end #862.

C. Family-side language contracts

• feat(xla): support multimodal 3D RoPE positions and decode deltas #864 — add explicit 3-axis M-RoPE positions and per-slot decode deltas. Blocked by feat(xla): seed single and batched IREE sessions from embeddings #860.
• feat(xla): add per-layer multimodal feature injection to prefill #867 — add sparse per-layer multimodal feature injection for Qwen3 DeepStack. Blocked by feat(xla): add a prefill-from-embeddings StableHLO entry #858 and feat(xla): seed single and batched IREE sessions from embeddings #860.
• feat(xla): add the Gemma3n text backbone and dense PLE prefill #876 — add the Gemma3n AltUp/Laurel text backbone and dense PLE embeddings-prefill contract. Blocked by feat(xla): parameterize context and multimodal prefill capacity #857.

D. Vision-family breadth

• feat(xla): support Qwen2-VL as the representative M-RoPE VLM #865 — support Qwen2-VL packed vision plus M-RoPE end to end. Blocked by test(xla): validate a LLaVA reference architecture end to end #862 and feat(xla): support multimodal 3D RoPE positions and decode deltas #864.
• feat(xla): support Qwen2.5-VL windowed vision attention #866 — support Qwen2.5-VL windowed/full vision attention and restoration ordering. Blocked by feat(xla): support Qwen2-VL as the representative M-RoPE VLM #865 and feat(xla): define operator-level MLX/IREE numeric contracts #932.
• feat(xla): support Qwen3-VL and DeepStack vision features #868 — support Qwen3-VL main plus DeepStack vision outputs, dense and qualified MoE text variants. Blocked by feat(xla): support Qwen2-VL as the representative M-RoPE VLM #865, feat(xla): add per-layer multimodal feature injection to prefill #867, and feat(xla): define operator-level MLX/IREE numeric contracts #932.
• feat(xla): support Gemma3 VLM additive prefill masks #869 — support Gemma3 VLM post-scale embeddings and additive f32 4D prefill masks. Blocked by feat(xla): emit SigLIP/ViT encoders and projectors through IREE #863 and feat(xla): define operator-level MLX/IREE numeric contracts #932.
• feat(xla): support Molmo sparse additive visual embeddings #870 — support Molmo processor-supplied sparse additive visual embeddings. Blocked by test(xla): validate a LLaVA reference architecture end to end #862 and feat(xla): define operator-level MLX/IREE numeric contracts #932.
• feat(xla): support Molmo2 indexed attention pooling #871 — support Molmo2 indexed attention pooling and additive image_patch_id merge. Blocked by test(xla): validate a LLaVA reference architecture end to end #862 and feat(xla): define operator-level MLX/IREE numeric contracts #932.
• feat(xla): support Youtu-VL windowed vision and MLA inference #872 — support Youtu-VL flattened patches, windowed vision/merger, and dense MLA text inference. Blocked by test(xla): validate a LLaVA reference architecture end to end #862 and feat(xla): define operator-level MLX/IREE numeric contracts #932.

E. Audio serving and XLA execution

• feat(xla): add bounded audio preprocessing and serving plumbing #873 — add bounded, cancellable, backend-neutral audio preprocessing and serving plumbing. Blocked by feat(xla): admit image requests in CLI and continuous-batch serving #861.
• feat(xla): support Phi4MM audio and per-slot speech adapters #877 — execute Phi4MM Conformer/projection and per-slot speech/vision LoRA modes through IREE. Blocked by feat(xla): add bounded audio preprocessing and serving plumbing #873 and feat(vlm): add a token-exact Phi4MM audio reference path #874.
• feat(xla): support Gemma3n audio with dense PLE prefill #878 — execute Gemma3n audio encoder/projection through IREE and feed exact dense PLE prefill. Blocked by feat(xla): add bounded audio preprocessing and serving plumbing #873, feat(vlm): add a token-exact Gemma3n audio reference path #875, feat(xla): add the Gemma3n text backbone and dense PLE prefill #876, and feat(xla): define operator-level MLX/IREE numeric contracts #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: epic: OpenXLA backend architecture-coverage parity with the MLX engine #493
• Scoping design: design: scope the multimodal / VLM track (Qwen2-VL/2.5-VL/3-VL, Gemma3n, Phi4MM, Molmo, and more) #503
• Design document: spike/openxla/MULTIMODAL_VLM_DESIGN.md
• Shared attention and mask foundations: refactor: share the full per-layer attention core across all emitter graph kinds #494 and feat: implement Gemma2 sliding-window attention (local/global alternation) #495