Repository: huggingface/transformers
Pull request: :rotating_light: [`Flash Attention`] Fix sliding window size

Fixes #39408

This issue will occur especially during the prefill phase when sliding window < prefill size. Newer cache refactors automatically remove kvs beyond the window so inference would/should have been less affected. Small script to show the effects:
```python
import torch
from transformers import AutoTokenizer, AutoConfig, MistralForCausalLM


torch_device = "cuda"

config = AutoConfig.from_pretrained("mistralai/Mistral-7B-v0.1")
config.sliding_window = 4
model = MistralForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1", torch_dtype=torch.bfloat16, config=config, device_map={"": torch_device}
)

prompt = "My favourite condiment is "
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1", use_fast=False)
inputs = tokenizer(prompt, return_tensors="pt").to(torch_device)

# greedy generation outputs
generated_ids = model.generate(**inputs, max_new_tokens=20, do_sample=False)
text_sdpa = tokenizer.decode(generated_ids[0], skip_special_tokens=True)

model.set_attn_implementation("flash_attention_2")
generated_ids = model.generate(**inputs, max_new_tokens=20, do_sample=False)
text_fa = tokenizer.decode(generated_ids[0], skip_special_tokens=True)

print("sdpa")
print(text_sdpa)
print()
print("flash")
print(text_fa)
``` 
Outputs should match with the fix, before they were not. Will add a common test so it won't happen again. Edit: modified the existing test instead

cc @Cyrilvallez @ArthurZucker 

Changed files:
filename: src/transformers/modeling_flash_attention_utils.py
patch: @@ -532,7 +532,12 @@ def _process_flash_attention_kwargs(
         flash_kwargs["dropout_p"] = dropout
 
     if supports_mapping["window_size"] and sliding_window is not None and key_length > sliding_window:
-        flash_kwargs["window_size"] = (sliding_window, sliding_window)
+        # The flash attention API sets inclusive boundaries, i.e. (4, 0) would take 4 tokens to the left
+        # and the current token for a total size of 5. However, we usually define our window sizes by
+        # their total window size (when causal). Encoder models as of now seldom use SWA and when they
+        # do, they have a custom workaround (e.g. ModernBERT) which would align with this symmetric logic, i.e.
+        # for a total of `2*sliding_window + 1`.
+        flash_kwargs["window_size"] = (sliding_window - 1, sliding_window - 1)
 
     if supports_mapping["deterministic"]:
         flash_kwargs["deterministic"] = (

filename: tests/test_modeling_common.py
patch: @@ -3494,6 +3494,10 @@ def flash_attn_inference_equivalence(self, attn_implementation: str, padding_sid
             # flash attention variants does not always support arbitrary headim
             config = self._prepare_config_headdim(config, 16)
 
+            # forcing the prefill size to go over sliding window size to check for SWA correctness
+            if getattr(config, "sliding_window", None):
+                config.sliding_window = 2
+
             # TODO it is unclear why saving and reloading with dtype works while
             # casting with `.to(dtype=..., device=...)` does not.
             # Discovered on tests with `Bart` models.