[Neuron] Compilation-friendly generation loop
Author: dacorvo
Pull request: no

## Context

`GenerationMixin._sample` grows `input_ids`, `attention_mask`, and `position_ids` via `torch.cat` on every decode step. This is problematic for the Neuron backend where dynamic tensor shapes carry a cost: static shapes are required for graph caching — dynamic shapes cause retracing, **even in eager mode**. 

## What works today

These four steps give `model.generate()` an **almost** stable, compile-friendly path on Neuron. None of them require upstream changes — this issue is about what's left after them.

1. **Use the Neuron compile backend.**
   ```py
   gen_config.compile_config = CompileConfig(backend="neuron")
   ```
2. **Build the `StaticCache` yourself, initialize it, pass it in.**
   ```py
   cache = StaticCache(config=model.config, max_cache_len=max_seq_len)
   cache.early_initialization(batch_size, num_heads, head_dim, dtype, device)
   ```
3. **Set `gen_config.prefill_chunk_size`** to the largest value the device fits.
4. **Pad the prompt to a multiple of `prefill_chunk_size`.**

Once these are in place, the remaining problem are the dynamic ops inside `_sample`.

## What was tried

We compared the results of a generation on Llama-3.1-8B with four states:
- **cold** — empty cache, short prompt, mnt=128.
- **warm** — same as cold; clean cache hit.
- **warm-diff-in** — warm cache, long prompt (2× chunk count).
- **warm-diff-mnt** — warm cache, short prompt, mnt=256.

We compared the current situation (vanilla) with two modified transformers versions:
- one that keeps input_ids and all associated tensors on CPU and moves them to device only before forward.
- one that uses static tensors in generate.

Each cell is `total_s / tps`. The 🟢 / 🟡 / 🔴 on the warm-* columns
rates wall-clock vs that mode's own warm baseline — 🟢 ≤ 1.5×, 🟡
1.5–10×, 🔴 > 10×. Cold cells aren't colored.

| mode           | cold          | warm                 | warm-diff-in            | warm-diff-mnt           |
|----------------|---------------|----------------------|-------------------------|-------------------------|
| vanilla        | 2694 s / 0.05 | 🟢 35.5 s / 3.60     | 🔴 2149 s / 0.06        | 🔴 2130 s / 0.12        |
| cpu            |  360 s / 0.35 | 🟢 30.9 s / 4.15     | 🟢 33.1 s / 3.87        | 🟢 38.5 s / 6.64        |
| static_tensors |  592 s / 0.22 | 🟢 21.6 s / 5.91     | 🟡  115 s / 1.11        | 🟢 28.0 s / 9.15        |

So it seems the "CPU" path is the happy path. This issue is to discuss how we could leverage that in `transformers` with the less disruption possible.

## Related

* Non-generation Neuron changes: [#44741](https://github.com/huggingface/transformers/issues/44741)
* Auto-StaticCache on Neuron device: [#44748](https://github.com/huggingface/transformers/issues/44748)


Comments:
Auto-compilation gate fix for Neuron: #44757

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.

@Cyrilvallez would love to have your feedback on this subject. Is the extension of the StaticCache capabilities a way to improve the generation loop efficiency when dealing with static shapes ?

cc @tengomucho 

Hey @dacorvo! Thanks for opening the issue. First I'd like to perfectly understand the issue: Neuron needs static shapes all the time if I understand correctly right? Because currently, we do compile only the `forward` pass, and only after prefill, so we have static shapes but indeed only inside the `forward` and during the decoding step (i.e. not in.-between `forward` calls).

In general, as you said the critical point is the mask. It would be relatively straightforward to use fixed-shape `input_ids` and `position_ids`, but the mask would require more work. `position_ids` cannot be retrieved from `cache.cumulative_length`, as the latter does not track padding tokens (however as I said, the `position_ids` could be tracked independently, so this would not be an issue in general).

What I'm most concerned about here is maintainability. This would add significant code paths only for the benefit of 1 type of hardware. The advantage of the current approach is that it's agnostic to cache/no-cache and static/non-static. We then only slice according to what we need.

I believe the best approach is to instead use `generate_batch`, where all shapes are fixed by design. It will also provides better perfs, and only supports a subset of `generate` that matches yours (as your approach could not be used for e.g. speculative decoding for example). cc @remi-or here for viz/help.

Let me know if that answers you @dacorvo!

@Cyrilvallez I will let @dacorvo answer, but i think that without these changes, working with transformers on hardware platforms that trigger compilation for every shape change such as Trainium or TPUs will be very challenging.
`generate_batch` does not solve all problems, and it seems to me that the design is more to allow continuous batching (for example for `transformers serve`).
 

It can be used for inference as well! What are the problems left if you switch to `generate_batch` ? 

@Cyrilvallez thanks for the answer.
First, an apology, by rereading the code I realize I missed the fact that `prepare_inputs_for_generation` is already cache-aware (it already creates a causal mask tensor of fixed shape).

So the only issue is actually the growing input_ids, position_ids and attention_mask tensors (and a few others for different model types).
Since they are on device, every operations triggers a recompile (the Neuron eager mode is fake: it is actually a per-op compilation).
For instance this leads to a recompilation:
``
input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
``
Even if we eventually pass `input_ids[-1] to the `forward.

So if we look at this only from a neuron perspective (purely hypothetical, not saying we should do that), either we keep them on CPU and pay the sync cost every time we decode (update 4D causal mask is one sync + sliced input_ids is another one, both will probably collapse), or we make them to use static shapes to at least avoid the recompilation.
I understand the latter is something you are reluctant to maintain, could we consider the other option ?

I updated the issue and the description on this narrower scope.

Regarding `generate_batch`, it looks like:
- it supports a lot less features than generate, 
- it has PagedAttention support as a prerequisite.

The neuron integration in transformers is not to provide the best inference performance (for that the target is rather vLLM), but to provide the widest set of models and features possible, so that it becomes appealing to machine learning engineer to prototype models and tools on Neuron.

Yes, `generate_batch` supports a lot less features than generate, but I wonder what kind of models you would like to see added? We already support all texts models, and this could be the occasion to broaden the support of some models. 
As for PagedAttention, it is not a requirement: you can use `eager`, `sdpa` or `flash`. It is more efficient to use Paged for decoding (a lot more actuallly) but not required at all. 
Big picture, it might be a good thing to change `generate_batch` to support "most" use cases, but it will surely be easier to add some small features to generate, despite the fact that they could only be useful for Neuron. I think it depends on what direction we want to follow in the future.

@remi-or I think you are right on the fact that `generate_batch` can be a faster API for inference when serving. I think I will look at that later.
For now I am trying to make the `generate` loop work better on TPU. In this context, I think that making the generation loop compiler-friendly for hardware systems that only handle static shape can be good. 

Sure! I just didn't want you to have to do the work twice!

> [@remi-or](https://github.com/remi-or) I think you are right on the fact that `generate_batch` can be a faster API for inference when serving. I think I will look at that later. For now I am trying to make the `generate` loop work better on TPU. In this context, I think that making the generation loop compiler-friendly for hardware systems that only handle static shape can be good.

This issue is about neuron and generate, which is the main customer-facing API that we must support for neuron users. @tengomucho I suggest you open a separate issue for TPU, and possibily for generate_batch if you find it relevant for your use case.

@Cyrilvallez I did more tests to find the best compromise I could to avoid as much as possible touching the generation code. It seems there is a path forward by keeping the problematic ops on CPU. Please take a look at the updated description and give me your feedback. I can share code if you want.

@Cyrilvallez I plan to submit this change:
https://github.com/huggingface/transformers/commit/2dab55c320d82e6d4c9496446b81472f0cd71647
Please let me know if that is OK for you.