Skip to content

Commit 833fd92

Browse files
[AUTOGENERATED] [release/2.7] Change triton package name depending on rocm version (#2528)
Cherry-pick of #2518 Co-authored-by: Ethan Wee <[email protected]>
1 parent 38278cd commit 833fd92

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
@@ -97,8 +97,12 @@ def build_triton(
9797
triton_pythondir = triton_basedir / "python"
9898
triton_repo = "https://github.com/openai/triton"
9999
if device == "rocm":
100-
triton_pkg_name = "pytorch-triton-rocm"
101100
triton_repo = "https://github.com/ROCm/triton"
101+
rocm_version = get_rocm_version() # e.g., "7.0.1"
102+
if tuple(map(int, rocm_version.split("."))) > (7, 0, 0):
103+
triton_pkg_name = "triton"
104+
else:
105+
triton_pkg_name = "pytorch-triton-rocm"
102106
elif device == "xpu":
103107
triton_pkg_name = "pytorch-triton-xpu"
104108
triton_repo = "https://github.com/intel/intel-xpu-backend-for-triton"

0 commit comments

Comments
 (0)