Repository: huggingface/transformers
Pull request: [CB] Ensure parallel decoding test passes using FA

This PR fixes a flaky test in continuous batching: the test was checking two streams of parallel decoding matched, but used `sdpa` as the attention implementation. But `sdpa` is not batch invariant, so despite a well formed attention mask, the test failed. With flash-attention, it is not the case anymore. It tested the following snippet: 
```
# Imports and implementations [...]
 
num_heads, head_dim, kv_per_query = 32, 64, 32

failures = 0
iters = 1000
for i in range(iters):

    # Create two identical query tokens
    Q = torch.randn(1, num_heads, 2, head_dim)
    Q[:, :, 1] = Q[:, :, 0]
    # Create identical KV cache for both queries
    K_or_V = torch.randn(1, num_heads, 2 * kv_per_query, head_dim)
    K_or_V[:, :, kv_per_query:] = K_or_V[:, :, :kv_per_query]

    # Attention mask to avoid cross-query interactions
    attention_mask = torch.zeros(1, 2, 2 * kv_per_query)
    attention_mask[0, 0, :kv_per_query] = 1
    attention_mask[0, 1, kv_per_query:] = 1

    # Use SDPA to compute attention
    output = compute_attention(Q, K_or_V, K_or_V, attention_mask)

    # Check using absolute and relative differences
    query_0_output, query_1_output = output[:, :, 0], output[:, :, 1]
    is_ok = torch.isclose(query_0_output, query_1_output).all().item()
    failures +=  1 - is_ok
```

and the number of failure for each implementation was:
- `sdpa`: 597
- `eager`: 156
- `kernels-community/flash-attn`: 0

So I went with flash attention for the test.

Changed files:
filename: tests/generation/test_continuous_batching.py
patch: @@ -38,6 +38,7 @@
 from transformers.testing_utils import (
     Expectations,
     require_deterministic_for_xpu,
+    require_flash_attn,
     require_torch_accelerator,
     slow,
     torch_device,
@@ -601,10 +602,9 @@ def test_block_sharing_with_hybrid_model(self) -> None:
 
         return self._test_block_sharing(model_id, num_layer_groups, input_msg, expected_generated_tokens)
 
-    # The test always passes on H100 with torch 2.9, but only passed case 0 on A100 with torch 2.6 and fails on A100
-    # with torch 2.9. This might be due to a GPU diff, so test might be flaky on the CI which runs on A10.
     @parameterized.expand([True, False])
     @require_torch_accelerator
+    @require_flash_attn  # otherwise the test can fail because attention bias has a very slight impact on SDPA and eager
     def test_num_return_sequences(self, allow_block_sharing: bool) -> None:
         model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
         tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")
@@ -616,7 +616,7 @@ def test_num_return_sequences(self, allow_block_sharing: bool) -> None:
         input_ids = [(x if isinstance(x, list) else x["input_ids"]) for x in tokenized]
 
         # Generation with continuous batching
-        model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="sdpa")
+        model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="flash_attention_2")
         model = model.to(torch_device).eval()
         model.generation_config.max_new_tokens = 30
         model.generation_config.do_sample = False