diff --git a/.ci/docker/ci_commit_pins/triton.txt b/.ci/docker/ci_commit_pins/triton.txt
index f0849cc7d8f6..877ce758db35 100644
--- a/.ci/docker/ci_commit_pins/triton.txt
+++ b/.ci/docker/ci_commit_pins/triton.txt
@@ -1 +1 @@
-4ed888920c5a0871957f1cf912e557bc79fbe56c
+0f380657dbf3ee86eb57558ff71df24f03b5d4e7
diff --git a/.ci/docker/common/install_triton.sh b/.ci/docker/common/install_triton.sh
index b2fdebdcc474..752d68976cc6 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 fa5f81bf06d9..437cb800c002 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"
@@ -111,14 +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/{ver}.{rev}.x"], cwd=triton_basedir
+                ["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