Fix L4/L5 e2e harness for Gemma4 per_layer_inputs + BOOL audio mask - #318
Merged
Conversation
… mask After #296 moved Gemma4's per-layer input computation from the text decoder into the embedding sub-model, the embedding model now emits a second output (per_layer_inputs) and the decoder accepts it as a required input. The e2e harness only wired the first embedding output (inputs_embeds) into the decoder, so every multi-model Gemma4 path (text-only on multi-model, VL prefill, VL generation, speech-language prefill, speech-language generation) failed with: ValueError: Required inputs (['per_layer_inputs']) are missing from input feed (['inputs_embeds', ...]) Fix by passing any extra embedding outputs through to the decoder by name. This is generic — for models without per_layer_inputs the extra loop iteration just no-ops. Separately, the speech-language audio encoder builds Gemma4's input_features_mask as tensor(bool), but the harness unconditionally cast feature-extractor outputs to float32 (and constructed the fallback all-True mask as np.bool_ that then crashed ort_easy's DLPack path, which has no bool type code). Fix by: - Honoring the session's declared input dtype for each audio-encoder input (BOOL stays BOOL, FLOAT becomes float32, etc.). - Routing bool numpy arrays through OrtValue.ortvalue_from_numpy directly in OnnxModelSession, bypassing ort_easy's DLPack-first path. Verified locally on H200: L4 text-generation/gemma-4-e2b PASS L4 image-text-to-text/gemma-4-e2b-it PASS L4 speech-language/gemma-4-e2b-it-audio PASS L5 image-text-to-text/gemma-4-e2b-it PASS L5 speech-language/gemma-4-e2b-it-audio PASS (L5 text-generation/gemma-4-e2b is gated by integration markers and skipped under fast runs.) ruff check + format pass. Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Fixes the L4/L5 end-to-end golden test harness for Gemma4 multi-model pipelines by correctly wiring the embedding model’s extra outputs into the decoder feed and by respecting declared ORT input dtypes for audio inputs (notably tensor(bool) masks). Also updates the ORT inference wrapper to safely handle bool NumPy arrays without going through a DLPack conversion path that can’t represent bool.
Changes:
- Pass through additional embedding outputs (e.g.
per_layer_inputs) to the decoder by matching decoder input names. - Cast audio encoder inputs (including fallback
input_features_mask) to the session-declared input dtype instead of forcingfloat32. - Bypass
onnxruntime_easy’s DLPack-first conversion for bool NumPy arrays by constructingOrtValuedirectly.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| tests/e2e_golden_test.py | Wires embedding outputs into decoder feeds across L4/L5 paths and casts audio inputs/masks using session-declared dtypes. |
| src/mobius/_testing/ort_inference.py | Adds bool-aware NumPy→OrtValue conversion to avoid DLPack limitations for tensor(bool) inputs. |
Performance Comparison
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
justinchuby
added a commit
that referenced
this pull request
May 28, 2026
## Summary After #296 ("Move Gemma4 per-layer embeddings to embedding model") the Gemma4 embedding sub-model emits a second output `per_layer_inputs` (`[B, S, num_layers * per_layer_dim]`) that the decoder consumes on every step. `examples/gemma4_multimodal.py` wasn't updated to forward this output, so any Gemma 4 build with `hidden_size_per_layer_input > 0` (e.g. `google/gemma-4-E2B-it`) crashes at the first `decoder.run()`: ``` ValueError: Required inputs (['per_layer_inputs']) are missing from input feed (['inputs_embeds', ..., 'attention_mask', 'position_ids']) ``` ## Fix - `prepare_decoder_feeds()`: accept an optional `per_layer_inputs` argument and add it to the feeds dict when present. Builds without per-layer inputs (`hidden_size_per_layer_input == 0`, e.g. larger Gemma 4 variants) keep working — the kwarg is optional. - `generate()`: pull `embed_out.get("per_layer_inputs")` and pass it to `prepare_decoder_feeds` on every step. ## Verification ``` $ python examples/gemma4_multimodal.py --mode text --prompt "What is 2+2?" ... 📝 TEXT-ONLY GENERATION ================================================================ Prompt: What is 2+2? ---------------------------------------------------------------- 2 + 2 = **4** ``` ## Related This is the example-side counterpart to #318 which fixed the same gap in the L4/L5 e2e test harness. --------- Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
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
L4 / L5 e2e tests have been failing on
mainfor every Gemma4 case since#296 (Move Gemma4 per-layer embeddings to embedding model). Two distinct
root causes:
1.
per_layer_inputsmissing from decoder feedAfter #296 the embedding sub-model emits a second output
per_layer_inputsand the decoder accepts it as a required input. Thee2e harness only wired the first embedding output (
inputs_embeds)into the decoder, so every multi-model Gemma4 path failed with:
Affected: text-only on multi-model (
text-generation/gemma-4-e2b),VL prefill (
image-text-to-text/gemma-4-e2b-it,…-e4b-it), VLgeneration (L5), speech-language prefill+generation.
2.
input_features_maskdtype mismatchThe Gemma4 audio encoder (
_gemma4.py:367) declaresinput_features_maskastensor(bool), but the harness unconditionallycast every feature-extractor output to
np.float32, producing:Even when the fallback all-True mask path was taken (line 1267), the
bool numpy array crashed
ort_easy's DLPack-first conversion pathbecause DLPack has no native bool type code.
Fix
tests/e2e_golden_test.pyIn every decoder-feed setup (5 sites: text-only multi-model prefill,
VL prefill, VL generation, speech-language prefill, speech-language
generation), wire any extra embedding outputs through to the decoder by
name. For models without
per_layer_inputsthe extra loop iteration isa no-op:
For both audio-encoder setup sites (L4 + L5 paths), honor the session's
declared input dtype:
And use the session-declared dtype for the fallback all-True mask too.
src/mobius/_testing/ort_inference.pyRoute bool numpy arrays through
OrtValue.ortvalue_from_numpydirectlyin
_numpy_to_ort_value, bypassingort_easy's DLPack-first path(which has no bool type code).
Verification (locally on H200)
text-generation/gemma-4-e2bimage-text-to-text/gemma-4-e2b-itspeech-language/gemma-4-e2b-it-audioimage-text-to-text/gemma-4-e2b-itspeech-language/gemma-4-e2b-it-audio5 passed, 1 skipped (L5 text-gen is gated by integration markers under
fast runs), 0 failed.
ruff check+ruff format --checkboth pass.L1 + L3 + non-gemma e2e suites are unaffected (the embedding-output
loop is a generic no-op for models that don't emit extra outputs;
audio-mask changes only trigger when the session declares a bool input).
Why this didn't get caught earlier
The CI matrix only runs L4 / L5 in the affected-models lane, and #296
was tested standalone before this lane started exercising the multi-
model Gemma4 path through the e2e harness. The bool-mask issue is even
older — it was masked by the harness's unconditional
astype(np.float32)until the audio_encoder was upgraded to take a real BOOL mask.