lablup / lablup/mlxcel

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

Open
#878 2 comments 0 reactions 0 assignees View on GitHub
area:architecture area:inference area:models priority:low status:blocked type:enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.