Repository: huggingface/transformers
Pull request: update flex attention to use `return_aux` instead of `return_lse` when torch verison >= 2.9

…

# What does this PR do?
In torch versions >= 2.9.0, it requests the lse from flex_attenetion using `AuxRequest` instead of the deprecated `return_lse`, which triggers a warning and can break tracing.

Fixes #44683


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?
@vasqu @ArthurZucker @Cyrilvallez 


Changed files:
filename: src/transformers/integrations/flex_attention.py
patch: @@ -26,19 +26,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Union
+from typing import Optional, Union
 
 import torch
 from packaging import version
 
 from ..utils import is_torch_flex_attn_available, logging
-from ..utils.import_utils import get_torch_version, is_torch_less_or_equal, is_torchdynamo_compiling
+from ..utils.import_utils import (
+    get_torch_version,
+    is_torch_greater_or_equal,
+    is_torch_less_or_equal,
+    is_torchdynamo_compiling,
+)
+
+
+_TORCH_FLEX_USE_AUX = is_torch_greater_or_equal("2.9.0")
 
 
 if is_torch_flex_attn_available():
     from torch.nn.attention.flex_attention import _DEFAULT_SPARSE_BLOCK_SIZE as flex_default_block_size
     from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention
 
+    if _TORCH_FLEX_USE_AUX:
+        from torch.nn.attention.flex_attention import AuxRequest
+    else:
+        AuxRequest = None
+
 
 logger = logging.get_logger(__name__)
 
@@ -84,6 +97,20 @@ def __call__(self):
         return self._compiled_flex_attention
 
 
+def get_flex_attention_lse_kwargs(return_lse: bool) -> dict[str, bool | Optional["AuxRequest"]]:
+    """
+    Requests the LSE from flex_attention in a version-agnostic fashion.
+
+    Before torch 2.9, the LSE was requested via the boolean return_lse field. However, starting with
+    torch 2.9, an AuxRequest object must be passed via the aux_request field. This method conditionally
+    returns the correct form based on the python version.
+    """
+    if _TORCH_FLEX_USE_AUX:
+        return {"return_aux": AuxRequest(lse=True) if return_lse else None}
+
+    return {"return_lse": return_lse}
+
+
 def compile_friendly_flex_attention(
     query: torch.Tensor,
     key: torch.Tensor,
@@ -298,12 +325,21 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx):
         kernel_options=kernel_options,
         # Last time checked on PyTorch == 2.5.1: Flex Attention always computes the lse regardless.
         # For simplification, we thus always return it as no additional computations are introduced.
-        return_lse=return_lse,
         training=module.training,
+        # inject the lse args
+        **get_flex_attention_lse_kwargs(return_lse),
     )
-    # lse is returned in float32
+
     if return_lse:
-        attention_output, lse = flex_attention_output  # type: ignore[misc]
+        # before torch 2.9, return_lse returns the LSE directly as a second tuple element
+        # in torch 2.9 and later, return_aux returns AuxOutput as a second tuple element -- the LSE must be extracted
+        if _TORCH_FLEX_USE_AUX:
+            attention_output, aux = flex_attention_output  # type: ignore[misc]
+            lse = aux.lse
+        else:
+            attention_output, lse = flex_attention_output  # type: ignore[misc]
+
+        # lse is returned in float32
         lse = lse.to(value.dtype)
 
         if s_aux is not None: