Enable test_attention

Signed-off-by: wenjh <wenjh@sugon.com>

tests/pytorch/attention/test_attention.py

@@ -50,6 +50,7 @@ from transformer_engine.pytorch.quantized_tensor import (
     prepare_for_saving,
     restore_from_saved,
 )
+from torch.utils.cpp_extension import IS_HIP_EXTENSION
 
 _current_file = pathlib.Path(__file__).resolve()
 sys.path.append(str(_current_file.parent.parent))
@@ -65,11 +66,12 @@ from utils import (
 fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True)
 fp8_attn_available, reason_for_no_fp8_attn = fp8_available, reason_for_no_fp8
 device_compute_capability = get_device_compute_capability()
-if fp8_available and (device_compute_capability < (9, 0) or device_compute_capability >= (12, 0)):
-    fp8_attn_available = False
-    reason_for_no_fp8_attn = (
-        "FP8 attention is not supported for compute capability ="
-        f" sm{device_compute_capability[0] * 10 + device_compute_capability[1]}"
+if not IS_HIP_EXTENSION:
+    if fp8_available and (device_compute_capability < (9, 0) or device_compute_capability >= (12, 0)):
+        fp8_attn_available = False
+        reason_for_no_fp8_attn = (
+            "FP8 attention is not supported for compute capability ="
+            f" sm{device_compute_capability[0] * 10 + device_compute_capability[1]}"
     )
 
 # Reset RNG seed and states