Don't pass masks in model_inputs when not needed

src/transformers/generation/configuration_utils.py

@@ -379,6 +379,8 @@ class GenerationConfig(PushToHubMixin):
     _pad_token_tensor: "torch.Tensor | None"
     _decoder_start_token_tensor: "torch.Tensor | None"
 
+    _can_drop_attention_mask: bool = False
+
     # Hash to detect whether the instance was modified after loading
     _original_object_hash: int | None
 

src/transformers/generation/utils.py

@@ -43,13 +43,14 @@
     resolve_trust_remote_code,
 )
 from ..integrations.deepspeed import is_deepspeed_zero3_enabled
-from ..masking_utils import create_masks_for_generate
+from ..masking_utils import create_masks_for_generate, fast_all
 from ..tokenization_python import ExtensionsTrie
 from ..utils import (
     ModelOutput,
     TransformersKwargs,
     has_file,
     is_accelerate_available,
+    is_tracing,
     logging,
 )
 from ..utils.generic import is_flash_attention_requested
@@ -653,6 +654,7 @@ def prepare_inputs_for_generation(
         attention_mask: torch.LongTensor | None = None,
         inputs_embeds: torch.FloatTensor | None = None,
         is_first_iteration: bool | None = False,
+        can_drop_attention_mask: bool = False,
         **kwargs,
     ):
         """
@@ -728,7 +730,7 @@ def prepare_inputs_for_generation(
                 is_first_iteration=is_first_iteration,
             )
 
-        if attention_mask is not None:
+        if attention_mask is not None and not can_drop_attention_mask:
             model_inputs[attention_mask_key] = attention_mask
 
         if encoder_attention_mask is not None:
@@ -2766,6 +2768,15 @@ def generate(
             generation_config, model_kwargs, generation_mode, batch_size, max_cache_length
         )
 
+        attention_mask = model_kwargs.get("attention_mask")
+        generation_config._can_drop_attention_mask = (
+            attention_mask is not None
+            and not self.config.is_encoder_decoder
+            and not is_tracing(attention_mask)
+            and not self._valid_auto_compile_criteria(model_kwargs, generation_config)
+            and bool(fast_all(attention_mask))
+        )
+
         if self.device.type != input_ids.device.type:
             warnings.warn(
                 "You are calling .generate() with the `input_ids` being on a device type different"
@@ -3025,7 +3036,10 @@ def _sample(
                 if prefill_consumed:
                     next_sequence_length = 1 if model_kwargs["use_cache"] else None
                     model_inputs = self.prepare_inputs_for_generation(
-                        input_ids, next_sequence_length=next_sequence_length, **model_kwargs
+                        input_ids,
+                        next_sequence_length=next_sequence_length,
+                        can_drop_attention_mask=generation_config._can_drop_attention_mask,
+                        **model_kwargs,
                     )
                     outputs = model_forward(**model_inputs, return_dict=True)
                 prefill_consumed = True
@@ -3527,7 +3541,10 @@ def _beam_search(
                 flat_running_sequences = self._flatten_beam_dim(running_sequences[:, :, :cur_len])
                 next_sequence_length = 1 if model_kwargs["use_cache"] else None
                 model_inputs = self.prepare_inputs_for_generation(
-                    flat_running_sequences, next_sequence_length=next_sequence_length, **model_kwargs
+                    flat_running_sequences,
+                    next_sequence_length=next_sequence_length,
+                    can_drop_attention_mask=generation_config._can_drop_attention_mask,
+                    **model_kwargs,
                 )
                 model_outputs = self(**model_inputs, return_dict=True)
             prefill_consumed = True
@@ -3875,6 +3892,7 @@ def _assisted_decoding(
                 candidate_input_ids,
                 next_sequence_length=next_sequence_length,
                 is_first_iteration=is_first_iteration,
+                can_drop_attention_mask=generation_config._can_drop_attention_mask,
                 **candidate_kwargs,
             )
 
@@ -4104,6 +4122,7 @@ def _prefill(
                 input_ids,
                 next_sequence_length=next_sequence_length,
                 is_first_iteration=is_first_iteration,
+                can_drop_attention_mask=generation_config._can_drop_attention_mask,
                 **model_kwargs,
             )
             return self(**model_inputs, return_dict=True)
@@ -4135,7 +4154,9 @@ def _prefill(
                     model_kwargs["attention_mask"] = attention_mask[:, :current_length]
                 if position_ids is not None:
                     model_kwargs["position_ids"] = position_ids[:, past_length:current_length]
-                model_inputs = self.prepare_inputs_for_generation(input_chunk, **model_kwargs)
+                model_inputs = self.prepare_inputs_for_generation(
+                    input_chunk, can_drop_attention_mask=generation_config._can_drop_attention_mask, **model_kwargs
+                )
 
                 outputs = model_forward(**model_inputs, return_dict=True)