Skip to content

Commit

Permalink
Add triton version for nightly and release (pytorch#1703)
Browse files Browse the repository at this point in the history
  • Loading branch information
atalman committed Feb 13, 2024
1 parent 80b420a commit 5314b20
Showing 1 changed file with 13 additions and 10 deletions.
23 changes: 13 additions & 10 deletions manywheel/build_cuda.sh
Original file line number Diff line number Diff line change
Expand Up @@ -262,17 +262,20 @@ else
exit 1
fi

# No triton dependency for now on 3.12 since we don't have binaries for it
# and torch.compile doesn't work.
if [[ $(uname) == "Linux" && "$DESIRED_PYTHON" != "3.12" ]]; then
TRITON_VERSION=$(cat $PYTORCH_ROOT/.ci/docker/triton_version.txt)
TRITON_REQUIREMENT="triton==${TRITON_VERSION}; platform_system == 'Linux' and platform_machine == 'x86_64'"

if [[ -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${TRITON_REQUIREMENT}"
else
export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | ${TRITON_REQUIREMENT}"
fi
TRITON_VERSION=$(cat $PYTORCH_ROOT/.ci/docker/triton_version.txt)
# Only linux Python < 3.12 are supported wheels for triton
TRITON_CONSTRAINT="platform_system == 'Linux' and platform_machine == 'x86_64' and python_version < '3.12'"
TRITON_REQUIREMENT="pytorch-triton==${TRITON_VERSION}; ${TRITON_CONSTRAINT}"
if [[ -n "$OVERRIDE_PACKAGE_VERSION" && "$OVERRIDE_PACKAGE_VERSION" =~ .*dev.* ]]; then
TRITON_SHORTHASH=$(cut -c1-10 $PYTORCH_ROOT/.github/ci_commit_pins/triton.txt)
TRITON_REQUIREMENT="pytorch-triton==${TRITON_VERSION}+${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
fi

if [[ -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${TRITON_REQUIREMENT}"
else
export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | ${TRITON_REQUIREMENT}"
fi

# builder/test.sh requires DESIRED_CUDA to know what tests to exclude
Expand Down

0 comments on commit 5314b20

Please sign in to comment.