[FA2] Fix it finally - revert fa kwargs preparation
Author: Cyrilvallez
Pull request: yes
Labels: name: for patch

# 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!

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_40161). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Yes! The following snippet:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
torch.manual_seed(0)
model_id = "openai/gpt-oss-20b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
  model_id, device_map="auto", use_kernels=True, attn_implementation="kernels-community/vllm-flash-attn3",
)

messages = [
  {"role": "user", "content": "How many rs are in the word 'strawberry'?"},
]

inputs = tokenizer.apply_chat_template(
  messages,
  add_generation_prompt=True,
  return_tensors="pt",
  return_dict=True,
).to(model.device)

generated = model.generate(**inputs, max_new_tokens=1024, do_sample=True)
print(tokenizer.decode(generated[0][inputs['input_ids'].shape[-1] :]))
```

returns

```

analysis
User asks: "How many rs are in the word 'strawberry'?". We need to count number of 'r' letters. Let's analyze: The word is 'strawberry'. Letters: s, t, r, a, w, b, e, r, r, y. Actually 'strawberry' is spelled s t r a w b e r r y. Yes, there are three r's: at positions 3, 8, 9. So answer: 3.

Could also ask for 'Rs' capital? But question: "How many rs are in the word 'strawberry'?" If case-insensitive, indeed 3. So answer: 3. We might respond simply.

But user might want just number. I'll respond: 3.

Add mention explanation.

Let's provide final answer.

assistant
final
There are **3** “r” letters in the word **strawberry** (s t **r** a w b e **r** **r** y).

```

which is the expected output!