feat(xla): support Gemma3n audio with dense PLE prefill

[inureyes]

Parent and dependencies

• Parent epic: epic: OpenXLA multimodal / VLM architecture support (Qwen2-VL, Gemma3n, Phi4MM, Molmo, and more) #566
• Blocked by feat(xla): add bounded audio preprocessing and serving plumbing #873, feat(vlm): add a token-exact Gemma3n audio reference path #875, and feat(xla): add the Gemma3n text backbone and dense PLE prefill #876
• Qualified MLX/oracle contract: feat(vlm): add a token-exact Gemma3n audio reference path #875
• Required XLA language/PLE contract: feat(xla): add the Gemma3n text backbone and dense PLE prefill #876

Problem

Gemma3n audio combines a specialized mel/Conformer encoder with the already-specialized Gemma3n AltUp/Laurel/dense-PLE language path. The audio frame mask drives convolutional subsampling and chunked local attention; the resulting projected tokens and per-layer inputs must align with expanded prompt positions before #876 seeds KV caches.

A host-only encoder fallback would leave a substantial second compute graph/runtime and would not complete the epic's OpenXLA parity target.

Goal

Execute the qualified Gemma3n audio encoder and multimodal projection through StableHLO/IREE, feed exact merged embeddings plus dense PLE to #876, and serve bounded audio requests through CLI and continuous batching.

Required implementation

• Keep waveform decode/resample and mel extraction in the bounded feat(xla): add bounded audio preprocessing and serving plumbing #873 host stage unless a measured, independently validated IREE front end is proposed. The audio.main graph accepts canonical mel features and mask/actual-frame metadata from feat(vlm): add a token-exact Gemma3n audio reference path #875.
• Define static frame buckets up to the pinned checkpoint's duration limit. Include mel bins, frame policy, convolution shapes/strides/padding, hidden size, Conformer depth, chunk/context sizes, projection/pooling config, and dtype in artifact identity.
• Emit both subsample convolution blocks with manual padding, cumulative group normalization, ReLU, mask/length propagation, and output projection exactly.
• Emit all Conformer components from feat(vlm): add a token-exact Gemma3n audio reference path #875: feed-forward halves, chunked local attention, relative position projection/shift, per-dimension query scaling, logit softcap, light causal/depthwise convolution, residuals, clipping behavior relevant to inference, and final normalization/pooling/projection.
• Carry masks through every stage; padded mel frames and padded output tokens must not contribute to cumulative statistics, attention, pooling, placeholder count, or PLE generation.
• Produce projected audio tokens plus explicit per-clip output lengths. Apply the exact feat(vlm): add a token-exact Gemma3n audio reference path #875 placeholder expansion/replacement and construct the post-scale merged embeddings and dense per-layer-input tensor required by prefill_embeddings_ple.main.
• Preserve multiple-clip order and request-scoped PLE. Submit effective expanded length to feat(xla): parameterize context and multimodal prefill capacity #857 admission, then release mel/encoder/projected/PLE buffers after successful prefill or on cancellation/error.
• Load audio weights conditionally. Capability is true only when the checkpoint contains compatible audio config/weights, audio.main, and the feat(xla): add the Gemma3n text backbone and dense PLE prefill #876 language bundle. Text/image-only checkpoints must not pay the audio load/memory cost.
• Integrate CLI/server audio, mixed image/audio if qualified by feat(vlm): add a token-exact Gemma3n audio reference path #875, cancellation, metrics, and slot reset. No audio/PLE state may survive slot reuse.
• Do not retain or invoke a full MLX Gemma3n decoder/audio tower on the production XLA path.

Non-goals

• Gemma4 audio or Phi4MM adapters.
• Speech output/streaming audio input.
• Claiming arbitrary Gemma3n checkpoint compatibility without the pinned processor/config oracle.

Validation

• Compare every stage selected by feat(vlm): add a token-exact Gemma3n audio reference path #875: mel/mask input, both subsample blocks/cumulative stats, relative-position and chunk masks, selected Conformer blocks, pooled/projected audio tokens, placeholder merge, dense PLE, selected Gemma3n language layers, KV/logits, and greedy tokens.
• Cover duration/frame bucket boundaries, padded batches, one/multiple clips, malformed masks, all-padded rejection, context overflow after audio expansion, audio-only, text/image-only, qualified mixed media, cancellation, and slot reuse.
• Interleave requests with different clip lengths and modalities; prove masks, effective lengths, and dense PLE remain slot-local.
• Run deterministic ASR through CLI and streaming server on a production-relevant IREE target with token-exact output.
• Run feat(vlm): add a token-exact Gemma3n audio reference path #875/feat(xla): add the Gemma3n text backbone and dense PLE prefill #876 and all existing Gemma3n/text/image/XLA gates.

Acceptance criteria

• Gemma3n audio encoder/projection executes entirely through IREE math after host mel extraction.
• Cumulative normalization, local attention, relative shift, masks, pooling, and projected lengths match feat(vlm): add a token-exact Gemma3n audio reference path #875 intermediates.
• Merged embeddings and dense PLE match feat(vlm): add a token-exact Gemma3n audio reference path #875 and seed feat(xla): add the Gemma3n text backbone and dense PLE prefill #876 KV/logits exactly.
• Frame/context bounds and all mask/length/cardinality invariants fail safely before native misuse.
• Audio capability and conditional weight loading reflect the actual loaded bundle.
• CLI/server deterministic output is token-exact and per-slot state is isolated.
• Existing Gemma3n text/image behavior remains token-exact.

[inureyes]

Wave 4A implementation is locally complete but is not ready for review or merge because the required parity gate remains unresolved.

Evidence:

• The production E4B run matched the first 11 generated tokens and first diverged at token 12.
• A bounded audio-only comparison matched projected lengths at [87] but found projected_audio max abs 0.125 / RMS 0.01098145222; hard_audio independently differed by max abs 0.09375.
• The branch now captures encoded_reduced and the soft/hard norm-linear-postnorm stages without loading the 30 GB language bundle, and all focused compile, cache/finite, Phi4MM, admission, and ownership tests pass.
• The first CUDA diagnostic binary was built for sm_90a because sandboxed nvidia-smi detection failed; it cannot run on the GB10 target (no kernel image is available).
• An explicit MLX_CUDA_ARCHITECTURES=121a rebuild completed native compilation but timed out in the final root link at the 5-minute bound. A single cached incremental link also timed out at 2 minutes. No compatible diagnostic binary was published and no further retries were made.

The local branch remains unpushed. Resume by producing the 121a diagnostic example in a build environment with a suitable linker/time budget, then run the graph-disabled audio-only stage gate once. Do not run the full E4B language bundle until the first divergent audio operation is identified and corrected.