Repository: pytorch
Issue: [BC BREAKING] Change default behavior of scaled_dot_product_attention's causal masking alignment
Labels: module: bc-breaking,module: nn,triaged,topic: bc breaking,module: sdpa

## Current Behavior:
The `torch.nn.functional.scaled_dot_product_attention` function currently applies causal masking to the top left corner of the attention matrix. This well defined when seqlen_q equals seqlen_kv, but an alignment choice needs to made when they are not.

#### Current alignment choice
If seqlen_q != seqlen_kv and causal=True, the causal mask is aligned to the the top-left corner.

For example, if seqlen_q = 2 and seqlen_kv = 5, the causal mask (1 = keep, 0 = masked out) is:
```
    1 0 0 0 0
    1 1 0 0 0
```

If seqlen_q = 5 and seqlen_kv = 2, the causal mask is:
```
    1 0
    1 1
    1 1
    1 1
    1 1
```

## Proposal:
I propose changing the default behavior of the scaled_dot_product_attention function's causal masking to align with the bottom right corner of the attention matrix. This change would be backward incompatible for inputs when seqlen_q does not equal seqlen_kv, as it would shift the masking pattern.

For example, if seqlen_q = 2 and seqlen_kv = 5, the causal mask (1 = keep, 0 = masked out) is:
```
    1 1 1 1 0
    1 1 1 1 1
```
If seqlen_q = 5 and seqlen_kv =  2, the causal mask is:
```
    0 0
    0 0
    0 0
    1 0
    1 1
 ```
_If the row of the mask is all zero, the output will be zero._

## Benefits:

This choice of mask is beneficial when performing autoregressive decoding for LLMs. 
For example lets say that your prefill prompt consisted of 4 tokens and you had a max KV cache size of 8. Then at step 1 of decoding, your query sequence length would be size 1 (index position 5) and your KV input would be of size 5. The query's seqlen is 1 but it really represents the 5 token in the sequence. So the mask for this attention should be:
```
1 1 1 1 1 0 0 0
```

With the existing implementation of causal, the produced attention mask would be. 
```
1 0 0 0 0 0 0 0 
``` 

## Drawbacks:

Backward compatibility: The proposed change would break backward compatibility for cases where seqlen_q doesn't equal seqlen_kv.
Existing code: Users who rely on the current masking alignment would need to update their code if they want to adopt the new behavior.

## Implementation:
To implement this change, the following steps need to be taken:

- Modify the `scaled_dot_product_attention_math` function to align the causal masking with the bottom right corner of the attention matrix.
- Change the default call `_scaled_dot_product_efficient_attention` to https://github.com/pytorch/pytorch/blob/3488837ec118cc8f8bf9fe10a8c5a32726ba0a7b/aten/src/ATen/native/transformers/cuda/attention.cu#L734 to use the bottom right
- After updating to FlashAttentionV2 we would apply the changes from this commit: https://github.com/Dao-AILab/flash-attention/commit/9e5e8bc91e30af5cdc321362b553f6c0da332e30 
- Update the fused CPU implementation to obey the above semantics.
- Update the function's documentation and any relevant examples to reflect the new masking behavior.

## Additional Context:
Both Xformers and FlashAttention have made this behavior the new default



cc @ezyang @gchanan @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @bhosmer @cpuhrsch @erichan1