Support GraniteMoeHybrid dense variant (granite-4.0-1b) with L4/L5 tests - #440
Merged
Conversation
Extend the existing GraniteMoeHybrid model to support the dense
granite-4.0-1b variant (ORT-genai #2124): all full_attention layers,
RoPE, and no routed experts. Adds working L4/L5 golden tests.
What changed and why:
- Config (_base.py): GraniteMoeHybridConfig.from_transformers now
rebuilds layer_types from GraniteMoeHybrid's own "full_attention" /
"linear_attention" naming instead of inheriting Bamba's
"attention"/"mamba" mapping, which mis-classified every
granite-4.0-1b layer as mamba2. num_local_experts, RoPE params
(rope_theta from rope_parameters), and the Granite scaling
multipliers are already carried through the base config.
- Model (granitemoehybrid.py):
- MoE is now optional: block_sparse_moe is built only when
num_local_experts > 0; otherwise the feedforward is the dense
shared_mlp alone (mirrors HF GraniteMoeHybridDecoderLayer).
- RoPE support: the text model builds a rotary embedding via
initialize_rope (present when position_embedding_type='rope',
None for NoPE variants like tiny-preview) and threads
(cos, sin) to the attention layers.
- Applies the previously missing Granite multipliers:
embedding_multiplier (after embed), attention_multiplier (as the
attention scale), and logits_scaling (final logits).
- Golden generation (generate_golden.py / torch_reference.py): all
full_attention hybrid models trip transformers' recurrent-mask
builder ("has_previous_state can only be called on LinearAttention
layers"). Both the single-pass forward and greedy generate now fall
back to use_cache=False on that specific error; greedy/logits output
is cache-independent.
- Tests: new testdata/cases/causal-lm/granite-4_0-1b.yaml (L4+L5) with
generated golden. L4 prefill argmax and L5 20-token generation both
match the HuggingFace reference on CUDA.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
🟡 Not ready to approve
granitemoehybrid.py now passes position_embeddings=None for NoPE variants but keeps non-optional tuple type annotations in layer forward() signatures, which is internally inconsistent and should be corrected before merge.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Pull request overview
Adds mobius support for the dense GraniteMoeHybrid variant ibm-granite/granite-4.0-1b (all full_attention, RoPE, num_local_experts=0) and introduces L4/L5 golden coverage to validate CUDA parity against HuggingFace.
Changes:
- Fix
GraniteMoeHybridConfig.layer_typesreconstruction to correctly interpret Granite’sfull_attention/linear_attentionnaming. - Update
GraniteMoeHybridmodel implementation: make MoE optional (shared MLP-only whennum_local_experts == 0), initialize/apply RoPE when configured, and apply Granite scaling multipliers (embedding/attention/logits). - Make golden generation / HF reference execution resilient to transformers’ hybrid-cache mask error by retrying with
use_cache=Falseon the specifichas_previous_statefailure.
File summaries
| File | Description |
|---|---|
| testdata/golden/causal-lm/granite-4_0-1b.json | Adds L4 (prefill logits/top-k) golden reference for granite-4.0-1b. |
| testdata/golden/causal-lm/granite-4_0-1b_generation.json | Adds L5 (20-token greedy generation) golden reference for granite-4.0-1b. |
| testdata/cases/causal-lm/granite-4_0-1b.yaml | New test case enabling L4+L5 coverage for the new checkpoint. |
| src/mobius/models/granitemoehybrid.py | Implements dense/no-experts path, RoPE wiring, and Granite multipliers in the model. |
| src/mobius/_testing/torch_reference.py | Retries HF forward without cache on the has_previous_state hybrid-cache error. |
| src/mobius/_configs/_base.py | Rebuilds layer_types using Granite’s naming to avoid misclassification. |
| scripts/generate_golden.py | Retries HF generate() with use_cache=False on the same transformers hybrid-cache error. |
Review details
Comments suppressed due to low confidence (1)
src/mobius/models/granitemoehybrid.py:323
- This forward method accepts/passes through
position_embeddings=Nonefor NoPE variants, but the type annotation istuple(non-optional). Update the annotation totuple | Noneto matchAttention.forwardand avoid MyPy errors at the call site.
def forward(
self,
op: OpBuilder,
hidden_states: ir.Value,
attention_bias: ir.Value,
position_embeddings: tuple,
past_key_value: tuple | None,
):
- Files reviewed: 7/7 changed files
- Comments generated: 1
- Review effort level: Low
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
Address review feedback on GraniteMoeHybrid config conversion:
- NoPE regression fix: granite-4.0-tiny-preview declares
position_embedding_type='nope' but still ships a default
rope_parameters dict, so driving RoPE purely off initialize_rope()
would have wrongly enabled RoPE for the NoPE variant. Now clear
rope_type (=> initialize_rope returns None => NoPE) whenever
position_embedding_type != 'rope'. granite-4.0-1b (='rope') is
unaffected.
- layer_types mapping is now explicit: {full_attention, attention} ->
full_attention, {linear_attention, mamba, mamba2} -> mamba2, and any
unknown value raises instead of silently defaulting to mamba2.
- Widen decoder-layer forward() position_embeddings annotation to
`tuple | None` (matches Attention.forward and the NoPE None path).
L4/L5 granite-4.0-1b golden tests still pass on CUDA; granite/bamba
unit builds unaffected.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Member
Author
|
Addressed the review feedback in d64b960:
L4/L5 granite-4.0-1b golden tests still pass on CUDA. |
justinchuby
enabled auto-merge (squash)
July 30, 2026 19:04
Contributor
There was a problem hiding this comment.
🟢 Ready to approve
The changes are cohesive (config + model + golden/reference infra) and are validated by new L4/L5 goldens that exercise the newly supported dense GraniteMoeHybrid behavior.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
- Files reviewed: 7/7 changed files
- Comments generated: 0 new
- Review effort level: Low
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
xiaoyu-work
approved these changes
Jul 30, 2026
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
Adds mobius support for the dense GraniteMoeHybrid variant
ibm-granite/granite-4.0-1b, the model targeted by onnxruntime-genai#2124. Unlikegranite-4.0-tiny-preview, this variant is allfull_attention(no Mamba), uses RoPE, and has no routed experts (num_local_experts=0) — so it runs natively in transformers and a golden can be generated in CI.L4 (prefill argmax) and L5 (20-token greedy generation) golden tests both pass on CUDA with exact parity to the HuggingFace reference.
What was broken for this variant
GraniteMoeHybridConfiginherited Bamba's"attention"/"mamba"naming, mapping everyfull_attentionlayer tomamba2.block_sparse_moeasserted experts exist; granite-4.0-1b has none.position_embeddings=None, but granite-4.0-1b usesposition_embedding_type='rope'(θ=1e7).embedding_multiplier,attention_multiplier, andlogits_scalingwere not applied (latent, since tiny-preview was skipped).Changes
_configs/_base.py: rebuildlayer_typesfrom GraniteMoeHybrid'sfull_attention/linear_attentionnaming.models/granitemoehybrid.py:shared_mlp-only feedforward whennum_local_experts == 0(mirrors HF).initialize_rope(present forrope,Nonefor NoPE variants), threaded to attention layers.embedding_multiplier,attention_multiplier(attention scale),logits_scaling.scripts/generate_golden.py/_testing/torch_reference.py: all-attention hybrids trip transformers' recurrent-mask builder (has_previous_state can only be called on LinearAttention layers); both forward and generate now fall back touse_cache=Falseon that specific error (greedy/logits output is cache-independent).testdata/cases/causal-lm/granite-4_0-1b.yaml(L4+L5) + generated golden.Verification
TestL4CheckpointVerified+TestL5GenerationE2Efortext-generation/granite-4_0-1b— 2 passed on CUDA.build_graph_test.py -k "granitemoehybrid or granite or bamba"— 26 passed, 6 skipped (NoPE tiny-preview path unaffected).ruff format+ruff checkclean.Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com