[RFC]: Deprecate the FlexAttention backend

[yewentao256]

Motivation.

FlexAttention is no longer selected by default on ROCm and ranks behind TritonAttention on CUDA.

It has a lot of test coverage and is not used in practice, and sometimes breaks CI like #46531

TritonAttention now covers the standard FlexAttention use cases, including FP32, encoder-only, non-causal, sliding window, multimodal prefix, batch invariance, and Unlimited-OCR R-SWA.

Removing FlexAttention would eliminate roughly 1,900 lines of backend-specific implementation and tests, while reducing reliance on private PyTorch FlexAttention APIs.

Proposed Change.

• Now: Mark FLEX_ATTENTION as deprecated.
• v0.29.0/v0.30.0: Remove the backend implementation, enum entry, Flex-specific configuration fields, and tests.

Feedback Period.

No response

CC List.

@liangel-02 @Isotr0py @MatthewBonanni @mgoin @LucasWilkinson @andakai

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[zou3519]

cc @drisspg

[AndreasKaratzas]

Flex attention had good potential here to be a batch invariant backend. It was cross platform and accurate. I know that it is not widely used anymore so if it is creating CI problems then perhaps we should indeed deprecate it. I haven't captured flex attention problems tbh on CI.
cc @khluu

[yewentao256]

Thanks @AndreasKaratzas. Note that for batch invariance it is not a must have as well, we have triton attention or FA as more reliable backends.

[drisspg]

🫡

[mgoin]

Are there examples of flaky CI other than batch invariance? We can just stop advertising flex attention as supporting that feature. I haven't seen other issues?

[MatthewBonanni]

I haven't really run into too many CI issues with it; I'm not necessarily opposed to removing it but I don't feel like it's been much of a maintenance burden

[yewentao256]

@mgoin @MatthewBonanni Here are some issues related

• #27724：fixed by #27853
• #27344：Bumping to PyTorch 2.9 ，Flex _build_block_mask_direct returns the wrong shape
• #41342：AMD pooling nightly with FLEX_ATTENTION case failures
• And of course [CI Flaky test]: Flaky batch invariant test with FLEX_ATTENTION #46531

[MatthewBonanni]

I mean, 2 of those were from last October and #41342 was just an issue with the test itself. Like I said, not opposed to deprecating, I just don't think CI is the reason necessarily

[yewentao256]

Thanks @MatthewBonanni ! There are some issues with some feature migration as well like v1 to v2 #39353.
Yes CI is not the only reason, the main thing is that it might not worth the effort to maintain a backend that is seldom use in practice by most of the users. And that's why I raise this RFC, hoping to get more feedbacks from community, eg. some users might say they use it a lot.

[liangel-02]

the main thing is that it might not worth the effort to maintain a backend that is seldom use in practice by most of the users

hey @yewentao256, torchtitan depends on vLLM's flex attention to support the generator side of their RL loop (see example config here) and it has been very helpful for us to have an implementation that integrates with vLLM out of the box!

[Isotr0py]

FlexAttention provides a flexible interface to draft/invesitgate sparse attention mask support. For example, Gemma3 bidi attn and Unlimited-OCR R-SWA support are both firstly implemented on FlexAttention backend, then migrated to Triton/FA.

I think a more critical question is that whether there will be a sparse mask implementation, which is difficult to be implemented on FA/Triton but not for FlexAttn.

[aoshen02]

I think it does provide a flexible way to config customized attention, especially some people in RL use it to do some validation. Maybe we should introduce magi attention to replace it if we want to remove this.