Repository: huggingface/transformers
Pull request: [FA2] Fix it finally - revert fa kwargs preparation

# What does this PR do?

Preparing the fa kwargs in advance is currently useless as it's being recomputed anyway. And worst, it actually break things, see following explanation.

When `_flash_attention_forward` is called, the following paths are taken inside the function:

```python
#  Contains at least one padding token in the sequence
if attention_mask is not None:
     # HERE IT ALWAYS RECOMPUTE THE FA KWARGS
    q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
        query_states, key_states, value_states, attention_mask, query_length, unpad_fn
    )
    ....
    out = flash_attn_varlen(...)

# Padding free, i.e. sequences flattened into one total sequence
elif is_fa_with_varlen_kwargs or is_fa_with_position_ids:
     # reshape everything
    q = query_states.reshape(-1, query_states.size(-2), query_states.size(-1))
     ...
    out = flash_attn_varlen(...)
    # reshape back
    out = out.view(query_states.size(0), -1, out.size(-2), out.size(-1))

# No padding
else:
    # THIS IS THE BATCHED VERSION OF FA2 -> we don't have padding here anyway, so no need to call varlen and worry about reshapes etc
    out = flash_fn(...)
```

So basically, we have 2 scenarios:
- attention mask is not None -> we recompute the fa kwargs anyway, so we're just wasting computation to get them before in generate
- attention mask is None -> why would we call `flash_varlen` if we don't have padding?? It forces us to reshape back and forth, and to compute fa kwargs where we don't need them -> better to not do it and simply call the batched fa version since we don't have padding anyway

Moreover, for hybrid models, different layer types could end up in different padding states. That is, with batched generation, if we have inputs of say (5k tokens, 4100 tokens) with a sliding window of 4096, then, after the prefill, the full layers have padding (they see the whole states), but the sliding layers don't have padding anymore (they only see the last 4096 - 1 tokens).
So, ahead-of-time-preparation would require careful handling of different layers.

But basically, for now preparing in advance is not only worst in performances (as we only compute things that are going to be recomputed/not needed), it's not logically correct! So reverting it!

I will explore later on if we can bring it back, as it would be useful for compilation!

Changed files:
filename: src/transformers/generation/utils.py
patch: @@ -46,7 +46,6 @@
 from ..integrations.deepspeed import is_deepspeed_zero3_enabled
 from ..integrations.fsdp import is_fsdp_managed_module
 from ..masking_utils import create_masks_for_generate
-from ..modeling_flash_attention_utils import prepare_fa_kwargs_from_position_ids
 from ..modeling_outputs import CausalLMOutputWithPast, Seq2SeqLMOutput
 from ..pytorch_utils import isin_mps_friendly
 from ..tokenization_utils import ExtensionsTrie
@@ -678,24 +677,12 @@ def prepare_inputs_for_generation(
         if encoder_attention_mask is not None:
             model_inputs["attention_mask"] = encoder_attention_mask
 
-        # 7. Prepare kwargs for flash attention to avoid recomputations
-        if "flash" in self.config._attn_implementation and self._supports_attention_backend:
-            (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = prepare_fa_kwargs_from_position_ids(
-                model_inputs["position_ids"], is_packed_sequence=False
-            )
-            model_inputs.update(
-                cu_seq_lens_q=cu_seq_lens_q.to(self.device),
-                cu_seq_lens_k=cu_seq_lens_k.to(self.device),
-                max_length_q=max_length_q,
-                max_length_k=max_length_k,
-            )
-
-        # 8. Forward ALL kwargs that are uninitialized (e.g. `use_cache`).
+        # 7. Forward ALL kwargs that are uninitialized (e.g. `use_cache`).
         for key, value in kwargs.items():
             if key not in model_inputs:
                 model_inputs[key] = value
 
-        # 9. Remove unexpected `generate` inputs (TODO @joao: fix trainer and examples)
+        # 8. Remove unexpected `generate` inputs (TODO @joao: fix trainer and examples)
         model_inputs.pop("labels", None)
         return model_inputs