transformers

[`FA`] Fix head dim manipulation for more special models

#49862

Open