Use custom SDPA for decoder-only HF Transformers

[guangy10]

• Register a custom SDPA attention to HF Transformers
• Use the custom SDPA for decoder-only text models (custom SDPA is optimized for ExecuTorch)
• Have to bump up transformers version to latest release (4.51.0) in order to use the AttentionInterface. This has been addressed in Bump Transformers verion #47
• optimum-cli export executorch supports custom SDPA

3x speedup using custom SDPA for HF smollm2 (XNNPACK fp32):

General applicable to all causal LMs. For encoder-decoder models, it may apply to the self attention layer in the decoder, can make an experiment in a follow-up PR.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[guangy10]

Will gate the access to custom_sdpa in 0.4.0 release.

[guangy10]

@larryliu0820 FYI, I'm running into the same the issue when using the custom sdpa with eager. The exported model is fine. https://github.com/huggingface/optimum-executorch/actions/runs/14324546808/job/40147624383?pr=46#step:5:137. Checked the shape of the attn output, and they are same. Will dig further by inspecting if output tensors are close tomorrow.

Updated: The issue is because the eager was using DynamicCache by default, and our custom SDPA can ONLY work with the StaticCache. After setting the cache_implementation correctly, both the eager and the ExecuTorch model can work as expected.

[guangy10]

Transformers version bump has been merged in #47

[guangy10]

cc: @larryliu0820 @kimishpatel for review

[larryliu0820]

Can you change this: https://github.com/huggingface/optimum-executorch/blob/main/optimum/executorch/modeling.py#L181 to use the new Python API: https://pytorch.org/executorch/stable/index.html

[guangy10]

Can you change this: https://github.com/huggingface/optimum-executorch/blob/main/optimum/executorch/modeling.py#L181 to use the new Python API: https://pytorch.org/executorch/stable/index.html

@larryliu0820 Yeah, I'm going to do it in a separate PR.

[guangy10]

Rebased and fixed conflicts

[guangy10]

@larryliu0820 @kimishpatel good to merge?

[guangy10]

Support export with custom_sdpa using optimum-cli export executorch