SDPA `is_causal=False` has no effect due to `LlamaModel._prepare_4d_causal_attention_mask_with_cache_position`
Author: ringohoffman
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.48.3
- Platform: Linux-5.15.0-130-generic-x86_64-with-glibc2.35
- Python version: 3.9.21
- Huggingface_hub version: 0.28.1
- Safetensors version: 0.5.2
- Accelerate version: 1.3.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.6.0+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA H100 80GB HBM3

### Who can help?

@ArthurZucker @Cyrilvallez 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Observe `is_causal=False` has no effect when using `attn_implementation="sdpa"` with an `attention_mask` with at least one `False` element:

```python
import torch
import transformers

device = torch.device("cuda:0")
input_ids = torch.tensor(
    [
        [
            128000, 128006,   9125, 128007,    271,     34,   7747,    553,    279,
            2768,   1495,    439,   1694,   5552,    311,   5557,     11,  17452,
            11,  10034,     11,    477,  11759,     13, 128009, 128006,    882,
            128007,    271,    791,    502,  77355,   3280,    690,  10536,   1022,
            449,    264,  72097,   2489,   1990,  35812,    323,  64921,     13,
            128009, 128006,  78191, 128007,    271,  42079, 128009, 128004, 128004,
            128004, 128004
        ]
    ],
    device=device,
)
attention_mask = torch.tensor(
    [
        [
            True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
            True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
            True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
            True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
            True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
            True,  True, False, False, False, False
        ]
    ],
    device=device,
)

with device:
    model = transformers.AutoModelForCausalLM.from_pretrained(
        "/models/meta-llama/Llama-3.2-1B-Instruct",  # https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct
        attn_implementation="sdpa",
        torch_dtype=torch.bfloat16,
    )

causal_logits = model(input_ids, attention_mask=attention_mask, is_causal=True).logits
noncausal_logits = model(input_ids, attention_mask=attention_mask, is_causal=False).logits

torch.testing.assert_close(causal_logits, noncausal_logits)  # shouldn't be true, otherwise what is_causal controlling?
```

Observe that mocking `LlamaModel._prepare_4d_causal_attention_mask_with_cache_position` with an implementation that just replicates the `attention_mask` also has no effect when using `is_causal=True`:

```python
from unittest import mock

def _prepare_4d_causal_attention_mask_with_cache_position(
    attention_mask: torch.Tensor,
    sequence_length: int,
    target_length: int,
    dtype: torch.dtype,
    device: torch.device,
    cache_position: torch.Tensor,
    batch_size: int,
    **kwargs,
):
    min_dtype = torch.tensor(torch.finfo(dtype).min, dtype=dtype, device=attention_mask.device)
    return ~attention_mask.view(batch_size, 1, 1, sequence_length).expand(batch_size, 1, sequence_length, sequence_length) * min_dtype

with mock.patch.object(model.model, "_prepare_4d_causal_attention_mask_with_cache_position", _prepare_4d_causal_attention_mask_with_cache_position):
    sdpa_causal_logits = model(input_ids, attention_mask=attention_mask, is_causal=True).logits

hf_causal_logits = model(input_ids, attention_mask=attention_mask, is_causal=True).logits

torch.testing.assert_close(sdpa_causal_logits, hf_causal_logits)  # shouldn't be true, otherwise what is _prepare_4d_causal_attention_mask_with_cache_position doing?
```

### Expected behavior

1. At the very least, `LlamaModel. _prepare_4d_causal_attention_mask_with_cache_position` should respect `is_causal=False`. Right now, it always returns a causal mask when using sdpa with sequence_length > 1 and an attention_mask with at least one False element.
2. It is not really clear to me why we aren't purely relying on SDPA's own `is_causal` parameter. My 2nd example demonstrates that the current implementation of `LlamaModel. _prepare_4d_causal_attention_mask_with_cache_position` definitely isn't always necessary... so when is it necessary? Or what parts are necessary? Looking at the equivalent implementation that PyTorch describes for [`scaled_dot_product_attention`](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html), it seems like we are replicating a bit of their handling of `attn_mask`. Also notably there are 4 separate CUDA allocations happening in the current implementation (`torch.full`, `torch.triu`, `torch.arange`, `Tensor.clone`) compared to my proposed 1.

Comments:
@gante 

Kinda related to https://github.com/huggingface/transformers/pull/36049. Not sure if the PR will solve the problem

> Kinda related to [#36049](https://github.com/huggingface/transformers/pull/36049). Not sure if the PR will solve the problem

#36049 deals purely with typing. It will definitely not solve this issue.

To fix this issue at a minimum, `LlamaModel._prepare_4d_causal_attention_mask_with_cache_position` needs to be passed `is_causal` and return a different result depending on that.

But I am asking what parts of the current implementation of `LlamaModel._prepare_4d_causal_attention_mask_with_cache_position` are necessary since it seems like `scaled_dot_product_attention` handles a lot of what it is doing already, as shown by my second code example. Maybe some of the `cache_position` stuff we might keep?

@ringohoffman yep, sorry! I meant related in a sense that `is_causal` was never an acceptable kwarg for the model's forward until the PR is merged. Therefore we don't expect the arg to have any effect. 

The `_prepare_4d_causal_attention_mask_with_cache_position` does indeed the same thing was SDPA's causal mask in most simple cases. But the SDPA mask doesn't handle padding ids which we add to the mask, as you mentioned. Also we have a variety of Key-Value caching methods, like Static Cache. When these methods are used, we need to prepare special mask by masking out the extra positions at the right end.



This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

This issue is still outstanding.

@ringohoffman I would like to echo @zucchini-nlp's comment:

> I meant related in a sense that is_causal was never an acceptable kwarg for the model's forward until the PR is merged. Therefore we don't expect the arg to have any effect.

Some encoder-decoder models (like Whisper) do accept `is_causal` in the initialization of the attention layers. But this is not the case of decoder-only models like Llama.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.