>3-5x faster torch.compile forward compilation for autoregressive decoder models
Author: fxmarty
Pull request: yes

As detailed in https://github.com/huggingface/transformers/pull/29114 and https://github.com/huggingface/transformers/issues/29165, the current implementation of `GenerationMixin.generate` passes 2D integer tensors of variable length as `attention_mask` to the models forward, that are transformed into 4D tensors of static shape in the models forward for decoder models.

This is an issue when using torch.compile with static cache, as then the model forward has inputs of various shape at each decoding step, which retriggers a cuda graph capture. This is especially bad when setting large enough `max_new_tokens` (>64) as cuda graphs are recaptured many many times.

Effectively as per https://pytorch.org/docs/stable/torch.compiler_cudagraph_trees.html#cudagraph-trees-integration, what is to be considered the compilation time is the `torch.compile` call, and the subsequent two forward calls (without guard invalidation, cudagraph invalidation).

For example, below, we set `max_new_tokens = 256`, with an 8B model on an A100, and on main the compilation using `model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)` takes:
```
----- GENERATE WITHOUT COMPILE
Non-compiled generate call took (no warmup): 12.546 s
decoded ['I would like (redacted)"]
compiling...
Finished compile call: 0.926 s
----- GENERATE WITH COMPILE
/home/felix/miniconda3/envs/fx/lib/python3.9/site-packages/torch/_inductor/compile_fx.py:150: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
  warnings.warn(
First generate call took: 89.536 s (the first run warms up each graph, which does things like CuBlas or Triton benchmarking)
-------- decoded 1 ['I would like (redacted)"]
Second generate call took: 267.591 s (The second run does a CUDA Graph recording, and replays it)
-------- decoded 2 ['I would like (redacted)"]
Third generate call took: 3.421 s (Finally we hit the optimized, CUDA Graph replay path)
-------- decoded 3 ['I would like (redacted)"]

```

while with this PR the compilation takes
```
----- GENERATE WITHOUT COMPILE
Non-compiled generate call took (no warmup): 11.672 s
decoded ['I would like (redacted)"]
compiling...
Finished compile call: 0.999 s
----- GENERATE WITH COMPILE
First generate call took: 85.505 s (the first run warms up each graph, which does things like CuBlas or Triton benchmarking)
-------- decoded 1 ['I would like (redacted)"]
Second generate call took: 4.845 s (The second run does a CUDA Graph recording, and replays it)
-------- decoded 2 ['I would like (redacted)"]
Third generate call took: 3.495 s (Finally we hit the optimized, CUDA Graph replay path)
-------- decoded 3 ['I would like (redacted)"]
```

repro here https://gist.github.com/fxmarty/74b4463935ada44ef7755a544cae8773, using `dynamic=True` is even worse

cc @LysandreJik @ArthurZucker @amyeroberts @gante I think we'll need a lot of cleanup in `modeling_attn_mask_utils`  at some point - we have many `_prepare_attention_mask_xxx` functions (even some in `PreTrainedModel`)

Closes https://github.com/huggingface/transformers/issues/29165

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_32227). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

> In accordance to transformers-philosophy, the causal mask creation should stay in the modeling file as it is CORE. Wrapping it in a function is not a problem !

Understood.

> This attention mask 4D created in prepare_input_for_generation should be done in generate automatically

I am not sure about this, as it is model specific (some models support `cache_position` & static cache, others don't). Do you mean that the 2D -> 4D update should preferably be handled in `_update_model_kwargs_for_generation` instead of `prepare_inputs_for_generation`?

Edit: one issue I see with using `prepare_inputs_for_generation` is that we regenerate the 4D mask at every step which is really not needed, not sure how this may be addressed

I'll add some tests here.

@ArthurZucker I see. Let me try that.

Edit: It is a bit involved to generate the 4d mask in generate due to dependencies on the 2d mask in `prepare_inputs_for_generation` and `_update_model_kwargs_for_generation`. So as to merge this today, let's leave a refactor to an other PR (should be easy with `make fix-copies`)

I checked the slow tests for a first few archs. There are no other tests failing due to this PR. Although on `main` there are a few tests failing unrelated to this PR which may be concerning (on A100, maybe some reference are for an other GPU, especially gemma):

```
tests/models/cohere/test_modeling_cohere.py::CohereModelTest::test_generate_compile_fullgraph
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_generate_compile_fullgraph
tests/models/gemma/test_modeling_gemma.py::GemmaIntegrationTest::test_model_2b_bf16
tests/models/gemma/test_modeling_gemma.py::GemmaIntegrationTest::test_model_2b_bf16_dola
tests/models/gemma/test_modeling_gemma.py::GemmaIntegrationTest::test_model_2b_eager
tests/models/gemma/test_modeling_gemma.py::GemmaIntegrationTest::test_model_2b_sdpa
tests/models/gemma/test_modeling_gemma.py::GemmaIntegrationTest::test_model_7b_4bit
tests/models/gemma/test_modeling_gemma.py::GemmaIntegrationTest::test_model_7b_bf16
tests/models/gemma2/test_modeling_gemma2.py::GemmaModelTest::test_generate_compile_fullgraph
tests/models/gemma2/test_modeling_gemma2.py::Gemma2ModelTest::test_sdpa_can_dispatch_on_flash
tests/models/gemma2/test_modeling_gemma2.py::Gemma2ModelTest::test_sdpa_equivalence
tests/models/llama/test_modeling_llama.py::LlamaModelTest::test_generate_compile_fullgraph
tests/models/llama/test_modeling_llama.py::LlamaIntegrationTest::test_model_7b_logits
tests/models/llama/test_modeling_llama.py::LlamaIntegrationTest::test_model_7b_logits_bf16
``` 

Testing for `mistral or mixtral or olmo or persimmon or phi or phi3 or qwen2 or qwen2_moe or stablelm or starcoder2 or whisper` now.

edit: a few other tests are failing with these archs, but nothing more from this PR. I think it is safe to merge.