Version constraint of flash-attn needs to be updated

[lucifer1004]

Version: latest stable

Currently, the version constraint for flash-attn is:

TransformerEngine/setup.py

Line 269 in b8eea8a

add_unique(install_reqs, ["torch", "flash-attn>=2.0.6,<=2.4.2,!=2.0.9,!=2.1.0"])

So most likely v2.4.2 is going to be installed. However, this version seems to have some issues when imported, claiming some missing symbols. I uninstalled it and manually installed the latest v2.5.5, and the import succeeded.

The question is: why do we set this upper bound for flash-attn?

[timmoon10]

Flash Attention is being rapidly developed and its API is somewhat unstable. We've found it safer to only bump the version constraint after validating that Flash Attention works as expected. We are open to improving our workflows though.

Related: #519

Pinging @cyanguwa.

[ptrendx]

We updated flash attention version to 2.5.8 (current latest).

[ianporada]

Is it possible to look into updating max to flash-attn version v2.6.3? The compile time of this version is much faster for me than the current max v2.5.8.

Edit: Never mind, realized I was looking at stable branch and main is updated to 2.6.3