Repository navigation
🚨 [gemma 3/4] Fix bidirectional attention masking crossing sliding window boundaries - #46850
Merged
zucchini-nlp merged 7 commits intoJun 30, 2026
Conversation
…ion tokens For Gemma 4, vision tokens should only attend bidirectionally within local (sliding window) layers, not global (full attention) layers. Previously, this implementation shared a single create_masks_for_vision_model function (defined in modular_gemma3.py) between both Gemma 3 and Gemma 4. This was incorrect because: - Gemma 3 uses a single unified mask with bidirectional attention baked in for ALL layers (including global). - Gemma 4 explicitly disables bidirectional attention on global layers. This commit: - Keeps Gemma 3's create_masks_for_vision_model unchanged (OR(causal, blockwise) on global layers). - Adds a Gemma 4-specific create_masks_for_vision_model in modular_gemma4.py that uses causal-only masking for global layers.
douglas-reid
force-pushed
the
fix/attention-masking-sliding-window-bidirectional
branch
from
June 25, 2026 19:50
01689c6 to
21a175b
Compare
zucchini-nlp
approved these changes
Jun 29, 2026
Member
There was a problem hiding this comment.
Oke, so we have different masks for Gemma3 and Gemma4 at the end. I think this is change the expected outputs in slow CI, let me trigger it. Some fast tests on assisted decoding are failing btw
Also lets mark the PR with 🚨 , as it is breaking even if in a good way
Member
|
run-slow: gemma3, gemma4, gemma4_unified |
Contributor
|
This comment contains models: ["models/gemma3", "models/gemma4", "models/gemma4_unified"] |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
douglas-reid
changed the title
Manually pad `block_sequence_ids` to `kv_length` before it is passed as a closure to `or_mask_function` in `create_causal_mask`. This ensures the shape matches the full cache length during autoregressive decoding, fixing a device-side assert crash during generation. Additionally, moved local `create_masks_for_vision_model` test imports to the top of the file as requested in PR review.
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: gemma3, gemma4, gemma4_unified |
Contributor
Author
|
Applied some test updates and added 🚨 to the PR title. |
Member
|
run-slow: gemma3, gemma4, gemma4_unified |
Contributor
|
This comment contains models: ["models/gemma3", "models/gemma4", "models/gemma4_unified"] |
1 task done
stevhliu
pushed a commit
to stevhliu/transformers
that referenced
this pull request
Jul 30, 2026
…ndow boundaries (huggingface#46850) * Fix bidirectional attention masking crossing sliding window boundaries in Gemma 3/4 * run make fix-repo * remove unused vars in test * update unified model as well * fix(gemma4): disable bidirectional attention on global layers for vision tokens For Gemma 4, vision tokens should only attend bidirectionally within local (sliding window) layers, not global (full attention) layers. Previously, this implementation shared a single create_masks_for_vision_model function (defined in modular_gemma3.py) between both Gemma 3 and Gemma 4. This was incorrect because: - Gemma 3 uses a single unified mask with bidirectional attention baked in for ALL layers (including global). - Gemma 4 explicitly disables bidirectional attention on global layers. This commit: - Keeps Gemma 3's create_masks_for_vision_model unchanged (OR(causal, blockwise) on global layers). - Adds a Gemma 4-specific create_masks_for_vision_model in modular_gemma4.py that uses causal-only masking for global layers. * Fix CUDA assert during generation in Gemma 3/4 sliding window attention Manually pad `block_sequence_ids` to `kv_length` before it is passed as a closure to `or_mask_function` in `create_causal_mask`. This ensures the shape matches the full cache length during autoregressive decoding, fixing a device-side assert crash during generation. Additionally, moved local `create_masks_for_vision_model` test imports to the top of the file as requested in PR review. * Fix misleading test comments: correct query/key direction in mask assertions --------- Co-authored-by: Douglas Reid <21148125+douglas-reid@users.noreply.github.com>
3 tasks done
1 of 9 tasks
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.
What does this PR do?
This fixes the behavior of attention masking for image tokens to properly respect sliding window boundaries for Gemma 3/4 models. This was needed to match the intended model behavior for local layers.
Before submitting
Pull Request checks?
to it if that's the case.
CI recap
Dashboard: View test results in Grafana
Latest run: 28405951213:2
Result:
success| Jobs: 7 | Tests: 2,337 | Failures: 0 | Duration: 32m 59s