feat(xla): seed single and batched IREE sessions from embeddings #860

Parent and dependencies

• Parent epic: epic: OpenXLA multimodal / VLM architecture support (Qwen2-VL, Gemma3n, Phi4MM, Molmo, and more) #566
• Blocked by feat(xla): parameterize context and multimodal prefill capacity #857, feat(xla): add a prefill-from-embeddings StableHLO entry #858, and refactor(vlm): extract a host multimodal preprocessor for XLA #859
• Runtime implementation: src/lib/mlxcel-xla/src/iree.rs, src/lib/mlxcel-xla/src/batch.rs, and src/lib/mlxcel-xla/csrc/xla_iree.c
• Session boundary: src/lib/mlxcel-core/src/session.rs and src/backend/session.rs

Problem

The IREE runtime currently compiles and loads only a token-prefill/decode pair. xla_llama_create accepts two modules, xla_llama_prefill_slot_logits uploads token IDs/positions/real_len, and XlaBatchEngine::Pending owns only prompt: Vec<i32>. There is no ABI or ownership path for merged embeddings and an explicit attention bias.

Separately, Session::Xla currently ignores supplied input embeddings and falls back to token IDs. That behavior can produce plausible but semantically wrong output and must not survive the new contract.

Goal

Load the token-prefill, embeddings-prefill, and decode modules as one compatible runtime bundle, and seed both single and continuous-batch XLA sessions from an owned prepared-embeddings payload.

Required implementation

• Replace the pair-only compile/load abstraction with a bundle that contains prefill.main, prefill_embeddings.main, and decode.main, plus a single compatibility fingerprint covering architecture, weights, dtype, context capacity, hidden size, KV layout, and argument schema.
• Extend xla_llama_create and the Rust wrapper to load all three modules together. Fail atomically and release already-created VM contexts/modules if any member is missing or incompatible.
• Add a C ABI entry point for embeddings prefill that accepts validated embeddings, positions, additive attention bias, real_len, and slot selection. Use explicit byte lengths/dtypes/shapes or a typed descriptor; do not infer buffer sizes from unchecked pointers.
• Reuse the existing device-side KV slot population path. Embeddings prefill must write the same rank-5 batch cache slot and cache length as token prefill without downloading/re-uploading KV tensors through the host.
• Extend the batch pending-input representation to an enum for token prefill versus prepared embeddings. The engine must retain logical token IDs separately from expanded embeddings so repetition penalties, token history, detokenization, cancellation, and usage accounting remain correct.
• Set cache length and first decode position from the prepared payload's effective expanded length, not from placeholder token count.
• Validate all shape, dtype, byte-count, real_len, context-capacity, and slot invariants in Rust before entering C. Repeat safety-critical bounds checks in C before constructing IREE buffer views.
• Preserve request ownership through queueing and cancellation. Dropping or cancelling a pending multimodal request must release its large host buffers exactly once and must not leak or poison a slot.
• Route the single-session embeddings API through the new entry. Until a call has a compatible prepared payload, return an explicit unsupported/invalid-input error; never ignore embeddings or silently execute token-only prefill.
• Keep token-only prefill and decode ABI behavior compatible unless a versioned symbol change is unavoidable.

Non-goals

• Image decoding, placeholder expansion, or vision encoding.
• Server admission of multimodal requests.
• M-RoPE/3D position semantics.
• Moving the vision encoder onto IREE.

Validation

• Unit-test descriptor validation with wrong dtype, rank, hidden size, byte count, mask shape, real_len, capacity, and slot index.
• Run token-prefill and embeddings-prefill parity through the public Rust API and through the C shim, comparing logits and every populated KV layer.
• Add batch tests that interleave text and embeddings requests across slots, cancel each form before and after prefill, and verify slot reuse and independent token histories.
• Add an overflow/canary test around host buffers and exercise native error cleanup under sanitizers where available.
• Run a real IREE continuous-batch decode after embeddings prefill on an available target.
• Prove the old token-only single and batched validation results are unchanged.

Acceptance criteria

• One runtime bundle atomically loads all three compatible modules.
• The embeddings C ABI has explicit, bounds-checked tensor descriptors.
• Embeddings prefill populates the selected device KV slot without a KV host round trip.
• Logical token history and expanded cache length remain distinct and correct.
• Mixed text/multimodal queueing, cancellation, and slot reuse are covered by tests.
• No XLA session path silently ignores embeddings.
• Existing text-only single and batch behavior remains token-exact.