feat(xla): support Gemma3n audio with dense PLE prefill
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Parent and dependencies
- Parent epic: #566
- Blocked by #873, #875, and #876
- Qualified MLX/oracle contract: #875
- Required XLA language/PLE contract: #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 #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 #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 #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 #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 #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 #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 #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 #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 #875/#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 #875 intermediates.
- [ ] Merged embeddings and dense PLE match #875 and seed #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.
Contributor guide
Research direction
Start by reading the qualified contracts in #875 and #876, then trace the audio.main and prefill_embeddings_ple.main entry points. Use the validation list to compare masks, projected lengths, merged embeddings, dense PLE, KV/logits, and greedy tokens. Done means token-exact CLI/server behavior, safe bounds handling, conditional audio loading, and isolated per-slot state.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- audio-video-rtc, cli, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100