Relax FA 2.0 checks for Ada (#331)

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

transformer_engine/pytorch/attention.py

@@ -885,10 +885,15 @@ class DotProductAttention(torch.nn.Module):
         if (query_layer.dtype not in [torch.bfloat16, torch.float16]
             or key_layer.dtype not in [torch.bfloat16, torch.float16]
             or value_layer.dtype not in [torch.bfloat16, torch.float16]
-            or (self.device_compute_capability in (8.6, 8.7, 8.9) and key_layer.shape[-1] > 64)
         ):
             use_flash_attention = False
 
+        if key_layer.shape[-1] > 64:
+            if self.device_compute_capability in (8.6, 8.7):
+                use_flash_attention = False
+            elif not _flash_attn_2_available and self.device_compute_capability == 8.9:
+                use_flash_attention = False
+
         if self.attn_mask_type == "padding" and attention_mask is not None:
             use_flash_attention = False
             use_fused_attention = False