Fix RecurrentGemma padded generation by declaring its `layer_types` by ydshieh · Pull Request #50577 · huggingface/transformers

ydshieh · 2026-11-10T15:18:08Z

Alternative to #50554, following @vasqu's review: fix the layer types instead of patching the call sites.

Problem

RecurrentGemmaConfig has no layer_types, so every layer is inferred as sliding_attention (sliding_window is set). The cache then gives the recurrent blocks KV slots too, and the mask / get_seq_length() read layer 0's slot, which belongs to a recurrent block and is always empty. With padding, the decode mask gets the wrong size: on GPU, CUDA SDPA raises (*bias): last dimension must be contiguous. On CPU, generation runs but the outputs are silently wrong. Regression from #48961.

Fix

Config: layers_block_type now returns standard layer types (recurrent → linear_attention, attention → sliding_attention) and is aliased as layer_types via attribute_map. Nothing changes in config.json, and the attention computation is unchanged (it reads config.sliding_window directly). Only the cache slots of the recurrent blocks change type, and those slots are never used: the recurrent block keeps its state on the module.
Modeling: accept the mask dict that generate pre-builds once layer_types exists (needed by the static-cache / compile path).
Tests: skip test_recurrent_layers_mask_padding_on_continued_forward. It now runs because the config declares a recurrent layer type, but the recurrent block has never supported a cached multi-token continuation (same skip as Mamba / Jamba).

Same generated tokens and scores as #50554 (tiny model, padded batch).

Results (A10G, torch 2.15)

Test	`main`	This PR
`test_generate_with_and_without_position_ids`	200/200 failed	200/200 passed
`test_generate_from_random_inputs_embeds`	29/1000 failed	1000/1000 passed
`tests/models/recurrent_gemma`	1 failed, 129 passed	130 passed

The slow integration tests' expected values still need updating (same as #50554).

🤖 Generated with Claude Code

Without `layer_types`, every layer is inferred as `sliding_attention`, so the cache allocates KV slots for the recurrent blocks too. The mask / `get_seq_length` then read the (always empty) slot of layer 0, a recurrent block, which gives a wrong decode mask with padding (CUDA SDPA error or silently wrong outputs). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ydshieh · 2026-11-10T15:19:51Z

+    @unittest.skip(
+        "RecurrentGemma's recurrent block has no chunked-continuation support: a cached multi-token forward is "
+        "treated as a new prefill, so the split-vs-single comparison cannot match regardless of padding masking."
+    )
+    def test_recurrent_layers_mask_padding_on_continued_forward(self):
+        pass
+


it was skipped before. I don't want to handle this in the current PR.

yea its expected tbh

mangeshraut712

Declaring layer_types so the generic generate/cache path can prepare a mask dict is the right fix, and reusing a mask that generate already built avoids building a second one. The fallback always reads the sliding_attention entry, which does not match a mixed or recurrent-only stack.

mangeshraut712 · 2026-11-10T15:29:18Z

-        )
+        # It may already have been prepared by e.g. `generate`
+        if isinstance(attention_mask, dict):
+            causal_mask = attention_mask["sliding_attention"]


layer_types maps recurrent blocks to "linear_attention" and attention blocks to "sliding_attention". When generate has already prepared attention_mask as a dict, this line always takes attention_mask["sliding_attention"] and uses that one tensor for every layer. A RecurrentGemma whose block_types are all "recurrent" (or a generate path that only filled linear_attention) will KeyError here. Even on the usual mixed stack, the linear/recurrent layers are then masked with the sliding-window tensor rather than the linear_attention mask that was prepared for them.

If recurrent blocks ignore the mask and only the attention blocks need the sliding window, please say so in the comment — but I'd still use .get("sliding_attention") and fall back to create_sliding_window_causal_mask(...) when that key is missing. If they do consume the mask, the dict should be passed through (or indexed per layer_types[i]) so the linear-attention entry is not dropped. The new skip on test_recurrent_layers_mask_padding_on_continued_forward is fair given the missing chunked-continuation support, but it means padded continued generation stays untested after this fix.

the model is so old, that it would require too much effort to fix it up properly imo

some of your initial points are more theoretical issues which are not really important

HuggingFaceDocBuilderDev · 2026-11-10T15:30:59Z

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.

vasqu · 2026-11-10T15:55:49Z

@mangeshraut712 not sure if it's fully agent comments but since I'm seeing you on multiple PRs now, could you keep the comments shorter (and human). Atm it feels very "agent just spammed" but I dont want to judge preemptively

Reviewing is hard enough but agents love to bloat comments to the end which makes it hard to extract the core messages at times

vasqu

Some comments but my biggest point is to avoid introducing yet another layer types property 👀

vasqu · 2026-11-10T15:56:29Z

        return (self.block_types * 100)[: self.num_hidden_layers]

+    @property
+    def layer_types(self):


hmm why not changed it directly in layers block types instead? Id like to avoid having so many layer types living side by side

vasqu · 2026-11-10T15:57:15Z

-        )
+        # It may already have been prepared by e.g. `generate`
+        if isinstance(attention_mask, dict):
+            causal_mask = attention_mask["sliding_attention"]


the model is so old, that it would require too much effort to fix it up properly imo

some of your initial points are more theoretical issues which are not really important

vasqu · 2026-11-10T15:57:23Z

+    @unittest.skip(
+        "RecurrentGemma's recurrent block has no chunked-continuation support: a cached multi-token forward is "
+        "treated as a new prefill, so the split-vs-single comparison cannot match regardless of padding masking."
+    )
+    def test_recurrent_layers_mask_padding_on_continued_forward(self):
+        pass
+


yea its expected tbh

…d property Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

vasqu

🚨 just to be sure, the model is not worth to fix on a deeper scale imo so good to go

ydshieh · 2026-11-10T16:59:20Z

run-slow: recurrent_gemma

github-actions · 2026-11-10T17:14:05Z

CI (Nvidia + AMD)

Nvidia · ❌ 2 tests with changed errors 😭 · ⚙️ Workflow · 📊 Dashboard

Failures (4 pre-existing)

Pre-existing failures (still failing) — 4 total:

Error changed (2):

        - recurrent_gemma:
            tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py::RecurrentGemmaIntegrationTest::test_2b_generate (❌ ⟹ ❌)
            tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py::RecurrentGemmaIntegrationTest::test_model_2b_8bit (❌ ⟹ ❌)

Same error (2):

        - recurrent_gemma:
            tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py::RecurrentGemmaIntegrationTest::test_2b_sample
            tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py::RecurrentGemmaIntegrationTest::test_longer_than_window

AMD · ❌ 2 tests with changed errors 😭 · ⚙️ Workflow · 📊 Dashboard

Failures (4 pre-existing)

Pre-existing failures (still failing) — 4 total:

Error changed (2):

        - recurrent_gemma:
            tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py::RecurrentGemmaIntegrationTest::test_2b_generate (❌ ⟹ ❌)
            tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py::RecurrentGemmaIntegrationTest::test_model_2b_8bit (❌ ⟹ ❌)

Same error (2):

        - recurrent_gemma:
            tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py::RecurrentGemmaIntegrationTest::test_2b_sample
            tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py::RecurrentGemmaIntegrationTest::test_longer_than_window

Values from the PR's CI run (Nvidia `captured_info.txt`, AMD MI300 failure log). They match the outputs of the daily CI before #48961. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

github-actions · 2026-11-10T18:30:03Z

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

run-slow: recurrent_gemma

Captured on an A10G with this PR; fp32 and fp16 give the same output. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ydshieh · 2026-11-10T18:39:47Z

The 2 remaining failing tests (test_2b_sample, test_longer_than_window) fail with the same error as on main. They're not touched in this PR.

ydshieh commented Nov 10, 2026

View reviewed changes

mangeshraut712 reviewed Nov 10, 2026

View reviewed changes

vasqu reviewed Nov 10, 2026

View reviewed changes

Use layers_block_type (aliased as layer_types) instead of a secon…

6820072

…d property Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ydshieh requested a review from vasqu November 10, 2026 16:32

vasqu approved these changes Nov 10, 2026

View reviewed changes

Update RecurrentGemma integration test expected values

c655463

Values from the PR's CI run (Nvidia `captured_info.txt`, AMD MI300 failure log). They match the outputs of the daily CI before #48961. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Update RecurrentGemma left-padding expected values

ba5a0d9

Captured on an A10G with this PR; fp32 and fp16 give the same output. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ydshieh merged commit 07cc8e8 into main Nov 10, 2026
39 of 40 checks passed

ydshieh deleted the fix_recurrent_gemma_layer_types branch November 10, 2026 18:44

Conversation

ydshieh commented Nov 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Fix

Results (A10G, torch 2.15)

Uh oh!

ydshieh Nov 10, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Nov 10, 2026

Choose a reason for hiding this comment

Uh oh!

mangeshraut712 left a comment

Choose a reason for hiding this comment

Uh oh!

mangeshraut712 Nov 10, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Nov 10, 2026

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Nov 10, 2026

Uh oh!

vasqu commented Nov 10, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

vasqu Nov 10, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Nov 10, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Nov 10, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

ydshieh commented Nov 10, 2026

Uh oh!

github-actions Bot commented Nov 10, 2026

CI (Nvidia + AMD)

Uh oh!

github-actions Bot commented Nov 10, 2026

Uh oh!

ydshieh commented Nov 10, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ydshieh commented Nov 10, 2026 •

edited

Loading