feat(xla): parameterize context and multimodal prefill capacity

[inureyes]

Parent and design context

• Parent epic: epic: OpenXLA multimodal / VLM architecture support (Qwen2-VL, Gemma3n, Phi4MM, Molmo, and more) #566
• Original design spike: design: scope the multimodal / VLM track (Qwen2-VL/2.5-VL/3-VL, Gemma3n, Phi4MM, Molmo, and more) #503
• Current implementation baseline: main at or after ac38f97b0541a49011c35078973b1739e5a069f4

Problem

The XLA path encodes a 256-token limit directly into graph and runtime shapes. src/lib/mlxcel-xla/src/emitter/model.rs defines MAX_SEQ = 256 and derives PREFILL_LP from it; src/lib/mlxcel-xla/src/iree.rs repeats the same capacity for prefill and ragged batches. KV cache shapes, token/position tensors, masks, slot copies, and admission checks all inherit that constant.

This blocks the parent epic before any vision graph is added. A standard LLaVA image already contributes 576 patch tokens at 336/14 or 729 patch tokens at 384/14, before text tokens and generation headroom. The current admission path also needs to reject based on the expanded prompt length plus the requested generation budget, rather than prompt length alone.

Goal

Make the compiled XLA context capacity an explicit, consistently propagated configuration value, and enforce it before native execution for text and future multimodal requests.

Required implementation

• Add one canonical context-capacity field to the XLA graph/runtime configuration. Use it to derive prefill length, per-layer KV cache shapes, continuous-batch slot shapes, ragged batch tensors, position ranges, and mask dimensions.
• Include the capacity in every compiled-artifact/cache identity so a module compiled for one shape cannot be loaded as another shape.
• Remove duplicated production constants from emitter/model.rs, iree.rs, and the C shim. A default of 256 may remain only as a documented compatibility default or test fixture.
• Thread the selected capacity through single-session compilation, XlaBatchEngine construction, and server startup. Fail during configuration/loading if paired modules or runtime buffers disagree.
• Centralize request admission around the effective sequence length. For token requests this is the token count; for multimodal requests it is the post-placeholder-expansion count supplied by the preprocessor.
• Reject requests when effective_prompt_len + max_new_tokens > context_capacity. Return a typed, actionable error containing all three values, and reject before acquiring or mutating a batch slot.
• Ensure decode position and cache-length arithmetic cannot cross the configured bound even if a caller bypasses high-level admission.
• Document how operators choose a capacity and the memory/compile-time consequences. Do not imply StableHLO dynamic shapes are supported.

Non-goals

• Adding an embeddings prefill entry point.
• Implementing image/audio preprocessing or a VLM architecture.
• Replacing static StableHLO shapes with fully dynamic shapes.

Validation

• Add pure Rust emitter/schema tests for at least capacities 256 and 1024. Assert every public graph input/output and KV cache dimension is internally consistent.
• Add admission boundary tests for exact fit, one-token overflow, expanded multimodal length, and max_new_tokens overflow.
• Add a regression test proving an admission failure does not consume a continuous-batch slot or alter another request.
• Compile and run the text prefill/decode pair with IREE at two capacities on an available target. Record target and commands in the PR.
• Run cargo test -p mlxcel-xla --lib and the relevant XLA validation script(s).
• Demonstrate that the compatibility default preserves the existing text-only token-exact result.

Acceptance criteria

• No production graph/runtime buffer shape depends on an unthreaded magic value of 256.
• Artifact identity and module loading reject capacity mismatches deterministically.
• Text and expanded multimodal lengths use the same admission invariant.
• Prompt-plus-generation overflow fails before native invocation and before slot mutation.
• Structural tests cover multiple capacities and a real IREE prefill/decode run passes.
• Operator-facing configuration and failure messages are documented.