Fix gemma4 audio-encoder audio_features export rank (#570) - #444
Draft
justinchuby wants to merge 1 commit into
Draft
Fix gemma4 audio-encoder audio_features export rank (#570)#444justinchuby wants to merge 1 commit into
justinchuby wants to merge 1 commit into
Conversation
The gemma4 audio dataflow edge
`audio_encoder.audio_features -> embedding.audio_features` was rank-3 on the
producer side but rank-2 on the consumer side, so onnx-genai pipeline admission
rejected every gemma4 audio (E2B/E4B) package:
'audio_encoder.audio_features -> embedding.audio_features' has incompatible
ranks: producer rank 3, consumer rank 2.
HF selects `audio_features[audio_mask]` (a boolean-mask gather) that both drops
padding frames and flattens the batch axis to rank-2
`[num_valid_frames, text_hidden_size]` before scattering rows into the
language-model embeddings.
Move that mask-selection into `_Gemma4AudioEncoderModel.forward`: after the
projector, flatten `[B, T//4, H] -> [-1, H]` and reuse the shared bf16-safe
`_dtype_safe_compress` helper to emit rank-2 `[num_valid_frames, H]`, matching
the `embedding` consumer and the encoder-free `_Gemma4UnifiedAudioEmbedderModel`.
The task builder `_build_audio` no longer duplicates the flatten+Compress. This
covers all gemma4 Conformer-audio variants since they share the module.
Add `TestGemma4AudioEncoderRank` locking the rank-2 contract: producer rank ==
consumer rank == 2 (trailing dim == hidden_size) and a Compress node is present
so the reduction is a real mask-select, not a bare reshape.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
Fixes Gemma4 audio package export incompatibility with onnxruntime-genai by ensuring the audio_encoder sub-model emits rank-2 audio_features ([num_valid_frames, hidden]) that match the embedding sub-model’s expected input rank, mirroring HuggingFace’s audio_features[audio_mask] semantics.
Changes:
- Move padding-frame stripping + batch/time flattening into
_Gemma4AudioEncoderModel.forward()via_dtype_safe_compress, so theaudio_encoderoutput is rank-2 when a mask is provided. - Remove duplicate flatten+Compress logic from
Gemma4Task._build_audio()and rely on the encoder module’s rank-2 contract. - Add regression tests asserting producer/consumer rank agreement and presence of a
Compressop in the exportedaudio_encodergraph.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| src/mobius/tasks/_gemma4.py | Removes task-level flatten+Compress; assumes audio_encoder already emits rank-2 features. |
| src/mobius/models/gemma4.py | Implements mask-based selection inside the audio encoder model to produce rank-2 audio_features. |
| src/mobius/models/gemma4_test.py | Adds tests to lock the rank-2 contract and ensure a real mask-select (Compress) exists in the audio graph. |
Comment on lines
2644
to
+2648
| Output: | ||
| - ``audio_features [B, T//4, text_hidden_size]``: projected audio tokens | ||
| - ``audio_features [num_valid_frames, text_hidden_size]``: projected audio | ||
| tokens with padding frames stripped so output rows align 1:1 with audio | ||
| placeholder tokens. This mirrors HF, which selects | ||
| ``audio_features[audio_mask]`` (a rank-2 gather) before scattering the |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Member
Author
|
✅ Coordinator-verified, ready for your merge (not self-merged, per mobius policy). Checked:
Suggest a quick manual e2e sanity (export |
justinchuby
marked this pull request as draft
August 1, 2026 04:29
Member
Author
|
This may break genai. Need to be careful |
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
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.
Problem (onnx-genai #570)
Every gemma4 (E2B/E4B) audio package is rejected at onnx-genai pipeline admission because the audio dataflow edge has mismatched ranks:
The
audio_encodersub-model emitted rank-3[B, T//4, H], whileembeddingconsumes rank-2[num_valid_frames, H](its scatter expects per-token rows). The shippedgemma4-e2b-onnxexport contains 0 Compress ops and ends at the projector MatMul — raw rank-3.HF reference semantics
HuggingFace gemma4 selects
audio_features[audio_mask]— a boolean-mask gather that both drops padding frames and flattens the batch axis, producing rank-2[num_valid_frames, text_hidden_size]rows that align 1:1 with audio placeholder tokens before they are scattered into the language-model embeddings.Fix
Move the mask-selection into
_Gemma4AudioEncoderModel.forward: after the projector, flatten[B, T//4, H] -> [-1, H]and reuse the shared bf16-safe_dtype_safe_compresshelper to emit rank-2[num_valid_frames, H], while still returningdownsampled_maskfor theaudio_features_maskdiagnostic output. The task builder_build_audiono longer duplicates the flatten+Compress._Gemma4UnifiedAudioEmbedderModel(which already emitted rank-2) and covers all gemma4 Conformer-audio variants since they share_Gemma4AudioEncoderModel— no E2B-specific hack.vision_encoderflattens batch); no change there.Test
New
TestGemma4AudioEncoderRank(non-vacuous):audio_encoder.audio_featuresrank ==embedding.audio_featuresinput rank == 2, trailing dim ==hidden_sizeon both ends;Compressnode exists in the audio graph, so the reduction is a real mask-select, not a bare reshape that would misalign rows with placeholder tokens.Validation
pytest src/mobius/models/gemma4_test.py→ 20 passed (incl. 2 new)pytest src/mobius/components/_gemma4_audio_test.py→ 31 passedruff checkon changed files → cleanCloses #570.