Pin triton-rocm to latest 2.3.1 commit (pytorch#126309)

* Revert "pin rocm" This reverts commit 45ebb10. * Revert "lint" This reverts commit 05860b9. * rocm_pin
ROCm · May 15, 2024 · 0365423 · 0365423
1 parent 03baf94
commit 0365423
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 12 deletions.
diff --git a/.ci/docker/ci_commit_pins/triton-rocm.txt b/.ci/docker/ci_commit_pins/triton-rocm.txt
@@ -1 +1 @@
-d08e16b738ab550c3af51305df624d5c823dc445
+c8ad905211f45e162102823149f0d7f2cfaa4418
diff --git a/.github/scripts/build_triton_wheel.py b/.github/scripts/build_triton_wheel.py
@@ -10,9 +10,6 @@
 SCRIPT_DIR = Path(__file__).parent
 REPO_DIR = SCRIPT_DIR.parent.parent
 
-# TODO: Remove me once Triton version is again in sync for vanilla and ROCm
-ROCM_TRITION_VERSION = "2.3.0"
-
 
 def read_triton_pin(rocm_hash: bool = False) -> str:
     triton_file = "triton.txt" if not rocm_hash else "triton-rocm.txt"
@@ -101,12 +98,9 @@ def build_triton(
         check_call(["git", "clone", triton_repo], cwd=tmpdir)
         if release:
             ver, rev, patch = version.split(".")
-            if build_rocm:
-                check_call(["git", "checkout", "release/2.3.x"], cwd=triton_basedir)
-            else:
-                check_call(
-                    ["git", "checkout", f"release/{ver}.{rev}.x"], cwd=triton_basedir
-                )
+            check_call(
+                ["git", "checkout", f"release/{ver}.{rev}.x"], cwd=triton_basedir
+            )
         else:
             check_call(["git", "checkout", commit_hash], cwd=triton_basedir)
 
@@ -165,7 +159,7 @@ def build_triton(
         patch_init_py(
             triton_pythondir / "triton" / "__init__.py",
             version=f"{version}",
-            expected_version=ROCM_TRITION_VERSION if build_rocm else None,
+            expected_version=None,
         )
 
         if build_rocm:
@@ -174,7 +168,7 @@ def build_triton(
                 triton_pythondir / "setup.py",
                 name=triton_pkg_name,
                 version=f"{version}",
-                expected_version=ROCM_TRITION_VERSION,
+                expected_version=None,
             )
             check_call("scripts/amd/setup_rocm_libs.sh", cwd=triton_basedir, shell=True)
             print("ROCm libraries setup for triton installation...")