Skip to content

Commit 0499e9d

Browse files
[rocm7.0_internal_testing] Change triton package name depending on rocm version (#2634)
Cherry-pick of #2518 Fixes https://ontrack-internal.amd.com/browse/SWDEV-554527 Co-authored-by: Divin Honnappa <[email protected]>
1 parent f096a8f commit 0499e9d

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

.github/scripts/build_triton_wheel.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,8 +101,12 @@ def build_triton(
101101
triton_pythondir = triton_basedir / "python"
102102
triton_repo = "https://github.com/openai/triton"
103103
if device == "rocm":
104-
triton_pkg_name = "pytorch-triton-rocm"
105104
triton_repo = "https://github.com/ROCm/triton"
105+
rocm_version = get_rocm_version() # e.g., "7.0.1"
106+
if tuple(map(int, rocm_version.split("."))) > (7, 0, 0):
107+
triton_pkg_name = "triton"
108+
else:
109+
triton_pkg_name = "pytorch-triton-rocm"
106110
elif device == "xpu":
107111
triton_pkg_name = "pytorch-triton-xpu"
108112
triton_repo = "https://github.com/intel/intel-xpu-backend-for-triton"

0 commit comments

Comments
 (0)