PyTorch FA test fix (#370)

* Corrected bugs

* Updated documentation

Author: Micky774

tests/pytorch/fused_attn/test_fused_attn.py

@@ -220,8 +220,8 @@ def test():
         )
         (
             use_flash_attention,
-            use_fused_attention,
             flash_attention_backend,
+            use_fused_attention,
             fused_attention_backend,
             use_unfused_attention,
             available_backends,
@@ -369,6 +369,7 @@ def test_dot_product_attention(
             and config.attn_mask_type in ["causal", "padding_causal"]
         )
         and (config.window_size[0] == -1 or FlashAttentionUtils.v2_3_plus)
+        and not is_mla
     ):
         flash_attn_supported = True
 

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -278,6 +278,8 @@ def get_attention_backend(
     ----------
     use_flash_attention: bool
         Whether the `FlashAttention` backend has been selected.
+    flash_attention_backend: PkgVersion
+        If `use_flash_attention = True`, the version of the selected `FlashAttention` backend.
     use_fused_attention: bool
         Whether the `FusedAttention` backend has been selected.
     fused_attention_backend: tex.NVTE_Fused_Attn_Backend