| 1 | #!/usr/bin/env bash |
| 2 | set -euxo pipefail |
| 3 | |
| 4 | if [[ $# -ne 2 ]]; then |
| 5 | echo "Usage: $0 <wheel-directory> <checksum-file>" >&2 |
| 6 | exit 1 |
| 7 | fi |
| 8 | |
| 9 | TRITON_DIST="$1" |
| 10 | TRITON_CHECKSUM="$2" |
| 11 | |
| 12 | test -d "${TRITON_DIST}" |
| 13 | test -f "${TRITON_CHECKSUM}" |
| 14 | (cd "${TRITON_DIST}" && sha256sum -c "${TRITON_CHECKSUM}") |
| 15 | |
| 16 | shopt -s nullglob |
| 17 | TRITON_WHEELS=("${TRITON_DIST}"/triton-*.whl) |
| 18 | if [[ ${#TRITON_WHEELS[@]} -ne 1 ]]; then |
| 19 | echo "Expected exactly one Triton wheel in ${TRITON_DIST}" >&2 |
| 20 | exit 1 |
| 21 | fi |
| 22 | |
| 23 | uv pip install --system --reinstall --no-deps "${TRITON_WHEELS[0]}" |
| 24 | |
| 25 | if [[ -n "${TRITON_SYSTEM_PTXAS_PATH:-}" \ |
| 26 | || -n "${TRITON_EXPECTED_PTXAS_VERSION:-}" ]]; then |
| 27 | : "${TRITON_SYSTEM_PTXAS_PATH:?TRITON_SYSTEM_PTXAS_PATH must be set}" |
| 28 | : "${TRITON_EXPECTED_PTXAS_VERSION:?TRITON_EXPECTED_PTXAS_VERSION must be set}" |
| 29 | PTXAS_PATH="${TRITON_SYSTEM_PTXAS_PATH}" |
| 30 | test -x "${PTXAS_PATH}" |
| 31 | TRITON_NVIDIA_BIN=$(python3 -c \ |
| 32 | 'from pathlib import Path; import triton; print(Path(triton.__file__).parent / "backends" / "nvidia" / "bin")') |
| 33 | mkdir -p "${TRITON_NVIDIA_BIN}" |
| 34 | ln -sfn "${PTXAS_PATH}" \ |
| 35 | "${TRITON_NVIDIA_BIN}/ptxas-blackwell" |
| 36 | |
| 37 | env -u TRITON_PTXAS_BLACKWELL_PATH \ |
| 38 | TRITON_EXPECTED_PTXAS_PATH="${PTXAS_PATH}" \ |
| 39 | python3 -c \ |
| 40 | 'import os; from triton.backends.nvidia.compiler import get_ptxas; tool = get_ptxas(107); assert tool.version == os.environ["TRITON_EXPECTED_PTXAS_VERSION"]; assert os.path.samefile(tool.path, os.environ["TRITON_EXPECTED_PTXAS_PATH"])' |
| 41 | fi |