feat(vlm): add a token-exact Gemma3n audio reference path

[inureyes]

Parent

• Parent epic: epic: OpenXLA multimodal / VLM architecture support (Qwen2-VL, Gemma3n, Phi4MM, Molmo, and more) #566
• Current gap: src/loading/vlm_gemma.rs::load_gemma3n_vlm loads only the text, MobileNetV5 vision, and multimodal vision embedder paths; no Gemma3n audio processor/tower/merge is present.

Normative upstream sources

• Official capability/format: https://ai.google.dev/gemma/docs/gemma-3n and https://ai.google.dev/gemma/docs/capabilities/audio
• Maintained model definition: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma3n/modular_gemma3n.py
• Official checkpoint: https://huggingface.co/google/gemma-3n-E4B
• Pin exact model and source revisions in tests/docs.

Problem

The current Gemma3n implementation covers its specialized text backbone, per-layer embeddings, and MobileNetV5 vision path, but it has no audio reference implementation. Official Gemma3n audio uses its own mel masks, convolutional subsampling/cumulative normalization, chunked Conformer attention with relative positions and softcapping, light convolution, projection/pooling, and multimodal token merge. It cannot be inferred from the existing Gemma4 audio module or from the vision PLE path.

Goal

Implement and qualify a complete MLX Gemma3n audio-to-text reference path, including official waveform processing, audio encoder, projection/merge, per-layer-input interaction, CLI, and serving.

Required implementation

• Implement the pinned official processor contract for mono 16 kHz float32 waveforms, mel features/masks, clip padding, placeholder expansion, and duration limits. The official docs currently state a 30-second maximum and approximately 6.25 input tokens per second for Gemma3n; assert the pinned processor's exact frame/token math rather than using those rounded values as code constants.
• Port the official audio subsample-convolution projection, manual padding, cumulative group normalization, activation, and mask/length propagation.
• Port Conformer feed-forward, chunked local attention, relative-position shift, per-dimension query scale, logit softcap, light causal/depthwise convolution, normalization, gradient-clipping semantics relevant to inference, and final projection/pooling.
• Load audio config/weights conditionally and validate every convolution, attention, normalization, and projection tensor. Text/vision-only loading must still be able to skip audio parameters intentionally.
• Reproduce official audio token expansion and merge into scaled text embeddings. Preserve multiple-clip boundaries and verify the number/order of valid projected audio frames against placeholder positions.
• Generate/project Gemma3n per-layer inputs exactly as the official multimodal forward does. Keep this state owned by the request/sequence and consume/release it once during prefill; do not reuse the single fallback slot across concurrent requests.
• Integrate audio prompt rendering inside the user turn, CLI/server preparation, capability reporting, cancellation, and no-media/image paths.
• Keep waveform, mel, encoder, projected-frame, and PLE tensor ownership explicit and release large intermediates after prefill.

Non-goals

• XLA/StableHLO emission.
• Reusing Gemma4 audio merely because class names overlap.
• Speech output generation.

Validation

• Pin an official Gemma3n checkpoint/revision and deterministic ASR waveform. Compare waveforms after normalization, mel/mask tensors, both subsample blocks, cumulative norm states, selected Conformer block outputs, pooled/projected audio tokens, merge positions, per-layer inputs, selected text-layer states, logits/KV, and greedy tokens against the pinned maintained implementation.
• Cover one/multiple clips, short and duration-boundary clips, padding, malformed/nonfinite audio, placeholder mismatch, audio-only, image-only, mixed image+audio if officially supported, text-only, and concurrent sequence-state isolation.
• Add negative fixtures for ordinary group norm instead of cumulative norm, global instead of chunked attention, missing relative shift/softcap, wrong token-rate rounding, and stale PLE state.
• Exercise CLI and OpenAI-compatible streaming server paths and run all existing Gemma3n image/text regressions.

Acceptance criteria

• Processor frame/mask/token math matches the pinned official revision.
• Audio subsampling, cumulative normalization, Conformer, pooling/projection, and merge intermediates match within documented tolerances.
• Per-layer input state is request-scoped and concurrency-safe.
• Deterministic ASR output is token-exact through CLI and streaming server.
• Audio parameters can be conditionally loaded/skipped without corrupting text/vision modes.
• No Gemma4-specific audio assumptions leak into Gemma3n.
• Existing Gemma3n text/image behavior remains token-exact.