Fix slow test_moshika_greedy_unconditional_fp16

src/transformers/cache_utils.py

@@ -95,7 +95,7 @@ def _sliding_cache_update(
     # Sliding window logic for generation phase or prefill < window
     slicing = torch.arange(max_cache_len, device=value_states.device)
     current_seq_len = cache_position[-1] + 1  # Use last position to determine current length
-    to_shift = current_seq_len > max_cache_len
+    to_shift = current_seq_len >= max_cache_len
     indices = (slicing + to_shift.sum()) % max_cache_len
 
     k_out_shifted = k_cache[:, :, indices]

src/transformers/generation/utils.py

@@ -2481,7 +2481,7 @@ def generate(
         # - `model_kwargs` may be updated in place with a cache as defined by the parameters in `generation_config`.
         # - different models have a different cache name expected by the model (default = "past_key_values")
         # - `max_length`, prepared above, is used to determine the maximum cache length
-        max_cache_length = generation_config.max_length - 1
+        max_cache_length = generation_config.max_length
         if (
             inputs_tensor.shape[1] != input_ids_length
             and model_input_name == "inputs_embeds"

src/transformers/models/moshi/modeling_moshi.py

@@ -1862,28 +1862,15 @@ def forward(
 
     def _prepare_attention_mask_for_generation(
         self,
-        input_ids: torch.LongTensor,
+        inputs_tensor: torch.Tensor,
         generation_config: GenerationConfig,
-        kwargs: dict[str, Any],
+        model_kwargs: dict[str, Any],
     ) -> torch.LongTensor:
-        pad_token_id = generation_config.pad_token_id
-        eos_token_id = generation_config.eos_token_id
-
-        default_attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=input_ids.device)
-        if pad_token_id is None:
-            return default_attention_mask
-
-        is_pad_token_in_inputs = (pad_token_id is not None) and torch.isin(input_ids, pad_token_id).any()
-        is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or ~torch.isin(
-            eos_token_id, pad_token_id
-        ).any()
-        can_infer_attention_mask = is_pad_token_in_inputs * is_pad_token_not_equal_to_eos_token_id
-        attention_mask_from_padding = input_ids.ne(pad_token_id).long()
-
-        attention_mask = (
-            attention_mask_from_padding * can_infer_attention_mask + default_attention_mask * ~can_infer_attention_mask
+        return super()._prepare_attention_mask_for_generation(
+            inputs_tensor=inputs_tensor,
+            generation_config=generation_config,
+            model_kwargs={},
         )
-        return attention_mask
 
     def _prepare_inputs_embeds_for_generation(
         self,
@@ -2101,12 +2088,6 @@ def generate(
         kwargs_depth_decoder = depth_decoder_generation_config
 
         attention_mask = kwargs.pop("attention_mask", None)
-        if attention_mask is None:
-            attention_mask = self._prepare_attention_mask_for_generation(
-                input_ids=input_ids,
-                generation_config=generation_config,
-                kwargs=kwargs,
-            )
         (
             inputs_embeds,
             input_ids,
@@ -2284,17 +2265,14 @@ def prepare_inputs_for_generation(
         if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
             if model_inputs["inputs_embeds"] is not None:
                 batch_size, sequence_length, _ = inputs_embeds.shape
-                device = inputs_embeds.device
             else:
                 batch_size, sequence_length = input_ids.shape
-                device = input_ids.device
 
             attention_mask = self.decoder.model._prepare_4d_causal_attention_mask_with_cache_position(
                 attention_mask,
                 sequence_length=sequence_length,
                 target_length=past_key_values.get_max_cache_shape(),
                 dtype=self.decoder.lm_head.weight.dtype,
-                device=device,
                 cache_position=cache_position,
                 batch_size=batch_size,
                 config=self.config,

tests/models/moshi/test_modeling_moshi.py

@@ -998,7 +998,11 @@ def test_moshiko_greedy_unconditional_fp16(self):
         some_expected_audio_tokens = [[1049, 127], [1700, 243], [1626, 457], [546, 290], [306, 306], [1443, 1443], [1871, 428], [2008, 1744]]  # fmt: skip
 
         model_outputs = model.generate(
-            do_sample=False, depth_decoder_do_sample=False, return_audio_codes=True, max_new_tokens=10
+            do_sample=False,
+            depth_decoder_do_sample=False,
+            return_audio_codes=True,
+            max_new_tokens=10,
+            # cache_implementation="dynamic"
         )
 
         # make sure audio encoded codes are correct