[CB] Ensure parallel decoding test passes using FA

[remi-or]

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.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[vasqu]

Thx also double checked locally, works 👍

[vasqu]

Thx for the comment