epic: OpenXLA backend architecture-coverage parity with the MLX engine

[inureyes]

Problem / Background

The MLX engine supports roughly 90 model_type values (see src/models/detection.rs). The OpenXLA/IREE backend (#449 M3 Stage 2d) currently serves only three text, dense architectures: Llama, Qwen2, and Gemma2, on both the single-sequence CLI path (MLXCEL_BACKEND=xla) and the continuous-batching serve path (mlxcel-server).

The gap is no longer infrastructure. The shared, architecture-agnostic base is already done and must not be re-created:

• Weight loading: single-file and sharded (index.json), bf16/f16/f32 widening, MLX 4bit/8bit affine dequant.
• IREE runtime, KV cache, and compile pipeline.
• Continuous-batching engine: XlaBatchEngine plus the ragged scheduler.
• Host-side sampling: temperature, top-k, top-p, min-p, seed.
• History penalties: repetition, frequency, presence, DRY.
• Streaming-safe stop strings.
• Server seam: BatchEngine trait, XlaServeWorker, /v1/completions, /metrics.

Goal

Bring the OpenXLA backend toward MLX-engine architecture coverage, planned in release windows. Because the common base is fixed, adding a new family needs only four per-family pieces:

1. Config::from_json arch detection and fields (emitter/config.rs).
2. weight_names per-layer order (iree.rs).
3. The emitter forward (emitter/model.rs).
4. Token-exact (or reference-exact) validation against an HF oracle on a real checkpoint.

The only exceptions are families that need a new shared graph primitive (MoE FFN, SSM/recurrent mixing); those are scoped as their own foundation issues so the per-family work stays cheap.

Validation pattern

The established gate for every architecture:

1. Dequant a 4bit/8bit checkpoint offline to f32 (or use a bf16 checkpoint directly).
2. Run HF fp32 as the oracle.
3. Run XLA on the quantized/bf16 model.
4. Compare: single-sequence path token-exact vs the HF fp32 oracle; serve path reference-exact via xla_batch_bench.

tanh-based architectures (soft-capping, GeGLU) may differ only at sub-0.01 logit near-ties; that is acceptable.

Sub-issues

Phase 1: Window A (Foundation)

• refactor: share the full per-layer attention core across all emitter graph kinds #494
• feat: implement Gemma2 sliding-window attention (local/global alternation) #495
• test: add a reusable per-architecture token-exact validation harness #496

Phase 2: Window B (Dense decoder-only families)

• feat: dense arch pack for qk-norm and Gemma family (Qwen3, Gemma1, Gemma3, SmolLM3, OLMo2/OLMo3) #497 (depends on refactor: share the full per-layer attention core across all emitter graph kinds #494)
• feat: dense arch pack for parallel-block and norm-variant families (Cohere/Cohere2, Phi3, StableLM, StarCoder2, Granite, MiniCPM/MiniCPM3) #498 (depends on refactor: share the full per-layer attention core across all emitter graph kinds #494)
• feat: dense arch pack for remaining Llama-family architectures (Mistral/Ministral, ExaOne, InternLM, Baichuan, GLM4, and more) #499 (depends on refactor: share the full per-layer attention core across all emitter graph kinds #494)

Phase 3: Window C (Mixture-of-Experts)

• feat: add the MoE FFN graph primitive (router + top-k expert dispatch) #500 (depends on refactor: share the full per-layer attention core across all emitter graph kinds #494)
• feat: MoE family coverage on the shared FFN primitive (Mixtral, Qwen2/Qwen3-MoE, OLMoE, DeepSeek-V2/V3, gpt-oss, and more) #501 (depends on feat: add the MoE FFN graph primitive (router + top-k expert dispatch) #500)

Phase 4: Window D (SSM / hybrid / recurrent)

• design: scope the SSM / hybrid / recurrent track (Mamba2, Jamba, RWKV7, Qwen3-Next, and more) #502

Phase 5: Window E (Multimodal / VLM)

• design: scope the multimodal / VLM track (Qwen2-VL/2.5-VL/3-VL, Gemma3n, Phi4MM, Molmo, and more) #503

Acceptance Criteria

• All sub-issues completed and merged.
• Each new family loads and generates via the CLI MLXCEL_BACKEND=xla path and serves via mlxcel-server.
• Each new family passes the validation gate: token-exact single-seq vs HF fp32 oracle, reference-exact serve via xla_batch_bench.
• Windows D and E either land a first reference architecture or record an explicit deferral into a dedicated follow-up epic.

Technical Considerations

• Parent milestone: feat: OpenXLA reference backend - export-route spike through 4-bit quantized decode #449 (OpenXLA reference backend, M3 Stage 2d).
• The shared attention-core refactor (refactor: share the full per-layer attention core across all emitter graph kinds #494, Window A) and the MoE FFN primitive (feat: add the MoE FFN graph primitive (router + top-k expert dispatch) #500, Window C) are prerequisites that unblock the cheap per-family work; explicit dependency edges are recorded in the sub-issue list above.
• Windows D (SSM/hybrid/recurrent) and E (multimodal/VLM) are separate tracks that introduce non-attention graph primitives or vision/audio encoders; each is expected to spin out into its own follow-up epic after its design issue lands.