`_update_causal_mask` very slow in latest transformer release
Author: alessandropalla
Pull request: no

### System Info


- `transformers` version: 4.38.1
- Platform: Windows-10-10.0.22621-SP0
- Python version: 3.11.2
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.2
- Accelerate version: 0.26.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.0+cpu (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help?

@ArthurZucker 

### Information

- [X] 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

Run any llama model

### Expected behavior

In latest `transformers` release (v4.38.1) the generation for llama model is significantly slower than before. After digging I found that is because of the `_update_causal_mask` function. Some heavy operations that runs at every inference on the fairly big `causal_mask`  (is a `config.max_position_embeddings` x `config.max_position_embeddings` array) buffer causes a significant slowdown for llama model. The same function is in `gemma` models

![image](https://github.com/huggingface/transformers/assets/28634533/4c68ca71-25f6-4a4e-bd1d-590fa410ea2b)

Proposed solution: have the causal mask registered buffer already having  `torch.finfo(dtype).min` instead of 1, repeat and cast only if needed


Comments:
Agree that it might not be optimal, and PRs are welcome. 
I benchmarked vs the `_prepare_causal_mask` and did not have significant slow down. This depends on the padding you have of course. 
The causal mask can be registered, but you end up with casting issues if you init with a different torch dtype and cast to another one using "model.to". Very brittle. 

The profiler seems to indicate `aten:min` is taking some time. If this can speed up would be nice. Not sure the produce is what is slowing down? 

Also if you want fast generation, like actually fast, would recommend using static kv cache

Alright so:
```python
                causal_mask = causal_mask.mul(~torch.all(causal_mask == causal_mask.min(), dim=-1)[..., None]).to(
                    dtype
                )
```
this can be improved IMO

```diff
+        min_dtype = torch.finfo(dtype).min
        # We use the current dtype to avoid any overflows
        causal_mask = self.causal_mask[None, None, :, :].repeat(batch_size, 1, 1, 1).to(dtype) * min_dtype

        causal_mask = causal_mask.to(dtype=dtype, device=device)
        if attention_mask is not None and attention_mask.dim() == 2:
            mask_length = attention_mask.shape[-1]
            padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)
            causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(
+                padding_mask, min_dtype
            )

        if self.config._attn_implementation == "sdpa":
            is_tracing = torch.jit.is_tracing() or isinstance(input_tensor, torch.fx.Proxy)
            if not is_tracing and attention_mask is not None and torch.any(attention_mask != 1):
+                causal_mask = causal_mask.mul(~torch.all(causal_mask == min_dtype, dim=-1)[..., None]).to(
                    dtype
                )
```

Interesting. I'm actually using a script that uses static shape and a custom generate (the one I posted here: https://github.com/huggingface/transformers/issues/28075) that used to work very well before. Do you have an example script to use the new static cache flow I'd love to try

https://gist.github.com/ArthurZucker/af34221def212259b43d55a2811d2dbb this should work

Could you benchmark again with proposed changes? 


I'll do it asap

We basically got rid of the tensor memory allocation, which was already slow before + improved the way we cope with right padding. 
Padding right introduced fully non attended: `if not is_tracing and attention_mask is not None and torch.any(attention_mask != 1):` this condition is triggered, so we have to update the entire mask with the multipilcation which is less costly than anything that was done before

Avoiding the `causal_mask.min()` helps a lot. I can gather more data on that function to see what can be done further

Awesome! Feel free to open a PR 🤗 

PR (in progress) here: https://github.com/huggingface/transformers/pull/29210