Plan B. Only add nv-triton build scripts to helion Dockerfile

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.

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

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 <install_dir>/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" <<EOF
+# Source this file to enable nvtriton TileIR backend.
+#   source ${INSTALL_DIR}/activate.sh
+export PYTHONPATH="${INSTALL_DIR}\${PYTHONPATH:+:\$PYTHONPATH}"
+export ENABLE_TILE=1
+export HELION_BACKEND=tileir
+echo "nvtriton activated."
+EOF
+
+# Generate deactivate.sh
+cat > "${INSTALL_DIR}/deactivate.sh" <<EOF
+# Source this file to revert to OSS triton.
+#   source ${INSTALL_DIR}/deactivate.sh
+if [ -n "\${PYTHONPATH:-}" ]; then
+    PYTHONPATH=\$(echo "\${PYTHONPATH}" | tr ':' '\n' | grep -v "^${INSTALL_DIR}\\\$" | paste -sd ':' || true)
+    [ -z "\${PYTHONPATH}" ] && unset PYTHONPATH || export PYTHONPATH
+fi
+unset ENABLE_TILE
+unset HELION_BACKEND
+echo "nvtriton deactivated. OSS triton is now active."
+EOF
+
+echo ""
+echo "Done! To activate:   source ${INSTALL_DIR}/activate.sh"
+echo "      To deactivate: source ${INSTALL_DIR}/deactivate.sh"
+echo "      To uninstall:  bash uninstall_nvtriton.sh ${INSTALL_DIR}"

docker/uninstall_nvtriton.sh

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+# Uninstall nvtriton completely.
+#
+# Usage:
+#   bash uninstall_nvtriton.sh              # removes ~/.local/triton_tileir
+#   bash uninstall_nvtriton.sh /my/path     # removes /my/path
+set -euo pipefail
+
+INSTALL_DIR="${1:-${HOME}/.local/triton_tileir}"
+
+# Check if nvtriton is still active in current shell
+if [ -n "${ENABLE_TILE:-}" ] || [ -n "${HELION_BACKEND:-}" ] || echo "${PYTHONPATH:-}" | grep -q "${INSTALL_DIR}"; then
+    echo "Error: nvtriton is still active. Please deactivate first:"
+    echo "  source ${INSTALL_DIR}/deactivate.sh"
+    exit 1
+fi
+
+if [ -d "${INSTALL_DIR}" ]; then
+    rm -rf "${INSTALL_DIR}"
+    echo "==> Removed ${INSTALL_DIR}"
+else
+    echo "==> ${INSTALL_DIR} not found (already clean)"
+fi
+
+echo "Done!"