Fix L4/L5 golden regressions across 7 broken models - #436
Merged
Conversation
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes multiple L4/L5 golden regressions by addressing model-specific build/runtime issues across several architectures, plus documenting intentional golden-test skips where a valid HF reference cannot be produced in CI.
Changes:
- Fix attention-mask/bias shape handling in DiffLlama and Gemma4 assistant to satisfy ORT Attention/KV-shape constraints.
- Fix Gemma4 audio masking/downsampling and
Rangescalar-limit issues to avoid shape-inference overflow and runtime errors. - Improve config resolution for composite configs (e.g., Qwen3-ASR
thinker_configdict) and harden linear-attention cache config validation; document skip reasons for unsupported CI golden generation cases.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| testdata/cases/speech/sensevoice-small.yaml | Adds skip_reason documenting why SenseVoice audio-ctc L4/L5 isn’t currently runnable in the e2e harness. |
| testdata/cases/causal-lm/nemotron-h-nano-4b.yaml | Enables trust_remote_code and documents golden-generation skip due to unavailable HF reference dependencies in CI. |
| src/mobius/tasks/_cache_utils.py | Adds explicit validation/error messaging when linear-attention cache dims are requested but config fields are unset. |
| src/mobius/models/diffllama.py | Fixes padding-mask bias rank/broadcasting by adding already-4D bias directly to attention scores. |
| src/mobius/components/_gemma4_audio.py | Fixes Range limit rank and bounds Slice ends to dynamic lengths to prevent shape-inference overflow. |
| src/mobius/components/_gemma4_assistant.py | Slices buffer-width attention mask to shared-KV length to avoid inconsistent total_sequence_length failures. |
| src/mobius/_builder.py | Converts dict-valued composite sub-configs to PretrainedConfig to reliably unwrap text_config. |
Comment on lines
49
to
+52
| head_k_dim = config.linear_key_head_dim | ||
| head_v_dim = config.linear_value_head_dim | ||
| conv_kernel = config.linear_conv_kernel_dim | ||
| missing = [ |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
An L4/L5 golden-test regression sweep found seven models failing at build or run time. Root-cause and fix each: - diffllama: create_padding_mask now returns a 4D bool [B,1,S,Sk] mask, but DiffLlamaAttention still unsqueezed it again, producing a rank-5 bias that broke ORT Split. Add the (already-4D) pad bias directly. - gemma-4 audio encoder (e2b + e4b): (1) op.Range got a 1-D [1] limit from op.Shape; squeeze to a scalar at both Range sites. (2) the mask-downsample Slice uses an INT64_MAX "to-the-end" sentinel with step 2; onnx-shape- inference < 0.3.1 did not normalize the sentinel for strided slices and derived a 2^62 extent that overflowed the downstream Reshape at load. Bump the dependency to onnx-shape-inference>=0.3.1 (which normalizes a positive- sentinel end to ceil(dim/step)) and keep the natural [INT64_MAX] slice. - gemma-4 assistant drafter: the generic Attention path built the sliding- window bias from the full buffer-width attention_mask (kv_len + 1) while K is only the shared KV (kv_len), so ORT rejected an inconsistent total_sequence_length. Slice the mask to kv_len first. - qwen3-asr / qwen3-asr-en: thinker_config ships as a plain dict, so the text_config-unwrap branch in build() missed it and the decoder config kept hidden_size=0. Convert the dict via _dict_to_pretrained_config before unwrapping text_config. - nemotron-h: build() without trust_remote_code mis-detected layer_types as linear_attention (TypeError). Add trust_remote_code and a clear None-guard in linear_attention_dims. The build now succeeds with correct mamba2 layer types, but the HF reference needs trust_remote_code custom modeling plus mamba-ssm CUDA kernels that are unavailable in CI golden generation (same as jamba-v0_1 / granitemoehybrid-tiny), and the previously committed golden was generated without trust_remote_code (wrong native config). Documented skip_reason added until a valid golden can be produced. - sensevoice-small: audio-ctc L4/L5 is not wired into the e2e harness (custom LFR-fbank input_features + language_id, FunASR reference, no audio-ctc golden generator). Documented skip_reason added. Verified on GPU (CUDA EP): diffllama, qwen3-asr, qwen3-asr-en, gemma-4-e2b-it-audio, gemma-4-e4b-it-audio, gemma-4-e2b-it-assistant all pass L4 and L5; nemotron-h and sensevoice-small skip. Regression-checked gemma-4-e2b, llama-3_2-1b, qwen3-0_6b (unchanged, still pass). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
fix/l4l5-broken-models
branch
from
July 30, 2026 16:32
18c4c8b to
db9a470
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
src/mobius/tasks/_cache_utils.py:69
- The new
missinglist check doesn't narrow theOptional[int]config fields for mypy, sokey_dim = head_k_dim * num_k_heads(and similar) will still be typed asOptional[int] * Optional[int]understrict_optional, causing a type-check failure. Also, the docstring still says it raisesTypeError, but the new code raisesValueError.
missing = [
name
for name, value in (
("linear_num_key_heads", num_k_heads),
("linear_num_value_heads", num_v_heads),
("linear_key_head_dim", head_k_dim),
("linear_value_head_dim", head_v_dim),
("linear_conv_kernel_dim", conv_kernel),
)
if value is None
]
if missing:
raise ValueError(
"Linear-attention cache inputs requested but the following config "
f"fields are unset: {', '.join(missing)}. This usually means the "
"architecture's layer_types were mis-detected as 'linear_attention' "
"(e.g. a hybrid model built without trust_remote_code)."
)
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
src/mobius/tasks/_cache_utils.py:64
- This guard now raises
ValueError, but the function docstring (above) documents aTypeErrorfor missing required config fields. To keep the public contract consistent (and avoid surprises for callers that may be catchingTypeError), consider raisingTypeErrorhere instead, or update the docstring in a follow-up.
raise ValueError(
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.
Summary
An L4/L5 golden-test regression sweep on
main(CI had not run for a while) found 7 models failing at build or run time. This PR root-causes and fixes each one, then verifies L4 (prefill argmax) and L5 (generation) on GPU (CUDA EP).Fixes
create_padding_maskreturns 4D[B,1,S,Sk], but attention re-Unsqueezed it → rank-5 bias broke ORTSplitop.Rangegot 1-D[1]limit fromop.Shape; (2) mask-downsampleSliceusedINT64_MAXends → shape inference derived 2^62 length →Reshapeoverflow at loadSliceends to the dynamic time dimkv_len+1) while K is only shared KV (kv_len) → inconsistenttotal_sequence_lengthkv_lenfirstthinker_configships as a plain dict, sotext_config-unwrap inbuild()missed it → decoderhidden_size=0_dict_to_pretrained_configbefore unwrappingbuild()withouttrust_remote_codemis-detectedlayer_typesaslinear_attention(TypeError)trust_remote_code+ clear None-guard inlinear_attention_dims; documentedskip_reason(HF ref needstrust_remote_code+ mamba-ssm CUDA kernels, unavailable in CI — same as jamba/granitemoehybrid; prior golden was generated with the wrong native config)language_id, FunASR reference, no golden generator)skip_reasonVerification (GPU, CUDA EP)
Notes
onnx_irAPIs only.ruff format/ruff checkclean.Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com