Compute flash data layout info once and for all when possible

[fxmarty]

As per title.

Try to remediate this happening at each layer:

It helps a bit, but not as much as I hoped (only a few %). It seems that cudaMemcpyAsync that the profiler was showing was mostly an artifact, and is not as long as the profiler suggests, and most of the actual overhead is in index_first_axis and unpad_input.

Now we still have overhead given that the unpad/pad still happens at each layer.

Things that would help:

• Fuse the query/key/value nn.Linear and apply unpadding only once on them (would help a bit but not much)
• Even better: unpad only once at the beginning at the model, and repad only once at the output of the model. The would remove all the unpad/pad overhead, and on top of that would allow us to save some compute / communication in the various nn.Linear, avoiding to do compute on padded tokens.

I hardly see any of these two options fitting transformers philosophy. The second solution basically involves a complete rewrite of the modeling. The second solution would be very helpful though for training with padding, and for batched inference where the prefill is expensive (maybe code models).

In the current state, the benefits from huggingface#25598 seem too limited compared to what PyTorch 2.1 will provide (aka flash v2 without padding) to justify the integration in transformers. But I think there's a real topic of whether we want to offload efficient training / inference to external libs (TGI, brrr) and have transformers more as a "playground", or whether we are keen to include optimizations as such in transformers (even e.g. in the longer term upstreaming SDPA in transformers, using torch.nested when possible, etc.).

Or if we keep the BT API (that has its limitation and flaws) and rely on that to monkey patch / overwrite transformers code.

cc @younesbelkada @LysandreJik