File tree Expand file tree Collapse file tree 2 files changed +3
-3
lines changed Expand file tree Collapse file tree 2 files changed +3
-3
lines changed Original file line number Diff line number Diff line change 9393
9494# Set triton via PYTORCH_EXTRA_INSTALL_REQUIREMENTS for triton rocm package
9595if [[ "$PACKAGE_TYPE" =~ .*wheel.* && -n "$PYTORCH_BUILD_VERSION" && "$PYTORCH_BUILD_VERSION" =~ .*rocm.* && $(uname) == "Linux" ]]; then
96- TRITON_REQUIREMENT="pytorch- triton-rocm ==${TRITON_VERSION}; ${TRITON_CONSTRAINT}"
96+ TRITON_REQUIREMENT="triton==${TRITON_VERSION}; ${TRITON_CONSTRAINT}"
9797 if [[ -n "$PYTORCH_BUILD_VERSION" && "$PYTORCH_BUILD_VERSION" =~ .*dev.* ]]; then
9898 TRITON_SHORTHASH=$(cut -c1-8 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton.txt)
99- TRITON_REQUIREMENT="pytorch- triton-rocm ==${TRITON_VERSION}+git${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
99+ TRITON_REQUIREMENT="triton==${TRITON_VERSION}+git${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
100100 fi
101101 if [[ -z "${PYTORCH_EXTRA_INSTALL_REQUIREMENTS:-}" ]]; then
102102 export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${TRITON_REQUIREMENT}"
Original file line number Diff line number Diff line change @@ -105,7 +105,7 @@ def build_triton(
105105
106106 triton_repo = "https://github.com/openai/triton"
107107 if device == "rocm" :
108- triton_pkg_name = "pytorch- triton-rocm "
108+ triton_pkg_name = "triton"
109109 triton_repo = "https://github.com/ROCm/triton"
110110 elif device == "xpu" :
111111 triton_pkg_name = "pytorch-triton-xpu"
You can’t perform that action at this time.
0 commit comments