Bump up FA to 2.8.1 (#1949)

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

qa/L3_pytorch_FA_versions_test/test.sh

@@ -18,10 +18,10 @@ sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); pri
 export FLASH_ATTN_CUDA_ARCHS=$sm_arch
 if [ $sm_arch -gt 90 ]
 then
-  FA_versions=(2.7.3)
+  FA_versions=(2.8.1)
 elif [ $sm_arch -eq 90 ]
 then
-  FA_versions=(2.5.7 2.7.3 3.0.0b1)
+  FA_versions=(2.7.3 2.8.1 3.0.0b1)
 fi
 
 for fa_version in "${FA_versions[@]}"

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -105,7 +105,7 @@ class FlashAttentionUtils:
     version = PkgVersion("0")
     version_required = PkgVersion("2.1.1")
     version_required_blackwell = PkgVersion("2.7.3")
-    max_version = PkgVersion("2.7.4.post1")
+    max_version = PkgVersion("2.8.1")
     v2_plus = False
     v2_1_plus = False
     v2_3_plus = False