generate: drop an all-ones attention mask from model_kwargs

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"
 
+    _mask_length: int | None = None
+
     # 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
@@ -2266,9 +2267,19 @@ def _valid_auto_compile_criteria(
         # Note: for some models that only use linear attention (e.g. Mamba), even a DynamicCache is compilable since all
         # Encoder-decoder models hold that cache in a subcache, so we unwrap it to check the cache the decoder actually generates with
         decoder_cache = cache.self_attention_cache if isinstance(cache, EncoderDecoderCache) else cache
-        # Note: for some models that only use linear attention (e.g. Mamba), even a DynamicCache is compilable since all
-        # layers are, but we don't want to ALWAYS compile when calling `generate`, so we check the type
-        using_compilable_cache = cache is not None and cache.is_compileable and type(decoder_cache) is not DynamicCache
+        if cache is not None:
+            # Note: for some models that only use linear attention (e.g. Mamba), even a DynamicCache is compilable since all
+            # layers are, but we don't want to ALWAYS compile when calling `generate`, so we check the type
+            using_compilable_cache = cache.is_compileable and type(decoder_cache) is not DynamicCache
+        else:
+            # There is no cache yet, so answer for the one `_prepare_cache_for_generation` will build.
+            # It builds none at all when `use_cache` is off or when the model makes its own, and of
+            # what it does build only a static implementation is compileable (the rest is a DynamicCache).
+            using_compilable_cache = (
+                bool(generation_config.use_cache)
+                and self._supports_default_dynamic_cache()
+                and generation_config.cache_implementation in ALL_STATIC_CACHE_IMPLEMENTATIONS
+            )
         can_compile = valid_hardware and using_compilable_cache
 
         # Exception 1: Some quantization methods do not support compilation
@@ -2700,6 +2711,25 @@ def generate(
         if not kwargs_has_position_ids and accepts_position_ids and not self.config.is_encoder_decoder:
             model_kwargs["position_ids"] = self._prepare_position_ids_for_generation(inputs_tensor, model_kwargs)
 
+        # A mask of all ones says nothing to the model, so drop it. Not under compilation though: dynamo
+        # guards on `attention_mask is None`, so dropping it costs a recompile.
+        attention_mask = model_kwargs.get("attention_mask")
+        decoding_name = GENERATION_MODES_MAPPING[generation_mode]
+        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))
+            # Hub recipes and the candidate generators index `model_kwargs["attention_mask"]` themselves
+            and "/" not in decoding_name
+            and decoding_method is getattr(GenerationMixin, decoding_name)
+            and generation_mode != GenerationMode.ASSISTED_GENERATION
+        )
+        if can_drop_attention_mask:
+            generation_config._mask_length = attention_mask.shape[-1]
+            del model_kwargs["attention_mask"]
+
         if self.config.is_encoder_decoder and "encoder_outputs" not in model_kwargs:
             # if model is encoder decoder encoder_outputs are created and added to `model_kwargs`
             model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(
@@ -4091,10 +4121,12 @@ def _prefill(
             if use_inputs_embeds:
                 next_sequence_length = model_kwargs["inputs_embeds"].shape[1] - past_length
             else:
-                attention_mask_key = "decoder_attention_mask" if self.config.is_encoder_decoder else "attention_mask"
-                attention_mask = model_kwargs.get(attention_mask_key)
-                # 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 attention_mask is not None and input_ids.shape[1] == attention_mask.shape[1]:
+                mask_key = "decoder_attention_mask" if self.config.is_encoder_decoder else "attention_mask"
+                attention_mask = model_kwargs.get(mask_key)
+                mask_length = attention_mask.shape[1] if attention_mask is not None else generation_config._mask_length
+                # Only slice when the inputs hold the whole sequence; if they hold only the new tokens there is
+                # nothing to do
+                if mask_length is not None and input_ids.shape[1] == mask_length:
                     # inputs will be sliced as `input_ids[:, -next_sequence_length :]` in `prepare_inputs_for_generation`
                     next_sequence_length = input_ids.shape[1] - past_length