From 4747380a4bad9094c4d6dc252d1aa12a070c7935 Mon Sep 17 00:00:00 2001
From: Jack Taylor <jack.taylor@amd.com>
Date: Tue, 2 Jun 2026 05:34:11 -0500
Subject: [PATCH 1/2] Pin Triton to ROCm internal 3.6 branch

Update the Triton pin to the tip of ROCm release/internal/3.6.x and route
non-XPU Triton wheel builds through the ROCm fork.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .ci/docker/ci_commit_pins/triton.txt  | 2 +-
 .github/scripts/build_triton_wheel.py | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.ci/docker/ci_commit_pins/triton.txt b/.ci/docker/ci_commit_pins/triton.txt
index f0849cc7d8f63..877ce758db352 100644
--- a/.ci/docker/ci_commit_pins/triton.txt
+++ b/.ci/docker/ci_commit_pins/triton.txt
@@ -1 +1 @@
-4ed888920c5a0871957f1cf912e557bc79fbe56c
+0f380657dbf3ee86eb57558ff71df24f03b5d4e7
diff --git a/.github/scripts/build_triton_wheel.py b/.github/scripts/build_triton_wheel.py
index fa5f81bf06d9b..d1da9f871fd62 100644
--- a/.github/scripts/build_triton_wheel.py
+++ b/.github/scripts/build_triton_wheel.py
@@ -101,9 +101,8 @@ def build_triton(
         triton_basedir = Path(tmpdir) / "triton"
         triton_pythondir = triton_basedir / "python"
 
-        triton_repo = "https://github.com/openai/triton"
+        triton_repo = "https://github.com/ROCm/triton"
         if device == "rocm":
-            triton_repo = "https://github.com/ROCm/triton"
             triton_pkg_name = "triton"
         elif device == "xpu":
             triton_pkg_name = "triton-xpu"
@@ -114,7 +113,8 @@ def build_triton(
         if release:
             ver, rev, patch = version.split(".")
             check_call(
-                ["git", "checkout", f"release/{ver}.{rev}.x"], cwd=triton_basedir
+                ["git", "checkout", f"release/internal/{ver}.{rev}.x"],
+                cwd=triton_basedir,
             )
         else:
             check_call(["git", "fetch", "origin", commit_hash], cwd=triton_basedir)

From e48db0ec35a483b468167964030953b0028039f5 Mon Sep 17 00:00:00 2001
From: Jack Taylor <jack.taylor@amd.com>
Date: Tue, 2 Jun 2026 04:41:37 -0500
Subject: [PATCH 2/2] Build Triton from internal branch in CI

Use release/internal/3.7.x for non-XPU Triton source builds in Docker and
wheel CI instead of checking out the triton.txt commit.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .ci/docker/common/install_triton.sh   |  8 ++++++--
 .github/scripts/build_triton_wheel.py | 13 ++++++++++---
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/.ci/docker/common/install_triton.sh b/.ci/docker/common/install_triton.sh
index b2fdebdcc4747..752d68976cc6c 100755
--- a/.ci/docker/common/install_triton.sh
+++ b/.ci/docker/common/install_triton.sh
@@ -26,7 +26,11 @@ else
 fi
 
 # The logic here is copied from .ci/pytorch/common_utils.sh
-TRITON_PINNED_COMMIT=$(get_pinned_commit ${TRITON_TEXT_FILE})
+if [ "${TRITON_TEXT_FILE}" = "triton" ]; then
+  TRITON_REF="release/internal/3.6.x"
+else
+  TRITON_REF=$(get_pinned_commit ${TRITON_TEXT_FILE})
+fi
 
 if [ -n "${UBUNTU_VERSION}" ];then
     apt update
@@ -49,7 +53,7 @@ pushd /var/lib/jenkins/
 
 as_jenkins git clone --recursive ${TRITON_REPO} triton
 cd triton
-as_jenkins git checkout ${TRITON_PINNED_COMMIT}
+as_jenkins git checkout ${TRITON_REF}
 as_jenkins git submodule update --init --recursive
 
 # Old versions of python have setup.py in ./python; newer versions have it in ./
diff --git a/.github/scripts/build_triton_wheel.py b/.github/scripts/build_triton_wheel.py
index d1da9f871fd62..437cb800c0023 100644
--- a/.github/scripts/build_triton_wheel.py
+++ b/.github/scripts/build_triton_wheel.py
@@ -110,15 +110,22 @@ def build_triton(
         else:
             triton_pkg_name = "triton"
         check_call(["git", "clone", triton_repo, "triton"], cwd=tmpdir)
-        if release:
+        if device == "xpu" and release:
             ver, rev, patch = version.split(".")
+            # XPU uses the patch version in the release branch name
             check_call(
-                ["git", "checkout", f"release/internal/{ver}.{rev}.x"],
+                ["git", "checkout", f"release/{ver}.{rev}.{patch}"],
                 cwd=triton_basedir,
             )
-        else:
+        elif device == "xpu":
             check_call(["git", "fetch", "origin", commit_hash], cwd=triton_basedir)
             check_call(["git", "checkout", commit_hash], cwd=triton_basedir)
+        else:
+            ver, rev, _patch = version.split(".")
+            check_call(
+                ["git", "checkout", f"release/internal/{ver}.{rev}.x"],
+                cwd=triton_basedir,
+            )
 
         # change built wheel name and version
         env["TRITON_WHEEL_NAME"] = triton_pkg_name