Bump flash attention 2

[francesco-bertolotti]

Description

Bump the maximum supported FlashAttention version from v8.2.3 to v8.2.4.

Currently, Transformer Engine (TE) caps FlashAttention at v8.2.3. However, newer commits on the FlashAttention main branch include changes that would be useful to TE users, such as those introduced in Dao-AILab/flash-attention#2260.

Allowing FlashAttention v8.2.4 enables users to use TE with newer FlashAttention releases. I ran the test suite with the updated version, and the results were largely successful. The few failures appear to be related to my specific environment rather than the version bump itself.

Was there a specific reason for guarding against newer FlashAttention releases, or was it mainly intended to guard against potential API changes?

Changes

• Increase the maximum supported FlashAttention version from v8.2.3 to v8.2.4.
• Validate the change with the existing test suite.

[greptile-apps]

Greptile Summary

The PR expands Transformer Engine’s accepted FlashAttention 2 range through version 2.8.4.

• Raises the exclusive maximum supported version from 2.8.4 to 2.8.5.
• Updates version-boundary and supported-version-message expectations.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Raises the exclusive FlashAttention 2 upper bound so version 2.8.4 is accepted.
tests/pytorch/attention/test_attention.py	Updates predicate and diagnostic-message expectations for the new upper bound.

_{Reviews (2): Last reviewed commit: "bump FA" | Re-trigger Greptile}