[`generate` + static cache + `torch.compile`] ability to pass statically shaped 4D `attention_mask` to the model forward
Author: fxmarty
Pull request: no
Labels: name: Generation

name: Compilation

name: Cache

### Feature request

Currently, the `attention_mask` passed is 2D and of dynamic shapes.

This causes issues when using a compiled model with `model.forward = torch.compile(model.forward, mode="reduce-overhead")`, see https://github.com/pytorch/pytorch/issues/120309 & https://github.com/huggingface/transformers/pull/29114.

```
----- in forward 0
name=input_ids, shape=torch.Size([2, 7]), stride=(7, 1), dtype=torch.int64, device=cuda:0
name=position_ids, shape=torch.Size([2, 7]), stride=(7, 1), dtype=torch.int64, device=cuda:0
name=cache_position, shape=torch.Size([7]), stride=(1,), dtype=torch.int64, device=cuda:0
name=past_key_values, value=None
name=use_cache, value=True
name=attention_mask, shape=torch.Size([2, 7]), stride=(7, 1), dtype=torch.int64, device=cuda:0
forward call latency: 1784.737 ms      <---------------------------- EXTREMELY SLOW.
----- in forward 1
name=input_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=position_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=cache_position, shape=torch.Size([1]), stride=(1,), dtype=torch.int64, device=cuda:0
name=past_key_values, value=None
name=use_cache, value=True
name=attention_mask, shape=torch.Size([2, 8]), stride=(8, 1), dtype=torch.int64, device=cuda:0
forward call latency: 1851.579 ms      <---------------------------- EXTREMELY SLOW.
----- in forward 2
name=input_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=position_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=cache_position, shape=torch.Size([1]), stride=(1,), dtype=torch.int64, device=cuda:0
name=past_key_values, value=None
name=use_cache, value=True
name=attention_mask, shape=torch.Size([2, 9]), stride=(9, 1), dtype=torch.int64, device=cuda:0
forward call latency: 1421.504 ms      <---------------------------- EXTREMELY SLOW.
----- in forward 3
name=input_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=position_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=cache_position, shape=torch.Size([1]), stride=(1,), dtype=torch.int64, device=cuda:0
name=past_key_values, value=None
name=use_cache, value=True
name=attention_mask, shape=torch.Size([2, 10]), stride=(10, 1), dtype=torch.int64, device=cuda:0
forward call latency: 1740.283 ms      <---------------------------- EXTREMELY SLOW.
```

Instead, we may want to pass directly 4D masks of static shape to the model, avoiding cuda graphs recomputation. This may allow the compile time to go down from >3 min to a mere 50s.

### Motivation

/

### Your contribution

/

Comments:
cc @gante 


@fxmarty I had a quick look at this -- we still have models (like `gpt2`) that exclusively accept 2D masks. We would have to rework that before making `generate` prepare a 4D mask.

We may want to work with padded tensors instead? In our TF/PT XLA implementation we have static shapes everywhere.

@gante thank you! What do you mean by work with padded tensors?

@fxmarty If we want to generate with `max_length=512`, the attention mask is always kept with sequence length = 512. Data is moved around inside the tensor as needed.

Oh yes, I think that's what I am suggesting here!