1#!/bin/bash 2# Updates Triton to the pinned version for this copy of PyTorch 3BRANCH=$(git rev-parse --abbrev-ref HEAD) 4DOWNLOAD_PYTORCH_ORG="https://download.pytorch.org/whl" 5 6if [[ -z "${USE_XPU}" ]]; then 7 # Default install from PyTorch source 8 9 TRITON_VERSION="pytorch-triton==$(cat .ci/docker/triton_version.txt)" 10 if [[ "$BRANCH" =~ .*release.* ]]; then 11 pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/test/ $TRITON_VERSION 12 else 13 pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ $TRITON_VERSION+$(head -c 10 .ci/docker/ci_commit_pins/triton.txt) 14 fi 15else 16 # The Triton xpu logic is as follows: 17 # 1. By default, install pre-built whls. 18 # 2. [Not exposed to user] If the user set `TRITON_XPU_BUILD_FROM_SOURCE=1` flag, 19 # it will install Triton from the source. 20 21 TRITON_VERSION="pytorch-triton-xpu==$(cat .ci/docker/triton_version.txt)" 22 TRITON_XPU_COMMIT_ID="$(head -c 10 .ci/docker/ci_commit_pins/triton-xpu.txt)" 23 if [[ -z "${TRITON_XPU_BUILD_FROM_SOURCE}" ]]; then 24 pip install --index-url ${DOWNLOAD_PYTORCH_ORG}/nightly/ ${TRITON_VERSION}+${TRITON_XPU_COMMIT_ID} 25 else 26 TRITON_XPU_REPO="https://github.com/intel/intel-xpu-backend-for-triton" 27 28 # force-reinstall to ensure the pinned version is installed 29 pip install --force-reinstall "git+${TRITON_XPU_REPO}@${TRITON_XPU_COMMIT_ID}#subdirectory=python" 30 fi 31fi 32