feat(xla): add Gemma3n audio execution#922
Draft
inureyes wants to merge 9 commits into
Draft
Conversation
Add fail-closed audio configuration and artifact identities for static frame buckets. Validate canonical mel masks, fixed placeholder spans, and owned prepared-prefill metadata before requests cross the IREE boundary.
Pin the converted Q4 checkpoint's 277 audio tensor names, dtypes, shapes, and deterministic order. Reject missing, unexpected, or incompatible tensors before enabling the audio capability.
Implement the complete Gemma3n Conformer audio path with exact cumulative normalization, local relative attention, masking, projection, row mapping, and shared dense-PLE rounding semantics. Split the production graph at the natural typed boundary: audio.encode owns audio tower/embedder weights and returns projected soft rows, hard audio embeddings, and lengths; audio.merge_ple owns only the shared language input-head weights. This avoids pinned IREE monolithic-module backend failures while keeping all model math in IREE. Validate the full mlxcel-xla library and compile the maximum 2997-frame/context-256 artifacts with pinned IREE for both local CPU and CUDA targets. Refs #878
Execute the split audio encode and merge graphs through IREE and wire them into CLI and continuous batching. Preserve pinned processor left-padding mel features while zeroing only the static bucket tail. Refs #878
Keep the rebased LLaVA vision runtime on the shared aligned auxiliary-weight storage API used by Gemma3n audio. Refs #878
MLX CUDA materializes the quantized product as BF16 before adding the BF16 affine bias, while the XLA loader previously fused both operations before a single rounding step. This changed roughly half of the pinned Gemma3n audio embedding and projection weights and explains the tower-independent hard-audio mismatch. Use sequential BF16 rounding in both Gemma3n language and audio weight loaders, rename the helper to reflect its semantics, and cover a real projection double-rounding counterexample. Validation: cargo test -p mlxcel-xla weights::tests --lib; cargo check -p mlxcel-xla --features diagnostics; cargo clippy -p mlxcel-xla --features diagnostics --lib --fix --allow-dirty --allow-staged; cargo fmt --all -- --check. Refs #878
Make the MLX checkpoint-dtype boundary explicit by rounding processor-produced F32 mel features to BF16 before the first SSCP convolution, independent of the graph-wide contraction precision. Add a synthetic regression whose pre-cast and post-convolution-only results differ, and expose the first SSCP output through the diagnostics-only MLX and IREE bundles for one bounded parity check. Validation: cargo fmt --all -- --check; git diff --check; focused Gemma3n MLX CPU test; focused Gemma3n audio emitter tests; mlxcel-xla diagnostics cargo check. Refs #878
inureyes
added
status:blocked
7 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Validation
Remaining blocker
This PR remains draft because the pinned MLX-to-IREE audio oracle still fails. The latest single bounded run localizes the first captured divergence to sscp_conv_0 with max abs 1.0 and RMS 0.02987759543; encoded_reduced remains max abs 8.0 and the final dense PLE remains outside tolerance. The explicit mel BF16 cast is now present, so the next investigation is the first convolution BF16 accumulation and operation schedule. Full language generation and identical diagnostic retries must not be repeated until that boundary is reconciled.
Closes #878
Parent epic: #566