feat(xla): add per-layer multimodal feature injection to prefill

[inureyes]

Parent and dependencies

• Parent epic: epic: OpenXLA multimodal / VLM architecture support (Qwen2-VL, Gemma3n, Phi4MM, Molmo, and more) #566
• Blocked by feat(xla): add a prefill-from-embeddings StableHLO entry #858 and feat(xla): seed single and batched IREE sessions from embeddings #860
• Existing Qwen3-VL state/injection behavior: src/vision/qwen3_vl.rs and src/vision/encoders/qwen3_vl.rs
• Shared XLA layer stack: src/lib/mlxcel-xla/src/emitter/model.rs

Problem

prefill_embeddings.main seeds only the initial hidden-state sequence. Qwen3-VL additionally projects selected intermediate vision-block outputs and injects them at visual token positions in configured language-model layers (DeepStack). Passing only the final merged visual embeddings drops these side branches and produces an incomplete model even if the initial prefill logits look plausible.

This capability belongs in the language-prefill contract rather than in a Qwen3-specific server shortcut.

Goal

Add a versioned XLA embeddings-prefill variant that can apply validated, sparse per-layer multimodal feature additions at selected sequence positions without changing ordinary embeddings prefill or decode.

Canonical injection contract

• Represent configured target language-layer indices explicitly and include them in the compiled artifact identity.
• Use a compact owned payload equivalent to visual_positions[Nv], layer_features[K, Nv, hidden_size], target layer_indices[K], and actual counts within documented static maxima. Do not allocate a dense [num_layers, Lp, hidden_size] host tensor.
• Visual positions refer to the expanded prefill sequence and must be sorted, unique, in range, and consistent with the placeholder expansion metadata.
• Features are post-projector, post-scaling residual additions in the model hidden-state dtype. Define exactly whether injection occurs before or after the target layer and match the existing MLX Qwen3 implementation.
• Injection is prefill-only. Incremental decode owns no lingering deepstack buffers after the KV cache is seeded.

Required implementation

• Emit a distinct, stable entry such as prefill_embeddings_deepstack.main; do not add dummy deepstack arguments to the ordinary text/LLaVA module.
• Refactor the shared transformer-layer loop to expose one reviewed residual-add hook while keeping all non-hooked layers byte/semantically identical.
• Implement bounded gather/scatter or mask-add StableHLO logic for compact visual positions. Reject duplicate/padded/out-of-range positions rather than relying on undefined scatter collision behavior.
• Extend compile/load compatibility, Rust tensor descriptors, C ABI, single-session, and batch pending/slot state for the optional module and payload. Load it only for architectures that declare the feature.
• Validate K, Nv, layer indices, feature shape/dtype/byte count, expanded length, and hidden size before native invocation and in the C safety boundary.
• Release large side tensors immediately after successful prefill and on every cancellation/error path.
• Keep cache layout and decode module unchanged.

Non-goals

• Emitting the Qwen3 vision encoder/mergers.
• Defining which Qwen3 checkpoint uses which target layers beyond reading validated config.
• General arbitrary callbacks inside the StableHLO layer loop.

Validation

• Add a small synthetic transformer fixture with known features injected at first/middle/last layers; compare every hooked layer output, final logits, and KV against an eager reference.
• Port a fixed Qwen3-VL deepstack fixture from the MLX implementation and compare selected residual states.
• Cover zero visual positions, multiple target layers, invalid/duplicate positions, invalid layer indices, static-max overflow, cancellation, and slot reuse.
• Prove the ordinary prefill_embeddings.main, token prefill, and decode IR/results are unchanged.
• Execute the deepstack module in a real IREE runtime, not only emitter tests.

Acceptance criteria

• Deepstack prefill has a distinct schema/module and does not burden ordinary prefill calls.
• Compact features are injected at the exact configured layers and expanded-sequence positions.
• All descriptor, position, and layer bounds are checked before native execution and again at the C boundary.
• Side tensors are released after prefill/cancellation and never become decode slot state.
• Synthetic and Qwen3 fixtures match eager intermediate states, logits, and KV.
• Existing token and ordinary embeddings prefill remain token-exact.