From fc97e6e9cf8b9b02999cfd7f10f58d65e17d4f97 Mon Sep 17 00:00:00 2001
From: Tarek Ziade <tarek@ziade.org>
Date: Tue, 27 Oct 2026 10:37:14 +0100
Subject: [PATCH] [generate] Slice cached decoder inputs of encoder-decoder
 models without decoder mask

When `generate` restarts from an existing cache, `_prefill` only slices the already-cached inputs when it knows the
mask length. Encoder-decoder models usually have no `decoder_attention_mask` (Whisper even pops it), and
`_mask_length` is only recorded for decoder-only models, so the FULL `decoder_input_ids` were processed again on
top of the cache.

This hits encoder-decoder assistants in assisted decoding: the assistant cache grows by the full sequence at every
drafting round, until Whisper's decoder indexes past its 448 learned positions (device-side assert on CUDA, which
then poisons every following test of the worker).

Fixes #50011

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/transformers/generation/utils.py | 13 ++++++----
 tests/generation/test_utils.py       | 36 ++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
index aa0570152ac6..462e2526365e 100644
--- a/src/transformers/generation/utils.py
+++ b/src/transformers/generation/utils.py
@@ -4291,11 +4291,14 @@ def _prefill(
             else:
                 attention_mask_key = "decoder_attention_mask" if self.config.is_encoder_decoder else "attention_mask"
                 attention_mask = model_kwargs.get(attention_mask_key)
-                mask_length = (
-                    attention_mask.shape[1]
-                    if attention_mask is not None
-                    else getattr(generation_config, "_mask_length", -1)
-                )
+                if attention_mask is not None:
+                    mask_length = attention_mask.shape[1]
+                elif self.config.is_encoder_decoder:
+                    # The decoder mask is usually absent (e.g. assistant models): the decoder inputs are the FULL sequence
+                    # if they are longer than the cache
+                    mask_length = input_ids.shape[1] if input_ids.shape[1] > past_length else -1
+                else:
+                    mask_length = getattr(generation_config, "_mask_length", -1)
                 # In this case we need to slice - if it's smaller than the mask, only the new inputs were passed -> no need to do anything
                 if mask_length == input_ids.shape[1]:
                     # inputs will be sliced as `input_ids[:, -next_sequence_length :]` in `prepare_inputs_for_generation`
diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
index a52f96f52111..09e84457e9ee 100644
--- a/tests/generation/test_utils.py
+++ b/tests/generation/test_utils.py
@@ -5673,6 +5673,42 @@ def test_generate_can_restart_from_cache_and_new_tokens_only(self):
                 torch.testing.assert_close(cache1.layers[idx].keys, cache2.layers[idx].keys)
                 torch.testing.assert_close(cache1.layers[idx].values, cache2.layers[idx].values)
 
+    @pytest.mark.generate
+    def test_encoder_decoder_generate_can_restart_from_cache_without_decoder_mask(self):
+        """
+        Test that encoder-decoder models can continue generating from past key values with the FULL `decoder_input_ids`
+        and no `decoder_attention_mask`, as assistant models do in assisted decoding. The cached part of the decoder
+        inputs must not be processed again (otherwise the cache grows twice as fast as the sequence).
+        """
+        model = AutoModelForSeq2SeqLM.from_pretrained("hf-internal-testing/tiny-random-bart").to(torch_device)
+        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart")
+
+        initial_inputs = tokenizer("Hello, world!", return_tensors="pt").to(torch_device)
+        # `forced_eos_token_id` would force EOS as the last token of the truncated generation
+        generate_kwargs = {
+            "use_cache": True,
+            "do_sample": False,
+            "return_dict_in_generate": True,
+            "forced_eos_token_id": None,
+        }
+
+        truncated_outputs = model.generate(**initial_inputs, **generate_kwargs, max_new_tokens=2, min_new_tokens=2)
+        full_outputs = model.generate(**initial_inputs, **generate_kwargs, max_new_tokens=4, min_new_tokens=4)
+
+        restarted_outputs = model.generate(
+            **initial_inputs,
+            **generate_kwargs,
+            decoder_input_ids=truncated_outputs.sequences,
+            past_key_values=copy.deepcopy(truncated_outputs.past_key_values),
+            max_new_tokens=2,
+            min_new_tokens=2,
+        )
+
+        self.assertListEqual(full_outputs.sequences.tolist(), restarted_outputs.sequences.tolist())
+        self.assertEqual(
+            restarted_outputs.past_key_values.get_seq_length(), full_outputs.past_key_values.get_seq_length()
+        )
+
     @pytest.mark.generate
     @parameterized.expand(
         [