Use AOTriton 0.9.2b gfx950 specialized package stored in compute-artifactory (#50)

* Use AOTriton 0.9.2b gfx950 specialized package stored in compute-artifactory

* Enable gfx950 SDPA UTs.

* Bump AOTriton 0.9.2b gfx950 specalized version to move gfx950 out of experimental.

(cherry picked from commit 67a9fb25f1664a185c01abb9195e9b798ed1f54a)

Author: xinyazhang

cmake/External/aotriton.cmake

@@ -22,22 +22,16 @@ if(NOT __AOTRITON_INCLUDED)
   # Replaces .ci/docker/aotriton_version.txt
   # Note packages information may have versions skipped (due to no ABI breaks)
   # But they must be listed from lower version to higher version
-  set(__AOTRITON_VER "0.9.2b")
+  set(__AOTRITON_VER "0.9.2b_607fcc616f")
   set(__AOTRITON_MANYLINUX_LIST
-      "manylinux_2_28"  # rocm6.2
-      "manylinux_2_28"  # rocm6.3
-      "manylinux_2_28"  # rocm6.4
+      "manylinux_2_35"  # rocm6.5 on Ubuntu 22.04
       )
   set(__AOTRITON_ROCM_LIST
-      "rocm6.2"
-      "rocm6.3"
-      "rocm6.4"
+      "rocm6.5"
       )
-  set(__AOTRITON_CI_COMMIT "b388d223d8c7213545603e00f6f3148c54d1f525")
+  set(__AOTRITON_CI_COMMIT "607fcc616ff6836c2fcc9ad0fd37e528191d36fe")
   set(__AOTRITON_SHA256_LIST
-      "08d84f96f4c984179f80f517c0431c7511ee26bb0ce9bd05a827573ddd78cc79"  # rocm6.2
-      "9094d59717e7e6eace9126ca100dd0e86510f07fc6c3a349569fc4e2d9056604"  # rocm6.3
-      "41190202c2736d5ff75b13a3abc0fb52ebfbb67226cf85dc3de7699c7000db44"  # rocm6.4
+      "6318ccd92c4eb943241ba5c4c72f07f9ceec4f26e3913d3cf346d9c98cfd4d22"  # rocm6.5
       )
   set(__AOTRITON_Z "gz")
 
@@ -98,8 +92,8 @@ if(NOT __AOTRITON_INCLUDED)
                                   "${__AOTRITON_VER}-${__AOTRITON_MANYLINUX}"
                                   "_${__AOTRITON_ARCH}-rocm${__AOTRITON_ROCM}"
                                   "-shared.tar.${__AOTRITON_Z}")
-    string(CONCAT __AOTRITON_URL "https://github.com/ROCm/aotriton/releases/download/"
-                                 "${__AOTRITON_VER}/${__AOTRITON_FILE}")
+    string(CONCAT __AOTRITON_URL "https://compute-artifactory.amd.com/artifactory/rocm-generic-local/aotriton/"
+                                 "${__AOTRITON_FILE}")
     ExternalProject_Add(aotriton_external
       URL "${__AOTRITON_URL}"
       URL_HASH SHA256=${__AOTRITON_SHA256}

torch/testing/_internal/common_cuda.py

@@ -49,23 +49,23 @@ def evaluate_gfx_arch_within(arch_list):
     return any(arch in effective_arch for arch in arch_list)
 
 def CDNA2OrLater():
-    return evaluate_gfx_arch_within(["gfx90a", "gfx942"])
+    return evaluate_gfx_arch_within(["gfx90a", "gfx942", "gfx950"])
 
 def evaluate_platform_supports_flash_attention():
     if TEST_WITH_ROCM:
-        arch_list = ["gfx90a", "gfx942", "gfx1100"]
+        arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx950"]
         if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "0") != "0":
-            arch_list += ["gfx1201", "gfx950"]
+            arch_list += ["gfx1201"]
         return evaluate_gfx_arch_within(arch_list)
     if TEST_CUDA:
         return not IS_WINDOWS and SM80OrLater
     return False
 
 def evaluate_platform_supports_efficient_attention():
     if TEST_WITH_ROCM:
-        arch_list = ["gfx90a", "gfx942", "gfx1100"]
+        arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx950"]
         if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "0") != "0":
-            arch_list += ["gfx1201", "gfx950"]
+            arch_list += ["gfx1201"]
         return evaluate_gfx_arch_within(arch_list)
     if TEST_CUDA:
         return True