From c9db053b5f15a4bd16bbebe4da7d19f1f2588b67 Mon Sep 17 00:00:00 2001 From: ydshieh Date: Tue, 10 Nov 2026 16:17:45 +0100 Subject: [PATCH 1/4] Fix RecurrentGemma padded generation by declaring its `layer_types` 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 --- .../configuration_recurrent_gemma.py | 6 ++++++ .../modeling_recurrent_gemma.py | 18 +++++++++++------- .../test_modeling_recurrent_gemma.py | 7 +++++++ 3 files changed, 24 insertions(+), 7 deletions(-) diff --git a/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py index da7a5d8ee0da..7f650ae26df1 100644 --- a/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py +++ b/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py @@ -99,5 +99,11 @@ def validate_architecture(self): def layers_block_type(self): return (self.block_types * 100)[: self.num_hidden_layers] + @property + def layer_types(self): + # Describes the layers for the generic cache / mask utilities (the modeling itself uses `layers_block_type`): + # recurrent blocks keep their own state (no KV cache), attention blocks always use local (sliding window) attention + return ["linear_attention" if x == "recurrent" else "sliding_attention" for x in self.layers_block_type] + __all__ = ["RecurrentGemmaConfig"] diff --git a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py index 5b276a58a06f..9b1ddb35b0f0 100644 --- a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py +++ b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py @@ -652,13 +652,17 @@ def forward( position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) - causal_mask = create_sliding_window_causal_mask( - config=self.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=past_key_values, - position_ids=position_ids, - ) + # It may already have been prepared by e.g. `generate` + if isinstance(attention_mask, dict): + causal_mask = attention_mask["sliding_attention"] + else: + causal_mask = create_sliding_window_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) hidden_states = hidden_states * self.normalizer.type(hidden_states.dtype) diff --git a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py index 055faa135fb2..f67f3e62ecc4 100644 --- a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py +++ b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py @@ -181,6 +181,13 @@ def test_assisted_decoding_matches_greedy_search(self): def test_left_padding_compatibility(self): pass + @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 + @pytest.mark.generate @unittest.skip(reason="Relies on `past_key_values` returned by the model. Not supported with recurrent gemma") def test_assisted_decoding_sample(self): From 6820072cfee41bd12a8c4df68a87bac349478d4c Mon Sep 17 00:00:00 2001 From: ydshieh Date: Tue, 10 Nov 2026 17:22:56 +0100 Subject: [PATCH 2/4] Use `layers_block_type` (aliased as `layer_types`) instead of a second property Co-Authored-By: Claude Opus 5.5 --- .../recurrent_gemma/configuration_recurrent_gemma.py | 12 ++++-------- .../recurrent_gemma/modeling_recurrent_gemma.py | 5 ++++- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py index 7f650ae26df1..6d6b70b1bf33 100644 --- a/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py +++ b/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py @@ -54,7 +54,7 @@ class RecurrentGemmaConfig(PreTrainedConfig): """ model_type = "recurrent_gemma" - attribute_map = {"sliding_window": "attention_window_size"} + attribute_map = {"sliding_window": "attention_window_size", "layer_types": "layers_block_type"} num_hidden_layers: int = 26 vocab_size: int = 256000 @@ -97,13 +97,9 @@ def validate_architecture(self): @property def layers_block_type(self): - return (self.block_types * 100)[: self.num_hidden_layers] - - @property - def layer_types(self): - # Describes the layers for the generic cache / mask utilities (the modeling itself uses `layers_block_type`): - # recurrent blocks keep their own state (no KV cache), attention blocks always use local (sliding window) attention - return ["linear_attention" if x == "recurrent" else "sliding_attention" for x in self.layers_block_type] + # Follow the `layer_types` conventions: the attention blocks always use local (sliding window) attention + block_types = (self.block_types * 100)[: self.num_hidden_layers] + return ["linear_attention" if x == "recurrent" else "sliding_attention" for x in block_types] __all__ = ["RecurrentGemmaConfig"] diff --git a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py index 9b1ddb35b0f0..5d7196b81f5f 100644 --- a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py +++ b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py @@ -465,7 +465,10 @@ def _setup_cache(self, batch, device, dtype): self.conv1d_state = torch.zeros((batch, self.hidden_size, self.conv1d_width - 1), device=device, dtype=dtype) -TEMPORAL_BLOCK_CLASSES = {"recurrent": RecurrentGemmaRecurrentBlock, "attention": RecurrentGemmaAttention} +TEMPORAL_BLOCK_CLASSES = { + "linear_attention": RecurrentGemmaRecurrentBlock, + "sliding_attention": RecurrentGemmaAttention, +} class RecurrentGemmaMlp(nn.Module): From c65546389ca01945095b89f5d67fb572b7193451 Mon Sep 17 00:00:00 2001 From: ydshieh Date: Tue, 10 Nov 2026 19:28:15 +0100 Subject: [PATCH 3/4] Update RecurrentGemma integration test expected values 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 --- .../models/recurrent_gemma/test_modeling_recurrent_gemma.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py index f67f3e62ecc4..d9e4c98a2f58 100644 --- a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py +++ b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py @@ -245,7 +245,7 @@ class RecurrentGemmaIntegrationTest(MemoryCleanupMixin, unittest.TestCase): model_id = "google/recurrentgemma-2b" def test_2b_generate(self): - EXPECTED_TEXTS = ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am looking for some information on the topic. I am looking for some information on the impact of the internet on the society. I am looking for some information on the impact of the internet on the society. I am looking for some', 'Hi today is a new app that allows you to make money by watching videos.\n\nThe app is very simple to use and you can earn money by watching videos.\n\nThe app is available for both Android and iOS devices and you can download it from the Google Play Store or the App Store.\n\nOnce you have downloaded the app'] # fmt: skip + EXPECTED_TEXTS = ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am looking for some information on the topic. I am looking for some information on the impact of the internet on the society. I am looking for some information on the impact of the internet on the society. I am looking for some', "Hi today the 19th of June 2019, I was in the office of the Ministry of Health and Social Protection of the Republic of Uzbekistan. I was there to discuss the issue of the implementation of the project “Strengthening the capacity of the Ministry of Health and Social Protection of the Republic of Uzbekistan to"] # fmt: skip model = AutoModelForCausalLM.from_pretrained( self.model_id, ).to(torch_device) @@ -304,7 +304,8 @@ def test_model_2b_8bit(self): EXPECTED_TEXTS = Expectations( { ("xpu", None): ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am stuck', "Hi today I'm going to show you how to make a simple and easy to make a 3D"], - (None, None): ['Hello I am doing a project on the topic of "The impact of social media on the society" and I am looking', "Hi today I'm going to show you how to make a simple and easy to make a 3D"], + ("rocm", (9, 4)): ['Hello I am doing a project on the 1960s and I am looking for a picture of a ', 'Hi today I am going to show you how to make a simple and easy to make a DIY'], + (None, None): ['Hello I am doing a project on the effects of the environment on the human body. I am looking for a few good', "Hi today I'm going to show you how to make a simple and easy to make a 3D"], } ) # fmt: on From ba5a0d9c50b288c2b208e7a57531706c6f5fe4f3 Mon Sep 17 00:00:00 2001 From: ydshieh Date: Tue, 10 Nov 2026 19:38:23 +0100 Subject: [PATCH 4/4] Update RecurrentGemma left-padding expected values Captured on an A10G with this PR; fp32 and fp16 give the same output. Co-Authored-By: Claude Opus 5.5 --- tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py index d9e4c98a2f58..a77c7994400a 100644 --- a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py +++ b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py @@ -261,7 +261,7 @@ def test_2b_generate(self): self.assertEqual(output_text, EXPECTED_TEXTS) tokenizer.padding_side = "left" - EXPECTED_TEXTS = ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am looking for some information on the topic. I am looking for some information on the impact of the internet on the society. I am looking for some information on the impact of the internet on the society. I am looking for some', 'Hi today I’m going to show you how to make a simple and easy to make a DIY DIY DIY DIY DIY DIY DIY DIY DIY DIY DIY DIY'] # fmt: skip + EXPECTED_TEXTS = ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am looking for some information on the topic. I am looking for some information on the impact of the internet on the society. I am looking for some information on the impact of the internet on the society. I am looking for some', "Hi today I'm going to show you how to make a simple and easy to make a 3D model of a house.\n\nYou will need:\n\n* A 3D printer\n* A 3D model of a house\n* A 3D printer filament\n* A 3D printer pen\n"] # fmt: skip inputs = tokenizer(self.input_text, return_tensors="pt", padding=True).to(torch_device) output = model.generate(**inputs, max_new_tokens=64, do_sample=False)