fix(gemma4): Restore KV sharing for the Gemma 4 assistants and their targets

src/transformers/models/gemma4/configuration_gemma4.py

@@ -226,9 +226,9 @@ def __post_init__(self, **kwargs):
 
     @property
     def kv_sharing_modes(self) -> list[str | None]:
-        # In this case there are no shared layers
-        if self.num_kv_shared_layers == 0:
-            return [None] * self.num_hidden_layers
+        # In this case all layers are shared
+        if self.num_kv_shared_layers == self.num_hidden_layers:
+            return ["reuse"] * self.num_hidden_layers
 
         # After this index, all layers are consumers
         first_shared_layer = self.num_hidden_layers - self.num_kv_shared_layers

tests/models/gemma4_assistant/test_modeling_gemma4_assistant.py

@@ -14,6 +14,8 @@
 
 import unittest
 
+from parameterized import parameterized
+
 from transformers import is_torch_available
 from transformers.testing_utils import (
     Expectations,
@@ -23,17 +25,66 @@
     torch_device,
 )
 
+from ...test_modeling_common import ids_tensor
 from ...test_processing_common import url_to_local_path
 
 
 if is_torch_available():
     from transformers import (
         AutoModelForCausalLM,
+        Gemma4AssistantConfig,
+        Gemma4AssistantForCausalLM,
+        Gemma4ForCausalLM,
         Gemma4ForConditionalGeneration,
         Gemma4Processor,
+        Gemma4TextConfig,
     )
 
 
+@require_torch
+class Gemma4AssistantTest(unittest.TestCase):
+    @parameterized.expand([0, 2])
+    def test_assisted_decoding_matches_greedy_search(self, num_kv_shared_layers):
+        text_config_kwargs = {
+            "vocab_size": 99,
+            "hidden_size": 32,
+            "intermediate_size": 64,
+            "num_attention_heads": 2,
+            "num_key_value_heads": 1,
+            "head_dim": 16,
+            "global_head_dim": 16,
+            "sliding_window": 8,
+        }
+        model_config = Gemma4TextConfig(
+            num_hidden_layers=4,
+            layer_types=["sliding_attention", "full_attention", "sliding_attention", "full_attention"],
+            num_kv_shared_layers=num_kv_shared_layers,
+            hidden_size_per_layer_input=16,
+            vocab_size_per_layer_input=99,
+            **text_config_kwargs,
+        )
+        assistant_config = Gemma4AssistantConfig(
+            text_config=Gemma4TextConfig(
+                num_hidden_layers=2,
+                layer_types=["sliding_attention", "full_attention"],
+                hidden_size_per_layer_input=0,
+                vocab_size_per_layer_input=0,
+                **text_config_kwargs,
+            ),
+            backbone_hidden_size=32,
+        )
+        model = Gemma4ForCausalLM(model_config).to(torch_device).eval()
+        assistant_model = Gemma4AssistantForCausalLM(assistant_config).to(torch_device).eval()
+        assistant_model.generation_config.num_assistant_tokens = 2
+
+        input_ids = ids_tensor([1, 7], vocab_size=99)
+        output_greedy = model.generate(input_ids, max_new_tokens=5, do_sample=False, eos_token_id=-1)
+        output_assisted = model.generate(
+            input_ids, assistant_model=assistant_model, max_new_tokens=5, do_sample=False, eos_token_id=-1
+        )
+        self.assertListEqual(output_greedy.tolist(), output_assisted.tolist())
+
+
 @slow
 @require_torch
 @unittest.skip(reason="Update after release")  # TODO @vasqu

tests/models/gemma4_unified_assistant/test_modeling_gemma4_unified_assistant.py

@@ -14,6 +14,8 @@
 
 import unittest
 
+from parameterized import parameterized
+
 from transformers import is_torch_available
 from transformers.testing_utils import (
     Expectations,
@@ -23,6 +25,7 @@
     torch_device,
 )
 
+from ...test_modeling_common import ids_tensor
 from ...test_processing_common import url_to_local_path
 
 
@@ -31,9 +34,53 @@
         AutoModelForCausalLM,
         Gemma4ForConditionalGeneration,
         Gemma4Processor,
+        Gemma4UnifiedAssistantConfig,
+        Gemma4UnifiedAssistantForCausalLM,
+        Gemma4UnifiedForCausalLM,
+        Gemma4UnifiedTextConfig,
     )
 
 
+@require_torch
+class Gemma4UnifiedAssistantTest(unittest.TestCase):
+    @parameterized.expand([0, 2])
+    def test_assisted_decoding_matches_greedy_search(self, num_kv_shared_layers):
+        text_config_kwargs = {
+            "vocab_size": 99,
+            "hidden_size": 32,
+            "intermediate_size": 64,
+            "num_attention_heads": 2,
+            "num_key_value_heads": 1,
+            "head_dim": 16,
+            "global_head_dim": 16,
+            "sliding_window": 8,
+        }
+        model_config = Gemma4UnifiedTextConfig(
+            num_hidden_layers=4,
+            layer_types=["sliding_attention", "full_attention", "sliding_attention", "full_attention"],
+            num_kv_shared_layers=num_kv_shared_layers,
+            **text_config_kwargs,
+        )
+        assistant_config = Gemma4UnifiedAssistantConfig(
+            text_config=Gemma4UnifiedTextConfig(
+                num_hidden_layers=2,
+                layer_types=["sliding_attention", "full_attention"],
+                **text_config_kwargs,
+            ),
+            backbone_hidden_size=32,
+        )
+        model = Gemma4UnifiedForCausalLM(model_config).to(torch_device).eval()
+        assistant_model = Gemma4UnifiedAssistantForCausalLM(assistant_config).to(torch_device).eval()
+        assistant_model.generation_config.num_assistant_tokens = 2
+
+        input_ids = ids_tensor([1, 7], vocab_size=99)
+        output_greedy = model.generate(input_ids, max_new_tokens=5, do_sample=False, eos_token_id=-1)
+        output_assisted = model.generate(
+            input_ids, assistant_model=assistant_model, max_new_tokens=5, do_sample=False, eos_token_id=-1
+        )
+        self.assertListEqual(output_greedy.tolist(), output_assisted.tolist())
+
+
 @slow
 @require_torch
 @unittest.skip(reason="Update after release")  # TODO(vasqu)