[Gemma3] compile ✨ 
Author: gante
Pull request: yes

# What does this PR do?

Enables compilation on Gemma3 (and re-enables it on Gemma2 / Cohere2).

Reverts #36620 
Supercedes #37433 (solves the same problem, but this PR is much cleaner)

### Performance
Measured on an RTX4090, excluding compile warmup time:
- [Example script from gemma 3 4b](https://huggingface.co/google/gemma-3-4b-it#running-the-model-on-a-singlemulti-gpu):  3.87s on `main` (not compiled) -> 2.39s this PR
- [Example script from gemma 2 9b](https://huggingface.co/google/gemma-2-9b-it#running-the-model-on-a-single--multi-gpu): 3.67s on `main` (not compiled) -> 2.18s this PR

### Tests
- [x] slow gemma 2 tests (9 failing tests from `main` -> need to be revisited)
- [x] slow gemma 3 tests (2 failing tests from `main`, `tests/models/gemma3/test_modeling_gemma3.py::Gemma3Vision2TextModelTest::test_eager_matches_sdpa_generate` gets fixed in this PR)


### Post-mortem: How did we break compile on Gemma 2?
1. Doing `git bisect`, compilation first "breaks" in the PR where the cache is initialized in the `meta` device (https://github.com/huggingface/transformers/pull/35164). "break" here doesn't mean "crash", but rather "becomes very slow". Curiously, this change doesn't slow down `StaticCache` + `llama` (why?), so it flew under the radar when we benchmarked before merging. Nevertheless, this specific PR has been reverted (https://github.com/huggingface/transformers/pull/36543).
2. Along the way, we corrected how the sliding window attention works, by slicing the attention mask correctly (https://github.com/huggingface/transformers/pull/35681). However, the solution here is not `torch.compile` friendly: `forward` now has an `int` argument that is different at each forward pass at generation time, causing recompilation ([reference](https://pytorch.org/docs/stable/torch.compiler_troubleshooting.html#wrapping-constants-with-tensors)). The changes in this PR work around this issue.

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

> Do you think we can add a slow test to compare generation time with compile or include HybridCache in benchmarks board, so we don't accidentally introduce graph breaks? Given that Gemma3 is a high usage model and supports only Hybrid cache, I think it's important to not break it

@zucchini-nlp Benchmarks are hard to put as a test: different devices/versions -> different speeds, needs multiple runs to avoid being flaky 💔  However, we should test that compilation only happens once for the entire forward pass, i.e. that there are no graph breaks nor recompilations. I'm going to explore torch docs to see if we can add this to a test, in this PR. At the moment, we are often doubting the quality of our compiled forward passes, which is not great

@gante yeah, we can do something similar to what diffusers is trying to do (in continuation to yesterday's thread discussions). Fine by me, as long as we check that there are no recompilations every step

@Cyrilvallez off by 1 comment addressed 👍 LMK if you'd like any further changes

@Cyrilvallez @zucchini-nlp There are three follow-up items to this PR, where each will be a separate PR:
1. Mixin test update to confirm that no recompilations, dynamic shapes, ... are happening. [Related comment](https://github.com/huggingface/transformers/pull/37447/#issuecomment-2812272952). [I've decided not to include it in this PR, since it needs to touch other models to fix failures];
2. TP + compile compatible caches. [Related comment](https://github.com/huggingface/transformers/pull/37447/#discussion_r2039726257)
3. Gemma 3 + Prefill chunking + more than one new token at generation time. [Related comment](https://github.com/huggingface/transformers/pull/37447/#discussion_r2044016869)



Hey @gante, can you point me how gemma3 is compiled, maybe a simple test? I looked through this PR but don't see a test for compile. BTW I'm working on add `torch.export` support for this model, and some changes in this PR are helpful with the export work as you are removing some of the data-dependent code which is also required by `torch.export`.

@guangy10 the test upgraded in [this PR](https://github.com/huggingface/transformers/pull/37629) is the base test for compilation -- it confirms compiled outputs == uncompiled outputs, and that there are no graph breaks in the forward pass (when called through `generate`)

FYI, in https://github.com/huggingface/transformers/pull/37728 `torch.export` is supported as well! AND the model can be lowered to ExecuTorch and run e2e w/o any issue!