Update triton wheel install script with new versioning (#141497)

This PR is a follow-on to #141410.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/141497
Approved by: https://github.com/huydhn
This commit is contained in:
Benjamin Glass
2024-11-25 18:14:13 +00:00
committed by PyTorch MergeBot
parent 91f7c547ec
commit 78491d6afc

View File

@ -10,7 +10,7 @@ if [[ -z "${USE_XPU}" ]]; then
if [[ "$BRANCH" =~ .*release.* ]]; then if [[ "$BRANCH" =~ .*release.* ]]; then
pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/test/ $TRITON_VERSION pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/test/ $TRITON_VERSION
else else
pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ $TRITON_VERSION+$(head -c 10 .ci/docker/ci_commit_pins/triton.txt) pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ $TRITON_VERSION+git$(head -c 8 .ci/docker/ci_commit_pins/triton.txt)
fi fi
else else
# The Triton xpu logic is as follows: # The Triton xpu logic is as follows:
@ -19,7 +19,7 @@ else
# it will install Triton from the source. # it will install Triton from the source.
TRITON_VERSION="pytorch-triton-xpu==$(cat .ci/docker/triton_version.txt)" TRITON_VERSION="pytorch-triton-xpu==$(cat .ci/docker/triton_version.txt)"
TRITON_XPU_COMMIT_ID="$(head -c 10 .ci/docker/ci_commit_pins/triton-xpu.txt)" TRITON_XPU_COMMIT_ID="git$(head -c 8 .ci/docker/ci_commit_pins/triton-xpu.txt)"
if [[ -z "${TRITON_XPU_BUILD_FROM_SOURCE}" ]]; then if [[ -z "${TRITON_XPU_BUILD_FROM_SOURCE}" ]]; then
pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ ${TRITON_VERSION}+${TRITON_XPU_COMMIT_ID} pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ ${TRITON_VERSION}+${TRITON_XPU_COMMIT_ID}
else else