Fix Gemma4 bidirectional attention + add gemma-4-12B (gemma4_unified) - #338
Merged
Conversation
Performance Comparison
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends Gemma 4 support in mobius in two directions: (1) correctly replicating HuggingFace’s “vision-block bidirectional” attention behavior for multimodal Gemma4 decoders, and (2) adding the encoder-free gemma-4-12B (“gemma4_unified”) multimodal architecture (decoder + vision/audio embedders + embedding fusion) with appropriate config extraction and tests.
Changes:
- Add
use_bidirectional_attention="vision"plumbing viablock_sequence_ids, enabling a blockwise bidirectional overlay baked into a float attention bias and forcingAttention(is_causal=0)when active. - Introduce
Gemma4UnifiedModel+Gemma4UnifiedTaskand config hooks for unified vision/audio embedders (encoder-free). - Add L1 graph-build coverage plus new unit/integration tests for block ids, bias correctness, and unified parity.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/integration_test.py | Adds slow integration parity tests for gemma-4-12B text/multimodal and an HF-vs-mobius bidirectional mask parity test. |
| tests/build_graph_test.py | Adds L1 graph-build tests for gemma4_unified_text and gemma4_unified, and updates registry completeness coverage. |
| src/mobius/tasks/_gemma4.py | Wires block_sequence_ids input into the decoder build and updates embedding builder to handle dict outputs; adds Gemma4UnifiedTask. |
| src/mobius/tasks/init.py | Exports/registers the new Gemma4UnifiedTask. |
| src/mobius/models/gemma4.py | Implements block-sequence-id computation, bidirectional overlay behavior, is_causal threading for Attention, and adds unified vision/audio embedders + Gemma4UnifiedModel. |
| src/mobius/models/gemma4_test.py | Adds unit tests for _compute_block_sequence_ids and end-to-end package wiring + is_causal=0 enforcement. |
| src/mobius/models/init.py | Exports Gemma4UnifiedModel. |
| src/mobius/components/_common.py | Extends create_attention_bias() to OR-in a same-block overlay from block_sequence_ids. |
| src/mobius/components/_common_test.py | Adds numerical tests validating the blockwise overlay behavior (full/sliding/padding/decode cases). |
| src/mobius/components/_attention.py | Adds is_causal parameter propagation to the ONNX Attention op to support fully baked float masks. |
| src/mobius/_registry.py | Registers gemma4_unified and gemma4_unified_text model types and adds a default model id mapping. |
| src/mobius/_configs/per_model/_gemma4_unified_vision.py | Adds a unified vision hook mapping HF unified vision fields into VisionConfig. |
| src/mobius/_configs/per_model/_gemma4_unified_audio.py | Adds a unified audio hook mapping HF unified audio fields into Gemma4AudioConfig. |
| src/mobius/_configs/per_model/init.py | Ensures the new unified config hooks are imported/registered. |
| src/mobius/_configs/_base.py | Adds use_bidirectional_attention to Gemma4Config and extracts it from HF configs. |
| attention_mask=torch.from_numpy(attention_mask), | ||
| position_ids=torch.from_numpy(position_ids), | ||
| ) | ||
| hf_logits = hf_out.logits.numpy() |
Member
Author
There was a problem hiding this comment.
Fixed in fb0ad6b — converted the HF reference tensors with .detach().cpu().numpy() so the parity check works regardless of CPU/CUDA placement.
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
jambayk
reviewed
Jun 5, 2026
Larger Gemma 4 models (12B/26B/32B) set `use_bidirectional_attention="vision"`: contiguous runs of image/audio placeholder tokens must attend bidirectionally within each block, on BOTH full-attention and sliding-window layers, while text stays causal. Mobius previously ignored this and built a purely causal decoder, producing incorrect attention for multimodal inputs. What changed: - `Gemma4Config` gains a `use_bidirectional_attention` field, extracted from the HF text config in `from_transformers`. - `create_attention_bias` gains an optional `block_sequence_ids` argument that ORs a "blockwise overlay" (same-block, id >= 0) onto the causal (and sliding) mask before the padding AND, exactly matching HF `blockwise_overlay` and the `(causal [AND window]) OR same_block AND padding` composition order. - `_apply_attention` gains an `is_causal` parameter. The bidirectional path bakes the full mask into the float bias and calls the Attention op with `is_causal=0`, since leaving the op's built-in causal mask on would re-mask the future same-block positions and cancel the overlay. This also forces the float-bias path (GroupQueryAttention cannot express the overlay). - The overlay is plumbed embedding -> decoder via a new `block_sequence_ids` tensor (like `per_layer_inputs`): the embedding sub-model computes it from `input_ids` (`_compute_block_sequence_ids`, mirroring HF `get_block_sequence_ids_for_mask`) and the decoder consumes it. The embedding `forward` now returns a dict of named outputs. The text-only path stays causal/GQA-eligible, matching HF (which only injects the overlay in the multimodal wrapper). Tests: - Numerical unit tests for the blockwise bias (multi-block, block wider than the sliding window, padding, single-query decode) and for `_compute_block_sequence_ids` (including adjacent image+audio = one block). - End-to-end wiring tests asserting the package exposes/consumes `block_sequence_ids` and emits `Attention` with `is_causal=0` (no GQA) when the overlay is active. - An integration test comparing mobius's bias against HuggingFace's real `create_causal_mask` / `create_sliding_window_causal_mask` for the actual gemma-4-26b-a4b-it config (config-only, no weights). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Register the gemma4_unified text backbone (model_type gemma4_unified_text) to Gemma4CausalLMModel + Gemma4Config. The gemma-4-12B unified text architecture is the same family as gemma4 (dual head_dim 256/512, attention_k_eq_v with a single global KV head, dual RoPE, final-logit softcapping, vision-block bidirectional attention), so the existing Gemma4CausalLMModel builds it directly. Tests: - L1 graph build (tiny dual-head_dim k_eq_v bidirectional config). - Real float32 prefill parity vs HuggingFace Gemma4UnifiedForConditionalGeneration text path on the 48-layer google/gemma-4-12B checkpoint (CUDA): max_abs_diff=0.01, cos=1.0, argmax match, no NaN. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Implement the full encoder-free multimodal gemma-4-12B model (`Gemma4UnifiedForConditionalGeneration`, model_type `gemma4_unified`), building on the already-landed text backbone and bidirectional vision-block attention fix. Unlike the existing gemma4 multimodal model, gemma-4-12B has no SigLIP vision tower or Conformer audio tower. Vision and audio are lightweight encoder-free embedders: - Vision: raw image patches (P^2*3 = 6912) → LayerNorm → Dense → LayerNorm → + factorized 2D position embeddings → LayerNorm → scale-free RMSNorm → Linear projection to the 3840-d text hidden size. - Audio: raw waveform-frame features (640) → scale-free RMSNorm → Linear projection to 3840. Both embedders strip padding patches/frames inside the ONNX graph via Compress, emitting `[num_valid, text_hidden]` directly consumable by the multimodal-fusion embedding sub-model (matching HF semantics exactly). The model reuses the gemma4 text decoder (dual head_dim, attention_k_eq_v, vision-block bidirectional attention via block_sequence_ids) and the multimodal-fusion embedding sub-model. `Gemma4UnifiedTask` builds a 3- or 4-model package (decoder, vision_encoder, embedding, and audio_encoder when an audio config is present). Config extraction hooks map the unified vision/audio sub-configs onto mobius VisionConfig / Gemma4AudioConfig. preprocess_weights maps the raw safetensors checkpoint names (`vision_embedder.*`, `embed_vision.embedding_projection.*`, `embed_audio.embedding_projection.*`) to the ONNX sub-model parameter names. Testing (real google/gemma-4-12B checkpoint, float32, CUDA): - L1 graph build test for the 4-model package. - Full multimodal prefill parity vs HuggingFace: vision cos_sim=1.000000, decoder last-token cos_sim=1.000000, max_abs_diff=1.10, argmax match, no NaN. The bidirectional reference requires passing mm_token_type_ids to HF (otherwise it falls back to a causal mask). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Address code review: the vision projector pre-norm and projection were sized with output_proj_dims while the forward pass feeds them mm_embed_dim activations (output of pos_norm). These are equal for gemma-4-12B (both 3840), so behavior is unchanged, but the coupling was implicit. Use mm_embed_dim for the projector input dim (matching HF embedding_projection: mm_embed_dim -> text_hidden) and assert output_proj_dims == mm_embed_dim to guard future variants. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Improve CI coverage of the gemma4_unified code (previously exercised only by the GPU integration test). Adds CPU-only unit tests for: - Gemma4UnifiedModel.preprocess_weights: full checkpoint name mapping (language_model -> decoder, tied embed/lm_head, vision_embedder -> vision_encoder with pos_embedding x/y split, embed_vision/embed_audio projection mapping, scale-free pre-projection norms dropped). - The standalone vision and audio embedder preprocess_weights methods. - The vision and audio config extraction hooks (field mapping plus the unrelated-model skip path). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Move inline numpy/onnx_ir/OnnxModelSession imports in the gemma4 block-wise attention tests to the top of the file, per the Google Python style guide. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
The gemma4_unified / gemma4_unified_text registry keys were tested only via dedicated build methods, so the model-coverage gates failed: they had no entry in _test_configs.py (L1/L3) and gemma4_unified_text had no test_model_id (L2). - Add gemma4_unified to VL_CONFIGS and gemma4_unified_text to CAUSAL_LM_CONFIGS so the shared parametrized build suites exercise them (encoder-free 3-model VL split; dual head_dim / k_eq_v text backbone). - Exclude gemma4_unified_text from synthetic parity: it is an internal alias with no AutoModelForCausalLM-registered HF model_type; numeric parity is covered by the real-weight integration test. - Add gemma4_unified_text test_model_id (google/gemma-4-12B) for L2. - Add L4/L5 YAML cases for gemma-4-12B (vision-language + text), gated and skipped in CI, mirroring the rest of the gemma-4 family. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
The gemma4 multimodal split previously computed block_sequence_ids in the embedding sub-model and forwarded it to the decoder as a separate INT64 graph tensor. onnxruntime-genai cannot forward an arbitrary integer tensor between the embedding and decoder sub-models (it can forward input_ids), so the bidirectional vision-block overlay never reached the decoder at runtime. Move the computation into the decoder: Gemma4TextModel.forward now derives block_sequence_ids from input_ids via _compute_block_sequence_ids when the model uses vision-block bidirectional attention and no overlay is supplied. The decoder graph declares input_ids (alongside inputs_embeds); the embedding model no longer emits block_sequence_ids. This computes the overlay exactly once (relocated, not duplicated), unifies the text and multimodal paths, and removes a cross-model tensor that genai cannot plumb. inputs_embeds still takes precedence for token embeddings, so input_ids is used only for the overlay (no embedding recomputation). The genai_config generator introspects graph I/O, so the decoder now lists input_ids and the embedding lists only inputs_embeds automatically. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Map the gemma-4-12B unified (encoder-free) checkpoint to the existing gemma4 ORT GenAI pipelines: model_type "gemma4_unified" -> "gemma4" (multimodal package) and "gemma4_unified_text" -> "gemma4_text" (standalone text backbone). Add both to _GEMMA4_MODEL_TYPES so the gemma4 vision/audio processor config writers apply. Covered by a new unit test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
justinchu/gemma4-bidirectional-attn
branch
from
June 5, 2026 18:56
2d2e629 to
4b723c0
Compare
The gemma4_unified (gemma-4-12B) encoder-free vision embedder's patch_dense projection produces activations whose magnitude (~77000, measured) exceeds the float16 range (max 65504). In float16 the dense output overflows to +inf and the following patch_ln2 LayerNorm emits NaN, which propagates to every output position (and triggers a CUDA illegal-memory-access in genai's topk). HF runs this embedder in bfloat16 (max ~3.4e38), so only float16 is out of range. Make the upcast dtype-aware: for float16 models compute patch_dense + patch_ln2 in float32 (via _F32Linear / _F32LayerNorm) then cast back to float16 after patch_ln2 normalizes into a safe range; bfloat16 and float32 models keep their native dtype (no casts), which also matches HF's bfloat16 behavior more faithfully. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…udio) Rewrite examples/gemma4_unified_ort_genai.py to demonstrate the full unified multimodal package (decoder + embedding + vision_encoder + audio_encoder) running text, image+text, and audio+text generation through onnxruntime-genai on GPU. Image / audio are preprocessed with the HuggingFace AutoProcessor and fed via Generator.set_inputs(NamedTensors), because genai's built-in Gemma4ImageTransform targets the SigLIP gemma4 contract (16px/768-dim), not the encoder-free unified model (48px merged patches/6912-dim). Text uses the native genai path with a manually prepended BOS (base checkpoint, no chat template). Documents the per-layer-KV genai build requirement for the mixed KV cache. Verified end-to-end on H200: image -> "A lake surrounded by mountains and trees."; audio (jfk.flac) -> "Ask not what your country can do for you...". Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Add an optional `--quantize {Q4_K_M,NF4}` flag to the gemma-4-12B
multimodal example that INT4-quantizes the text decoder with Olive's
OnnxKQuantQuantization / OnnxBnb4Quantization pass, then runs generation
against the quantized package.
The decoder holds >95% of the model's weights, so quantizing it alone
shrinks the f16 package ~23GB -> ~6.8GB (3.4x) while the embedding,
vision_encoder, audio_encoder and all tokenizer/processor/genai configs
are copied unchanged for a drop-in ORT GenAI package. Olive bookkeeping
files are removed and the `logits` output name is preserved (defensive
rename handling for Olive versions that emit `logits_Q4`).
Spot-checked on the base checkpoint: coherent text/image/audio output,
~0.986 last-token logit cosine and ~75% greedy top-1 agreement vs f16.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
justinchu/gemma4-bidirectional-attn
branch
from
June 5, 2026 20:16
4b723c0 to
ab8a56c
Compare
…ention _compute_block_sequence_ids previously OR'd in audio_token_id, granting audio placeholder runs block bidirectional attention. That diverges from HuggingFace: HF builds is_vision from mm_token_type_ids as (==1)|(==2) (image or video), while audio is token-type 3 and is deliberately excluded, so audio tokens keep plain causal attention. This was a real bug for gemma-4-12B (gemma4_unified), the only gemma4 model with both vision and audio, where audio_token_id is non-None. Drop the audio branch (and the audio_token_id parameter) so only image runs form blocks; gemma4_unified has no video modality. Update the two unit tests to assert audio is excluded. Reported by @jambayk in review. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
justinchu/gemma4-bidirectional-attn
branch
from
June 5, 2026 22:25
f57b291 to
131faa2
Compare
Address Copilot review feedback on PR #338: - use_bidirectional_attention docstring: correct the "vision" mode to say only image blocks are bidirectional (audio is token-type 3 in HF and excluded), and note that "all" mode is not implemented (treated as causal; only "vision" activates the block overlay). - gemma4_unified integration tests: convert HF reference tensors with .detach().cpu().numpy() so the parity checks work regardless of whether the reference model is placed on CPU or CUDA. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Contributor
Review synthesis (4-reviewer team: readability, code, critical, deep)Verdict: Strong PR, no Critical issues. The deep reviewer verified the bidirectional block-overlay mask math is bit-for-bit faithful to HF's blockwise overlay (grounded on Gemma3's identical algorithm, since Gemma4 isn't in the installed transformers). Findings below are deduplicated; the deep reviewer's grounded verdict wins on math/spec ties (it downgraded the cache-boundary concern from Major to Minor — safe for single-shot prefill where Major
Minor
Open questions for the author (HF-grounded, from the deep reviewer)
NitDuplicate Praise
Synthesized from a 4-model review team (readability, code, critical, deep reviewers). |
Member
Author
|
“Overlay disables GQA for the entire text/decode phase” --> this isn't applicable because the graph is static. |
The gemma-4-12B base checkpoint tends to emit the <end_of_image> /
<end_of_audio> structural tokens (258882 / 258883) during generation,
degenerating into repeated <image|> instead of describing the image.
HF's generation_config handles this via suppress_tokens, but genai has
no native equivalent, so the decode loop now masks those token ids to
-inf before sampling (get_logits -> mask -> set_logits -> sample).
Also switch the default image/audio prompts to completion-style leads
("This image shows" / "The audio says") since this is a base, not
instruction-tuned, model.
Verified on GPU (f16, greedy), matching HF model.generate:
- text -> " Tokyo."
- image -> " the Chinese Arch in the Chinatown of Sydney, Australia."
- audio -> " He hoped there would be stew for dinner, turnips and ..."
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Previously use_bidirectional_attention="all" (HF's fully-bidirectional mode) silently fell through to causal attention, since only "vision" activated the block overlay. Silently producing the wrong attention pattern is a correctness trap. The decoder now raises NotImplementedError for any value other than None or "vision". No supported Gemma4 checkpoint uses "all"; if one ever does, this surfaces the gap loudly instead of generating subtly wrong outputs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
| f"Building unified multimodal package for {model_id!r} " | ||
| f"(dtype={dtype}, ep={ep}) — this downloads ~25GB of weights ..." | ||
| ) | ||
| manifest = auto_export(model_id, output_dir, dtype=dtype, ep=ep) |
| ), | ||
| ) | ||
| fields: dict = {} | ||
| _gemma4_unified_vision(composite, None, "gemma4_unified", fields) |
|
|
||
| composite = SimpleNamespace(model_type="qwen2_vl", vision_config=object()) | ||
| fields: dict = {} | ||
| result = _gemma4_unified_vision(composite, None, "qwen2_vl", fields) |
| audio_token_id=258881, | ||
| audio_config=SimpleNamespace(audio_embed_dim=640), | ||
| ) | ||
| result = _gemma4_unified_audio(composite, None, "gemma4_unified", {}) |
…ors, audio eps) Triaged the latest review. Fixes: - TypeError in integration test: _compute_block_sequence_ids no longer takes audio_token_id; drop the stray kwarg (was @integration-gated so the fast suite never caught it). - auto_export emitted SigLIP-style image/audio processor configs (Gemma4ImageTransform, 128-dim Gemma4LogMel) for the encoder-free gemma4_unified* models, whose real inputs are raw 6912-dim merged patches / 640-dim frames. There is no ort-extensions transform for that contract, so omit image_processor.json / audio_processor.json for unified types (callers use HF processor + Generator.set_inputs, as the example does). Added skip tests. - Vision-block overlay now gates on a valid (non-zero) image_token_id. The gemma4_unified_text backbone carries use_bidirectional_attention= "vision" but has no image token, so it now keeps pure causal (GQA- eligible) attention instead of building an all-(-1) no-op overlay. Also removes the image_token_id-or-0 footgun (token 0 could be mis-marked). - Audio embedder now reads audio_config.rms_norm_eps (falling back to the text eps) so audio/text eps can differ; the hook now actually maps it. - Fixed a build_graph test docstring that claimed the embedding model emits block_sequence_ids (it does not; the decoder derives it), and rewrote the misleading "GQA-eligible is kept" decoder comment to state the real static-graph trade-off honestly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
The two gemma-4-12B (gemma4_unified) cases previously used skip_reason, which skips the golden test everywhere AND blocks golden generation, so no golden files existed. Switch to ci_skip_reason (skips only in CI, runs locally, does not block generation), generate the JSON goldens, and add a third case for audio. All verified locally on H200: - causal-lm/gemma-4-12b L4 PASS (L5 multi-model: framework-skipped) - vision-language/gemma-4-12b L4 + L5 PASS -> "a bobcat in the snow." - speech/gemma-4-12b-audio (new) L4 + L5 PASS -> LibriSpeech transcription Supporting test-infra fixes (general, not gemma-specific): - Golden generation + e2e harness assumed every processor has a usable chat template. Base checkpoints (gemma-4-12B) ship none, so build the multimodal prompt by manually prepending one image/audio placeholder token per media item (processor then expands to soft tokens). Added a shared _build_mm_prompt helper. - L5 generation now honors generation_config.suppress_tokens (via _load_suppress_token_ids/_suppress_logits), matching HF model.generate. The gemma-4-12B base model degenerates into <end_of_image>/<end_of_audio> without it; for models with no suppress_tokens it is a no-op. - Use completion-style prompts and eos_token_id for the base checkpoint so greedy decode stops where HF does (exact token match, no length- mismatch). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
ORT has no bfloat16 Compress kernel, so the gemma4_unified vision and audio embedders failed to load in bf16 packages (the padding-strip Compress on the projected patch/frame features ran in bf16). Wrap both Compress calls with a float32 cast and cast the selected rows back to the source dtype. The f16/bf16 round-trip through f32 is lossless, so the selection is exact for every build dtype. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Replace the hardcoded `_MAX_GQA_HEAD_DIM = 256` module constant in the
GroupQueryAttention rewrite rules with a per-EP capability,
`EpCapabilities.max_gqa_head_dim` (default 256), threaded through
`group_query_attention_rules(max_head_dim=...)`.
Released ORT GQA kernels (Flash / XQA / memory-efficient) support only
head_dim in {64, 128, 256}. Larger head dims (e.g. Gemma4
global-attention layers, head_dim=512) require the unfused
FP32-QK-accumulation fallback added by ORT PR #28198 (fixes #28195),
unreleased as of ORT 1.28.0 — verified to run head_dim=512 GQA correctly
(prefill + decode, f16) in isolation on an ORT 1.28.0 dev build.
The default stays 256 for all EPs (no behavior change) so a GQA node is
only emitted where the running ORT kernel can execute it. The limit is
now an EP/runtime capability decision rather than a magic constant: once
the head_dim>256 support ships in a released ORT, the cuda EP entry can
raise `max_gqa_head_dim` (or gate it on the runtime ORT version) without
touching the rewrite rules. Adds tests covering both the skip
(head_dim > limit) and fuse (head_dim <= limit) paths.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
justinchu/gemma4-bidirectional-attn
branch
from
June 8, 2026 23:58
2328c47 to
389fa47
Compare
ORT PR #28198 (fixes #28195) gives the CUDA GroupQueryAttention kernel an unfused FP32-QK-accumulation fallback that handles any head_dim. Set the CUDA EP's `max_gqa_head_dim=None` (no limit) so all decoder attention layers — including Gemma4's head_dim=512 global-attention layers — fuse to GroupQueryAttention instead of falling back to the standard Attention op. Verified head_dim=512 GQA runs correctly (prefill + decode, f16) in isolation on an ORT 1.28.0 dev build containing #28198. `max_gqa_head_dim` is now `int | None`; other EPs keep the conservative 256 cap until their runtimes ship larger-head-dim support. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
GQA fusion previously skipped Attention nodes whose head_dim exceeded a per-EP cap (default 256), keeping standard ONNX Attention for large-head-dim layers such as Gemma4 global attention (head_dim=512). The CUDA EP already lifted this cap because its GQA kernel handles any head_dim via an FP32-QK unfused fallback. Removing the cap entirely so GQA is emitted uniformly for every decoder Attention node regardless of head_dim. Whether a runtime's GQA kernel supports a given head_dim is an EP concern already gated by gqa_dtypes; the head_dim cap was redundant policy. This deletes: - _DEFAULT_MAX_GQA_HEAD_DIM and _head_dim_exceeds_gqa_limit - the max_head_dim parameter on the rule classes and the rules factory - the EpCapabilities.max_gqa_head_dim field and CUDA override The two cap-behavior tests are replaced by a positive test asserting a head_dim=512 model fuses to GroupQueryAttention. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Jun 12, 2026
…n renames, gemma3 multimodal (#350) ## Summary Ran the full integration + L4 (golden) + L5 (generation) suite on GPU (`MOBIUS_TEST_DEVICE=cuda`, 8× H200) to find regressions and long-standing bugs. Triaged 88 failures by root cause. **PR #338 (gemma4 bidirectional overlay + GQA-cap removal) introduced zero regressions** — every spot-checked failure reproduces identically on the parent commit `f972184`. This PR lands the verified, high-impact fixes. Remaining buckets (deep numeric/decode bugs, ref-API drift, env/external) are tracked separately and will follow. ## Fixes ### 1. Disable TF32 on GPU (`tests/conftest.py`) On Ampere+/Hopper the ORT CUDA EP uses TF32 for fp32 matmuls by default, while the PyTorch reference computes in true fp32. The ~1e-2 logit discrepancy spuriously failed **~35** fp32 numeric-parity tests (rtol/atol 1e-3). Set `NVIDIA_TF32_OVERRIDE=0` in conftest before any CUDA library initializes (uses `setdefault` so users can opt back in). Verified: with the env var unset, `gpt2` and `qwen2.5-0.5b` L4 now pass. ### 2. transformers ≥5.x flattened ViT/CLIP weight names (`src/mobius/models/vit.py`) transformers 5.x flattened the ViT state dict to `layers.N.*` with consolidated `attention.{q,k,v,o}_proj` / `mlp.fc1/fc2` names. The legacy rename map no longer matched, leaving graph initializers unfilled (ORT load failure). Added an **additive** new-naming branch (legacy 5.0–5.9 path preserved) and aligned the in-test torch reference modules with the mobius graph param names. ### 3. gemma3 multimodal vision encoder (`src/mobius/models/gemma3.py`) Two bugs broke the gemma-3 image-text-to-text pipeline: - The full-VLM `preprocess_weights` prefixed vision weights but didn't rename the vision MLP `fc1/fc2` → `up_proj/down_proj`, so FCMLP initializers were never filled (ORT load failure). - The vision encoder returned rank-3 `(batch, tokens, hidden)`, but the embedding sub-model declares `image_features` as rank-2 and gathers along axis 0. Squeeze the batch dim to honor the 2-D contract (matching the `PixtralVLTask` precedent and the ort-genai runtime). With both, **gemma-3-4b-it L4 golden passes** on CUDA. ## Verification - TF32 fix: `gpt2`, `qwen2.5-0.5b` L4 pass with env unset. - ViT/CLIP parity tests pass. - gemma-3-4b-it L4 golden passes; 25 gemma3 build/task unit tests pass. - `lintrunner` clean on all changed files. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ## Architecture-diff notes (review follow-up) Two entries in the `arch-diff-bot` comment (`e005489 -> 2f13c7d`) deserve explanation: ### falcon: `Gelu` -> `Sigmoid + Mul` The bloom-GELU fix changed falcon's MLP from a hardcoded `activation="gelu"` to `activation=config.hidden_act or "gelu"` so `FCMLP` can be shared with Bloom (which needs `gelu_pytorch_tanh`). For a **real** falcon config this is a no-op: `ArchitectureConfig.hidden_act` resolves to `"gelu"` via `config.activation` (HF `FalconConfig` has no `hidden_act` and defaults `activation="gelu"`). The arch-diff only showed SiLU (`Sigmoid+Mul`) because the **synthetic test config** used the generic `_base_config` default `hidden_act="silu"`. Fixed by setting `hidden_act="gelu"` in falcon's synthetic config so it matches real Falcon. ### `RotaryEmbedding: num_heads: 2 -> 4` This is a **false positive** from the diff tool's positional node matching, *caused by* the SiLU regression above -- not a real change. Each falcon layer emits two `RotaryEmbedding` nodes: Q-rotary (`num_heads=4`) and K-rotary (`num_heads=2`, GQA `kv_heads=2`). Inserting `Sigmoid+Mul` per layer shifted all subsequent node indices by +2, so the tool aligned base's K-rotary (`num_heads=2`) against head's Q-rotary (`num_heads=4`). Per-head counts are unchanged at both base and head. Restoring `Gelu` removes the extra nodes, realigns indices, and makes this spurious entry disappear. --------- Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+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.
Summary
Two related changes for the Gemma 4 family:
1. Fix bidirectional vision-block attention (existing gemma4 models)
The existing gemma4 (32B) and gemma4 MoE multimodal models always built a
purely causal decoder, ignoring
use_bidirectional_attention="vision".HuggingFace gemma4 makes contiguous vision-token spans attend
bidirectionally (via a blockwise overlay OR-ed onto the causal/sliding
mask, applied to both full- and sliding-attention layers). This adds:
use_bidirectional_attentiontoGemma4Config+ extraction.block_sequence_idsplumbed embedding → decoder;create_attention_biasgains a
block_sequence_idsblockwise-overlay path (forces the float-bias,non-GQA attention path with
is_causal=0).2. Add gemma-4-12B (
gemma4_unified)The full encoder-free multimodal
Gemma4UnifiedForConditionalGeneration:global 512),
attention_k_eq_vwith a single global KV head, dual RoPE,final-logit softcapping. Reuses
Gemma4CausalLMModel.→ factorized 2D posemb → LN → scale-free RMSNorm → Linear(→3840).
→ Linear(→3840).
Gemma4UnifiedModel+Gemma4UnifiedTaskbuild a 3/4-model package(decoder, vision_encoder, embedding, optional audio_encoder).
Testing
Fast suite: 2783 passed. L1 graph-build tests for text and 4-model
multimodal packages.
Real
google/gemma-4-12Bcheckpoint, float32 on CUDA (H200):decoder last-token cos_sim=1.000000, max_abs_diff=1.10, argmax match,
no NaN. (Bidirectional reference requires passing
mm_token_type_idsto HF.)