Fix sliding-window mask `layer_idx` in Gemma3/Gemma4 `create_masks_for_vision_model` by jiqing-feng · Pull Request #48482 · huggingface/transformers

jiqing-feng · 2026-09-02T08:56:45Z

Overview

create_masks_for_vision_model() pads block_sequence_ids with layer_idx=0, but builds the sliding mask via create_causal_mask() without passing layer_idx, which defaults to the first full_attention layer. On real checkpoints layer 0 is sliding_attention, so padding and mask are sized against different layers and blockwise_overlay indexes out of bounds.

The two lengths only diverge when max_cache_len > sliding_window; below that the bug is invisible.

Introduced by #46850.

Reproduction

google/shieldgemma-2-4b-it, StaticCache(max_cache_len=sliding_window + 256), one 896x896 image. Before the fix, on A100:

prompt_len: 419  sliding_window: 1024
kv @ layer 0 (sliding) = 1024
kv @ layer 5 (full)    = 1280

IndexKernel.cu:111: Assertion `-sizes[i] <= index && index < sizes[i]` failed. (x128)
torch.AcceleratorError: CUDA error: device-side assert triggered

After the fix: exit 0, generates 'No'. (On XPU the same run aborts with SIGABRT instead of raising — same root cause, different surfacing.)

CPU-only, no checkpoint, the bug has two faces depending on layer order:

`layer_types`	before	after
`["sliding_attention", "full_attention"]`	`IndexError: index 13 is out of bounds for dimension 1 with size 13`	ok
`["full_attention", "sliding_attention"]`	sliding mask `kv_length=32`, cache wants `13`, no exception	ok

The second face — a silently wrong-shaped mask — is the more dangerous one.

Fix

Pin both the padding and the sliding mask to a sliding layer:

if getattr(past_key_values, "is_sliding", None) and True in past_key_values.is_sliding:
    sliding_layer_idx = past_key_values.is_sliding.index(True)
else:
    sliding_layer_idx = 0

getattr(..., None) rather than hasattr: past_key_values=None and an empty is_sliding both fall back to 0, preserving current behaviour for models without sliding layers.

Affected models

gemma3 — gated only on token_type_ids is not None.
gemma4 — gated on use_bidirectional_attention == "vision", set for the 31B and 26B-A4B variants, which use layer_types = ["sliding_attention"] * 5 + ["full_attention"] and sliding_window=1024, i.e. the same trigger conditions. E2B/E4B have use_bidirectional_attention=None and never reach this path.
gemma4_unified — inherits gemma4's function.

Reproduced directly on both gemma3 and gemma4.

Tests

test_attention_mask_composition_with_hybrid_cache, added to Gemma3Vision2TextModelTest and Gemma4Vision2TextModelTest, parameterized over three scenarios. The assertions are just:

self.assertEqual(mask_dict["sliding_attention"].shape[-1], sliding_kv_length)
self.assertEqual(mask_dict["full_attention"].shape[-1], full_kv_length)

Everything before them exists to reach a state where sliding_kv_length != full_kv_length, since the bug is invisible otherwise. The assertNotEqual guard makes a future config change that turns those assertions into tautologies fail loudly.

scenario	q_len	sliding_kv	full_kv	before the fix
`prefill_static`	13	13	32	`IndexError: index 13 out of bounds ... size 13`
`decode_static`	1	4	32	`IndexError: index 14 out of bounds ... size 14`
`decode_dynamic`	1	4	14	shape assertion fails, no exception

Identical numbers on gemma3 and gemma4. Before the fix: 3 failed, 1 passed. After: 4 passed.

Decode is reachable through generate(): prepare_inputs_for_generation clears token_type_ids only after super() returns, and the mask is already built inside that super call. Tracing generate(max_new_tokens=8) with a StaticCache shows two calls into create_masks_for_vision_model, q_len=419 and q_len=1. decode_dynamic is only reachable through a manual model(...) call, but it is the only scenario covering the silent wrong-shape face.

Existing CI missed this because test_attention_mask_composition passes past_key_values=None (no cache, no divergence), and the gemma3 tester config happens to put full_attention at layer 0.

Regressions

tests/models/gemma3 + gemma4 + gemma4_unified modeling tests on A100: baseline 43 failed / 1101 passed, this PR 40 failed / 1107 passed. No regressions. The 40 failures are pre-existing on main. The 3-test delta is flaky, not fixed here — on the unpatched tree they pass 3/3 in isolation, and the failure is a 1.5e-5 mismatch against atol=1e-5.

ruff check, ruff format --check and utils/check_modular_conversion.py all pass.

Files changed

src/transformers/models/gemma3/modular_gemma3.py                    <- hand-edited
src/transformers/models/gemma4/modular_gemma4.py                    <- hand-edited
src/transformers/models/gemma3/modeling_gemma3.py                   <- generated
src/transformers/models/gemma4/modeling_gemma4.py                   <- generated
src/transformers/models/gemma4_unified/modeling_gemma4_unified.py   <- generated
tests/models/gemma3/test_modeling_gemma3.py
tests/models/gemma4/test_modeling_gemma4.py

…odel `create_masks_for_vision_model` pads `block_sequence_ids` using the KV sizes of cache layer 0, but then calls `create_causal_mask` without `layer_idx`, which defaults to the first `full_attention` layer. On Gemma3/Gemma4 checkpoints layer 0 is a `sliding_attention` layer, so with a hybrid cache the two lengths differ: * DynamicCache: the sliding mask is built with the full layer's kv_length, so it no longer matches the sliding KV cache. * StaticCache: the full layer's kv_length is the padded `max_cache_len`, so `blockwise_overlay` indexes `block_sequence_ids` out of bounds. This raises an IndexError on CPU and triggers a device-side assert (crashing the process) on CUDA and XPU. Pin both the padding and the mask itself to the same `sliding_attention` layer. Also add a regression test covering a decode step with both a dynamic and a static hybrid cache. The existing `test_attention_mask_composition` only checks prefill with `past_key_values=None`, and the tester config places `full_attention` at layer 0, so neither exercised this path.

The crash actually happens on the first (prefill) forward: Gemma3's `prepare_inputs_for_generation` clears `token_type_ids` after the first step, so decode never reaches `create_masks_for_vision_model`. With an empty StaticCache the sliding layer reports `kv_length == q_len` while the full layer reports the padded `max_cache_len`, which is exactly the divergence that made `blockwise_overlay` read out of bounds. Rename the test accordingly and parameterize it over prefill/static, decode/static and decode/dynamic.

The Gemma4 checkpoints that reach create_masks_for_vision_model (use_bidirectional_attention="vision": the 31B and 26B-A4B variants) use layer_types starting with sliding_attention and sliding_window=1024, so they hit exactly the same layer_idx mismatch as Gemma3. Mirror the gemma3 test on Gemma4Vision2TextModelTest. Verified: 3 of the 4 parameterizations fail before the fix and all pass after.

zucchini-nlp · 2026-09-02T12:37:53Z

run-slow: gemma3, gemma4, gemma4_unified

github-actions · 2026-09-02T12:39:19Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/gemma3", "models/gemma4", "models/gemma4_unified"]
quantizations: []

HuggingFaceDocBuilderDev · 2026-09-02T12:49:27Z

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.

github-actions · 2026-09-02T13:51:24Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	e773d31d	workflow commit (merge commit)
PR	a5efb462	branch commit (from PR)
main	e15d4674	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

zucchini-nlp · 2026-09-03T09:10:17Z

+    @parameterized.expand([("prefill_static",), ("decode_static",), ("decode_dynamic",)])
+    def test_attention_mask_composition_with_hybrid_cache(self, scenario: str):
+        """The sliding mask must be sized against a sliding layer, not the first full attention one.
+
+        Regression test: on real Gemma3 checkpoints layer 0 is a `sliding_attention` layer, and once the
+        cache outgrows the sliding window its `kv_length` differs from the first `full_attention` layer's.
+        Sizing the sliding mask against the wrong layer yields a mask that mismatches the cache or indexes
+        `block_sequence_ids` out of bounds.
+        """


so iiuc bug was found with a real ShieldGemma ckpt when we feed a very long input, and it throws a cuda index error

Why not just create a tiny test in shield gemma and call model with a long input, or add a slow test_generation_beyond_sliding_window for this failing ckpt? So we don't have to manually engineer inputs and rely on not-so-reliable create_masks_for_vision_model - utilities tend to change over time

Fixed. We use model level tests now.

…the mask utility

github-actions · 2026-09-04T02:55:54Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: gemma3, gemma4, gemma4_unified

github-actions · 2026-09-04T03:05:00Z

CI recap

Dashboard: View test results in Grafana
Latest run: 33631007448:1
Result: success | Jobs: 1 | Tests: 756 | Failures: 0 | Duration: 1h 38m

zucchini-nlp

Nice!

zucchini-nlp

meant to approve

…r_vision_model` (huggingface#48482) * Fix sliding mask layer_idx in Gemma3/Gemma4 create_masks_for_vision_model `create_masks_for_vision_model` pads `block_sequence_ids` using the KV sizes of cache layer 0, but then calls `create_causal_mask` without `layer_idx`, which defaults to the first `full_attention` layer. On Gemma3/Gemma4 checkpoints layer 0 is a `sliding_attention` layer, so with a hybrid cache the two lengths differ: * DynamicCache: the sliding mask is built with the full layer's kv_length, so it no longer matches the sliding KV cache. * StaticCache: the full layer's kv_length is the padded `max_cache_len`, so `blockwise_overlay` indexes `block_sequence_ids` out of bounds. This raises an IndexError on CPU and triggers a device-side assert (crashing the process) on CUDA and XPU. Pin both the padding and the mask itself to the same `sliding_attention` layer. Also add a regression test covering a decode step with both a dynamic and a static hybrid cache. The existing `test_attention_mask_composition` only checks prefill with `past_key_values=None`, and the tester config places `full_attention` at layer 0, so neither exercised this path. * Cover the prefill path in the vision sliding mask regression test The crash actually happens on the first (prefill) forward: Gemma3's `prepare_inputs_for_generation` clears `token_type_ids` after the first step, so decode never reaches `create_masks_for_vision_model`. With an empty StaticCache the sliding layer reports `kv_length == q_len` while the full layer reports the padded `max_cache_len`, which is exactly the divergence that made `blockwise_overlay` read out of bounds. Rename the test accordingly and parameterize it over prefill/static, decode/static and decode/dynamic. * Cover the vision sliding mask regression on Gemma4 as well The Gemma4 checkpoints that reach create_masks_for_vision_model (use_bidirectional_attention="vision": the 31B and 26B-A4B variants) use layer_types starting with sliding_attention and sliding_window=1024, so they hit exactly the same layer_idx mismatch as Gemma3. Mirror the gemma3 test on Gemma4Vision2TextModelTest. Verified: 3 of the 4 parameterizations fail before the fix and all pass after. * Test the vision sliding mask regression through the model instead of the mask utility

jiqing-feng added 3 commits September 2, 2026 06:28

jiqing-feng marked this pull request as ready for review September 2, 2026 08:56

Merge branch 'main' into fix-gemma3-vision-sliding-mask-layer-idx

a5efb46

github-actions Bot requested review from ydshieh and zucchini-nlp September 2, 2026 08:57

zucchini-nlp reviewed Sep 3, 2026

View reviewed changes

Test the vision sliding mask regression through the model instead of …

e9e0734

…the mask utility

zucchini-nlp reviewed Sep 4, 2026

View reviewed changes

zucchini-nlp approved these changes Sep 4, 2026

View reviewed changes

zucchini-nlp enabled auto-merge September 4, 2026 08:04

zucchini-nlp added this pull request to the merge queue Sep 4, 2026

Merged via the queue into huggingface:main with commit c676202 Sep 4, 2026
51 checks passed

Navigation Menu

Repository navigation

Fix sliding-window mask `layer_idx` in Gemma3/Gemma4 `create_masks_for_vision_model` - #48482

Fix sliding-window mask `layer_idx` in Gemma3/Gemma4 `create_masks_for_vision_model`#48482
zucchini-nlp merged 5 commits into
huggingface:mainfrom
jiqing-feng:fix-gemma3-vision-sliding-mask-layer-idx

jiqing-feng commented Sep 2, 2026 •

edited by github-actions Bot

Loading

Uh oh!

zucchini-nlp commented Sep 2, 2026

Uh oh!

github-actions Bot commented Sep 2, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Sep 2, 2026

Uh oh!

github-actions Bot commented Sep 2, 2026

Uh oh!

Uh oh!

zucchini-nlp Sep 3, 2026

Uh oh!

jiqing-feng Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

zucchini-nlp left a comment

Uh oh!

zucchini-nlp left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

jiqing-feng commented Sep 2, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Reproduction

Fix

Affected models

Tests

Regressions

Files changed

Uh oh!

zucchini-nlp commented Sep 2, 2026

Uh oh!

github-actions Bot commented Sep 2, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Sep 2, 2026

Uh oh!

github-actions Bot commented Sep 2, 2026

CI Results

Commit Info

Uh oh!

Uh oh!

zucchini-nlp Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

jiqing-feng Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

CI recap

Uh oh!

zucchini-nlp left a comment

Choose a reason for hiding this comment

Uh oh!

zucchini-nlp left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jiqing-feng commented Sep 2, 2026 •

edited by github-actions Bot

Loading