From ddeb9fcc2d500cf9939ba167a0da171407aaeb5d Mon Sep 17 00:00:00 2001 From: Ethan Wee Date: Fri, 15 Aug 2025 13:05:24 -0700 Subject: [PATCH] [release/2.6] Change triton package name depending on rocm version (#2518) Relates to https://github.com/ROCm/builder/pull/91 Release 7.0: http://rocm-ci.amd.com/job/pytorch2.6-manylinux-wheels_rel-7.0/40/ `pytorch_triton_rocm-3.2.0+rocm7.0.0.git20943800-cp313-cp313-linux_x86_64.whl pushed to http://compute-artifactory.amd.com/artifactory/compute-pytorch-rocm/compute-rocm-rel-7.0/6/triton_main_conditionalize/pytorch_triton_rocm-3.2.0+rocm7.0.0.git20943800-cp313-cp313-linux_x86_64.whl` Mainline: http://rocm-ci.amd.com/job/mainline-pytorch2.6-manylinux-wheels/244/ `triton-3.2.0+rocm7.1.0.git20943800-cp313-cp313-linux_x86_64.whl pushed to http://compute-artifactory.amd.com/artifactory/compute-pytorch-rocm/compute-rocm-dkms-no-npi-hipclang/16516/triton_main_conditionalize/triton-3.2.0+rocm7.1.0.git20943800-cp313-cp313-linux_x86_64.whl` --- .github/scripts/build_triton_wheel.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/scripts/build_triton_wheel.py b/.github/scripts/build_triton_wheel.py index 695b4a9c865a6..ec26d8cbefaf9 100644 --- a/.github/scripts/build_triton_wheel.py +++ b/.github/scripts/build_triton_wheel.py @@ -101,8 +101,12 @@ def build_triton( triton_repo = "https://github.com/openai/triton" if device == "rocm": - triton_pkg_name = "pytorch-triton-rocm" triton_repo = "https://github.com/ROCm/triton" + rocm_version = get_rocm_version() # e.g., "7.0.1" + if tuple(map(int, rocm_version.split("."))) > (7, 0, 0): + triton_pkg_name = "triton" + else: + triton_pkg_name = "pytorch-triton-rocm" elif device == "xpu": triton_pkg_name = "pytorch-triton-xpu" triton_repo = "https://github.com/intel/intel-xpu-backend-for-triton"