FA does not support head_dim > 64 on Ada (#328)

* FA does not support head_dim > 64 on Ada
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add cc8.7 to no FA list
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

transformer_engine/pytorch/attention.py

@@ -879,7 +879,7 @@ 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 == 8.6 and key_layer.shape[-1] > 64)
+            or (self.device_compute_capability in (8.6, 8.7, 8.9) and key_layer.shape[-1] > 64)
         ):
             use_flash_attention = False