Skip to content

Commit 4859b4a

Browse files
committed
Add opt-in nvtriton install/activate scripts to helion Dockerfile
Replace the direct nvtriton install with opt-in scripts at /opt/triton_tileir/. Users can install/activate/deactivate/uninstall at will without affecting the base OSS triton environment.
1 parent 5f29b3a commit 4859b4a

4 files changed

Lines changed: 112 additions & 7 deletions

File tree

docker/INSTALL_NVTRITON_TILEIR.md

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# nvtriton — Triton with TileIR Backend
2+
3+
Install nvtriton alongside your existing Triton. OSS Triton is never modified.
4+
5+
## Quick Start
6+
7+
```bash
8+
bash install_nvtriton.sh # install
9+
source ~/.local/triton_tileir/activate.sh # activate
10+
source ~/.local/triton_tileir/deactivate.sh # deactivate
11+
bash uninstall_nvtriton.sh # uninstall
12+
```
13+
14+
## Custom Install Path
15+
16+
```bash
17+
bash install_nvtriton.sh /my/custom/path
18+
source /my/custom/path/activate.sh
19+
source /my/custom/path/deactivate.sh
20+
bash uninstall_nvtriton.sh /my/custom/path
21+
```
22+
23+
> Deactivate before uninstalling — the script will remind you if you forget.

docker/helion.Dockerfile

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -50,13 +50,14 @@ RUN sudo uv pip install --system \
5050
torch \
5151
--index-url https://download.pytorch.org/whl/cu130
5252

53-
# nvtriton (Triton with TileIR backend — replaces upstream triton)
54-
RUN curl -L -o /tmp/nvtriton-3.6.0-cp313-cp313-linux_x86_64.whl \
55-
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 \
56-
&& sudo uv pip install --system /tmp/nvtriton-3.6.0-cp313-cp313-linux_x86_64.whl \
57-
&& rm /tmp/nvtriton-3.6.0-cp313-cp313-linux_x86_64.whl
58-
59-
ENV ENABLE_TILE=0
53+
# nvtriton (Triton with TileIR backend) — opt-in install/uninstall scripts.
54+
# bash /opt/triton_tileir/install.sh # install
55+
# source ~/.local/triton_tileir/activate.sh # activate
56+
# source ~/.local/triton_tileir/deactivate.sh # deactivate
57+
# bash /opt/triton_tileir/uninstall.sh # uninstall
58+
COPY docker/install_nvtriton.sh /opt/triton_tileir/install.sh
59+
COPY docker/uninstall_nvtriton.sh /opt/triton_tileir/uninstall.sh
60+
COPY docker/INSTALL_NVTRITON_TILEIR.md /opt/triton_tileir/README.md
6061

6162
# Helion
6263
RUN sudo uv pip install --system helion

docker/install_nvtriton.sh

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
#!/usr/bin/env bash
2+
# Install nvtriton (Triton with TileIR backend) alongside OSS triton.
3+
# OSS triton is untouched.
4+
#
5+
# Usage:
6+
# bash install_nvtriton.sh # installs to ~/.local/triton_tileir
7+
# bash install_nvtriton.sh /my/path # installs to /my/path
8+
#
9+
# After install, activate with:
10+
# source <install_dir>/activate.sh
11+
set -euo pipefail
12+
13+
RELEASE_URL="https://github.com/triton-lang/Triton-to-tile-IR/releases/download/v3.6.0-rc1"
14+
INSTALL_DIR="${1:-${HOME}/.local/triton_tileir}"
15+
16+
PY_TAG=$(python3 -c "import sys; print(f'cp{sys.version_info.major}{sys.version_info.minor}')")
17+
WHEEL="nvtriton-3.6.0-${PY_TAG}-${PY_TAG}-linux_x86_64.whl"
18+
19+
echo "==> Detected Python ${PY_TAG}"
20+
echo "==> Install directory: ${INSTALL_DIR}"
21+
22+
echo "==> Downloading ${WHEEL}..."
23+
curl -fSL -o "/tmp/${WHEEL}" "${RELEASE_URL}/${WHEEL}"
24+
25+
echo "==> Installing to ${INSTALL_DIR} (OSS triton is untouched)..."
26+
mkdir -p "${INSTALL_DIR}"
27+
python3 -m pip install --no-cache-dir --no-deps --target "${INSTALL_DIR}" "/tmp/${WHEEL}"
28+
rm -f "/tmp/${WHEEL}"
29+
30+
# Generate activate.sh
31+
cat > "${INSTALL_DIR}/activate.sh" <<EOF
32+
# Source this file to enable nvtriton TileIR backend.
33+
# source ${INSTALL_DIR}/activate.sh
34+
export PYTHONPATH="${INSTALL_DIR}\${PYTHONPATH:+:\$PYTHONPATH}"
35+
export ENABLE_TILE=1
36+
export HELION_BACKEND=tileir
37+
echo "nvtriton activated."
38+
EOF
39+
40+
# Generate deactivate.sh
41+
cat > "${INSTALL_DIR}/deactivate.sh" <<EOF
42+
# Source this file to revert to OSS triton.
43+
# source ${INSTALL_DIR}/deactivate.sh
44+
if [ -n "\${PYTHONPATH:-}" ]; then
45+
PYTHONPATH=\$(echo "\${PYTHONPATH}" | tr ':' '\n' | grep -v "^${INSTALL_DIR}\\\$" | paste -sd ':' || true)
46+
[ -z "\${PYTHONPATH}" ] && unset PYTHONPATH || export PYTHONPATH
47+
fi
48+
unset ENABLE_TILE
49+
unset HELION_BACKEND
50+
echo "nvtriton deactivated. OSS triton is now active."
51+
EOF
52+
53+
echo ""
54+
echo "Done! To activate: source ${INSTALL_DIR}/activate.sh"
55+
echo " To deactivate: source ${INSTALL_DIR}/deactivate.sh"
56+
echo " To uninstall: bash uninstall_nvtriton.sh ${INSTALL_DIR}"

docker/uninstall_nvtriton.sh

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
#!/usr/bin/env bash
2+
# Uninstall nvtriton completely.
3+
#
4+
# Usage:
5+
# bash uninstall_nvtriton.sh # removes ~/.local/triton_tileir
6+
# bash uninstall_nvtriton.sh /my/path # removes /my/path
7+
set -euo pipefail
8+
9+
INSTALL_DIR="${1:-${HOME}/.local/triton_tileir}"
10+
11+
# Check if nvtriton is still active in current shell
12+
if [ -n "${ENABLE_TILE:-}" ] || [ -n "${HELION_BACKEND:-}" ] || echo "${PYTHONPATH:-}" | grep -q "${INSTALL_DIR}"; then
13+
echo "Error: nvtriton is still active. Please deactivate first:"
14+
echo " source ${INSTALL_DIR}/deactivate.sh"
15+
exit 1
16+
fi
17+
18+
if [ -d "${INSTALL_DIR}" ]; then
19+
rm -rf "${INSTALL_DIR}"
20+
echo "==> Removed ${INSTALL_DIR}"
21+
else
22+
echo "==> ${INSTALL_DIR} not found (already clean)"
23+
fi
24+
25+
echo "Done!"

0 commit comments

Comments
 (0)