[BE] install_triton_wheel.sh update for internal dev (#154637)

internal devgpu gets mad at `pip install ...` but `python3 -m pip install ...` is fine
Pull Request resolved: https://github.com/pytorch/pytorch/pull/154637
Approved by: https://github.com/Skylion007, https://github.com/cyyever
This commit is contained in:
Nicolas Macchioni
2025-05-31 06:57:53 +00:00
committed by PyTorch MergeBot
parent 7368eeba5e
commit 15e9119a69

View File

@ -1,5 +1,7 @@
#!/bin/bash
# Updates Triton to the pinned version for this copy of PyTorch
PYTHON="python3"
PIP="$PYTHON -m pip"
BRANCH=$(git rev-parse --abbrev-ref HEAD)
DOWNLOAD_PYTORCH_ORG="https://download.pytorch.org/whl"
@ -8,9 +10,9 @@ if [[ -z "${USE_XPU}" ]]; then
TRITON_VERSION="pytorch-triton==$(cat .ci/docker/triton_version.txt)"
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
pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ $TRITON_VERSION+git$(head -c 8 .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
else
# The Triton xpu logic is as follows:
@ -21,11 +23,11 @@ else
TRITON_VERSION="pytorch-triton-xpu==$(cat .ci/docker/triton_version.txt)"
TRITON_XPU_COMMIT_ID="$(head -c 8 .ci/docker/ci_commit_pins/triton-xpu.txt)"
if [[ -z "${TRITON_XPU_BUILD_FROM_SOURCE}" ]]; then
pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ ${TRITON_VERSION}+git${TRITON_XPU_COMMIT_ID}
${PIP} install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ ${TRITON_VERSION}+git${TRITON_XPU_COMMIT_ID}
else
TRITON_XPU_REPO="https://github.com/intel/intel-xpu-backend-for-triton"
# force-reinstall to ensure the pinned version is installed
pip install --force-reinstall "git+${TRITON_XPU_REPO}@${TRITON_XPU_COMMIT_ID}#subdirectory=python"
${PIP} install --force-reinstall "git+${TRITON_XPU_REPO}@${TRITON_XPU_COMMIT_ID}#subdirectory=python"
fi
fi