feat(xla): add the Gemma3n text backbone and dense PLE prefill

[inureyes]

Parent and dependency

• 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
• Existing MLX text/PLE implementation: src/models/gemma3n.rs
• Existing XLA gap: src/lib/mlxcel-xla has no Gemma3n/AltUp/Laurel/per-layer-input emitter path.

Problem

Gemma3n audio cannot reuse the existing Gemma3 XLA decoder. The text backbone expands hidden state into multiple AltUp planes, normalizes plane magnitudes, applies AltUp prediction/correction, LAUREL residuals, per-layer input gates with GeGLU interaction, activation sparsity, alternating full/sliding attention, optional shared KV layers, and a final collapse/unembed/mean. Each token also has a per-layer embedding tensor projected together with the base embedding.

Therefore a standalone audio encoder issue would still have no compatible XLA language graph to receive its embeddings/PLE state.

Goal

Add token-exact Gemma3n text prefill/decode and an embeddings-plus-dense-per-layer-input prefill entry to the XLA backend before audio integration.

Required implementation

• Add an explicit Gemma3n architecture/config parser and weight map. Validate layer_types, shared-KV mapping, AltUp plane counts/active indices, Laurel rank, per-layer vocab/hidden sizes, activation sparsity, sliding window, softcaps, quantization metadata, and tied embedding/LM-head behavior.
• Emit scaled token embeddings and embed_tokens_per_layer; token IDs outside vocab_size_per_layer_input must map to zero exactly as get_per_layer_inputs does.
• Emit per_layer_model_projection, projection norm, 1/sqrt(2) combination, and canonical dense PLE shape [Lp, num_hidden_layers, hidden_size_per_layer_input].
• Port AltUp input projections, magnitude normalization, predict/correct equations, active-plane scaling, LAUREL, attention, activation-sparse MLP, per-layer gate plus compiled approximate GeGLU interaction, per-layer projection/norm, unembed projections, plane collapse/mean, final norm, tied head, and final logit softcap.
• Preserve full/sliding layer masks and shared-KV cache mapping. Define the public KV output order independently of logical layer count and prove decode reads the same shared cache indices.
• Add a distinct prefill_embeddings_ple.main entry accepting post-scale merged embeddings plus dense projected PLE. Token prefill may compute PLE internally; multimodal prefill receives it explicitly. Ordinary feat(xla): add a prefill-from-embeddings StableHLO entry #858 and feat(xla): add per-layer multimodal feature injection to prefill #867 schemas remain unchanged.
• Extend the IREE bundle/C ABI/single/batch runtime for the Gemma3n module set and dense PLE descriptor. Validate shape/dtype/byte counts and release PLE after prefill/cancellation.
• Keep effective sequence length, PLE length, cache length, and logical tokens aligned; padding/slicing must be explicit and never silently truncate a modality.
• Include every Gemma3n structural option in artifact compatibility identity.

Non-goals

• Gemma3n vision or audio encoder emission.
• Reusing sparse Qwen DeepStack injection as if it were dense Gemma3n PLE.
• Performance-specific MLX M5 dispatch choices that do not change math.

Validation

• Build small synthetic fixtures isolating AltUp predict/correct, magnitude normalization, LAUREL, activation sparsity, shared KV mapping, per-layer GeGLU interaction, plane collapse, and softcap.
• Pin a Gemma3n text checkpoint/revision and compare scaled embeddings, projected PLE, every layer's active/all AltUp planes at selected layers, K/V, final hidden state, logits, and greedy tokens against MLX and the maintained reference.
• Compare token prefill with embeddings-plus-PLE prefill using token-derived inputs.
• Cover full/sliding alternation, shared/nonshared KV layers, text tokens outside per-layer vocab, quantized and supported unquantized weights, near-capacity prefill, mixed batch slots, cancellation, and slot reuse.
• Execute prefill/decode in real IREE on a production-relevant target and run all existing XLA text gates.

Acceptance criteria

• Gemma3n token prefill/decode is token-exact on a pinned checkpoint.
• Every AltUp/Laurel/PLE/inter-layer invariant has an intermediate oracle assertion.
• Dense PLE prefill has a distinct schema and matches token prefill with token-derived inputs.
• Shared KV layer mapping is explicit, bounds-checked, and validated through decode.
• PLE tensors are request-owned, released after prefill/cancellation, and cannot bleed across slots.
• Artifact identity covers all math/shape-affecting Gemma3n configuration.
• Existing non-Gemma3n XLA architectures remain token-exact.