diff --git a/docker/INSTALL_NVTRITON_TILEIR.md b/docker/INSTALL_NVTRITON_TILEIR.md new file mode 100644 index 000000000..82afa86a9 --- /dev/null +++ b/docker/INSTALL_NVTRITON_TILEIR.md @@ -0,0 +1,23 @@ +# nvtriton — Triton with TileIR Backend + +Install nvtriton alongside your existing Triton. OSS Triton is never modified. + +## Quick Start + +```bash +bash install_nvtriton.sh # install +source ~/.local/triton_tileir/activate.sh # activate +source ~/.local/triton_tileir/deactivate.sh # deactivate +bash uninstall_nvtriton.sh # uninstall +``` + +## Custom Install Path + +```bash +bash install_nvtriton.sh /my/custom/path +source /my/custom/path/activate.sh +source /my/custom/path/deactivate.sh +bash uninstall_nvtriton.sh /my/custom/path +``` + +> Deactivate before uninstalling — the script will remind you if you forget. diff --git a/docker/helion.Dockerfile b/docker/helion.Dockerfile index 5a9f3664f..03c0772ee 100644 --- a/docker/helion.Dockerfile +++ b/docker/helion.Dockerfile @@ -50,13 +50,14 @@ RUN sudo uv pip install --system \ torch \ --index-url https://download.pytorch.org/whl/cu130 -# nvtriton (Triton with TileIR backend — replaces upstream triton) -RUN curl -L -o /tmp/nvtriton-3.6.0-cp313-cp313-linux_x86_64.whl \ - https://github.com/triton-lang/Triton-to-tile-IR/releases/download/v3.6.0-rc1/nvtriton-3.6.0-cp313-cp313-linux_x86_64.whl \ - && sudo uv pip install --system /tmp/nvtriton-3.6.0-cp313-cp313-linux_x86_64.whl \ - && rm /tmp/nvtriton-3.6.0-cp313-cp313-linux_x86_64.whl - -ENV ENABLE_TILE=0 +# nvtriton (Triton with TileIR backend) — opt-in install/uninstall scripts. +# bash /opt/triton_tileir/install.sh # install +# source ~/.local/triton_tileir/activate.sh # activate +# source ~/.local/triton_tileir/deactivate.sh # deactivate +# bash /opt/triton_tileir/uninstall.sh # uninstall +COPY docker/install_nvtriton.sh /opt/triton_tileir/install.sh +COPY docker/uninstall_nvtriton.sh /opt/triton_tileir/uninstall.sh +COPY docker/INSTALL_NVTRITON_TILEIR.md /opt/triton_tileir/README.md # Helion RUN sudo uv pip install --system helion diff --git a/docker/install_nvtriton.sh b/docker/install_nvtriton.sh new file mode 100755 index 000000000..4ef2dee44 --- /dev/null +++ b/docker/install_nvtriton.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# Install nvtriton (Triton with TileIR backend) alongside OSS triton. +# OSS triton is untouched. +# +# Usage: +# bash install_nvtriton.sh # installs to ~/.local/triton_tileir +# bash install_nvtriton.sh /my/path # installs to /my/path +# +# After install, activate with: +# source /activate.sh +set -euo pipefail + +RELEASE_URL="https://github.com/triton-lang/Triton-to-tile-IR/releases/download/v3.6.0-rc1" +INSTALL_DIR="${1:-${HOME}/.local/triton_tileir}" + +PY_TAG=$(python3 -c "import sys; print(f'cp{sys.version_info.major}{sys.version_info.minor}')") +WHEEL="nvtriton-3.6.0-${PY_TAG}-${PY_TAG}-linux_x86_64.whl" + +echo "==> Detected Python ${PY_TAG}" +echo "==> Install directory: ${INSTALL_DIR}" + +echo "==> Downloading ${WHEEL}..." +curl -fSL -o "/tmp/${WHEEL}" "${RELEASE_URL}/${WHEEL}" + +echo "==> Installing to ${INSTALL_DIR} (OSS triton is untouched)..." +mkdir -p "${INSTALL_DIR}" +python3 -m pip install --no-cache-dir --no-deps --target "${INSTALL_DIR}" "/tmp/${WHEEL}" +rm -f "/tmp/${WHEEL}" + +# Generate activate.sh +cat > "${INSTALL_DIR}/activate.sh" < "${INSTALL_DIR}/deactivate.sh" < Removed ${INSTALL_DIR}" +else + echo "==> ${INSTALL_DIR} not found (already clean)" +fi + +echo "Done!"