[inureyes]

Summary

• Add a dedicated Gemma3n XLA architecture, configuration and weight schema, token prefill, dense projected per-layer-input prefill, normal decode, and ragged batch decode.
• Port the Gemma3n text math across AltUp prediction and correction, magnitude normalization, LAUREL, alternating full and sliding attention, activation-sparse MLP, per-layer GeGLU interaction, plane collapse, final normalization, tied unembedding, and logit softcap.
• Make shared-KV cache ownership and public cache ordering explicit, align logical tokens, effective sequence length, dense PLE length, and cache length, and release request-owned PLE state across prefill, cancellation, completion, and slot reuse.
• Extend the IREE C ABI and Rust runtime for the Gemma3n module family, prepared dense-PLE descriptors, single-session execution, mixed batch execution, diagnostics, artifact compatibility, and quantized or supported unquantized weight loading.
• Embed the versioned kernel-v8 / abi-v6 CUDA Q4 projection and D=256 decode SDPA bundle, including dynamic decode-position push constants, exact scale-bit transport, bounded GQA/cache support, and materialized StableHLO fallbacks for unsupported dimensions, capacities, layouts, prefill, and truncated sliding windows.

Correctness boundaries

• Token IDs outside the per-layer vocabulary produce zero per-layer inputs, while token-derived dense PLE and externally prepared dense PLE use the same canonical shape and math.
• Normal and ragged decode share the same production attention core; shared full/sliding KV mappings are bounds-checked and validated across the layer transition.
• Native decode SDPA is restricted to one-query D=256 CUDA execution with valid GQA, cache capacity up to 1024, and an absent or capacity-spanning window; unsupported cases retain the existing materialized attention path.
• The native artifact identity includes kernel and ABI versions, minimum SM, PTX identity, and all Gemma3n graph compatibility inputs.

Validation

• cargo fmt --all -- --check
• cargo clippy --release -p mlxcel-xla --lib --features diagnostics -- -D warnings
• cargo clippy --release -p mlxcel --lib --features xla-iree,xla-diagnostics -- -D warnings
• cargo test --release -p mlxcel-xla --lib --features diagnostics: 185 passed, 0 failed, 10 ignored.
• Typed CUDA-capable IREE compilation verified dynamic tensor decode positions lower to scalar push constants.
• Actual CUDA Q4 group-64 projection matched its reference bit-exactly.
• Production CUDA SDPA matched the MLX vector oracle bit-exactly for live lengths 1, 2, 31, 32, 33, and 1024, a capacity-spanning window, exp2 underflow, and BF16 halfway rounding.
• Synthetic emitter coverage verifies kernel-v8 / abi-v6, normal and ragged native dispatch, materialized fallbacks, shared-KV mapping, AltUp, LAUREL, activation sparsity, GeGLU, plane collapse, and softcap.
• Real-IREE CUDA mixed lifecycle coverage verifies token versus dense-PLE generation, near-capacity prefill, cancellation, simultaneous token and PLE requests, and completed-slot reuse.
• Pinned E2B actual-input layer 0 and dense layer 10 probes matched MLX at every BF16 carrier; bounded layer traces for layers 10-19 and shared-KV layers 20-29 found no active or all-plane divergence.
• The pinned E2B canonical gate passed in one session: every prefill intermediate, final hidden state, diagnostic logits, token-prefill logits, and prepared-prefill logits matched exactly; prefix-decode logits remained within the pinned envelope with max absolute error 1.4765625, cosine 0.97841440, and normalized RMSE 0.26351152; all top-1 results matched at token 236798 and greedy output matched [236798, 236789].

Closes #876