diff --git a/.github/scripts/build_aiter_triton.sh b/.github/scripts/build_aiter_triton.sh new file mode 100755 index 0000000000..5ad68cbffa --- /dev/null +++ b/.github/scripts/build_aiter_triton.sh @@ -0,0 +1,37 @@ +#!/bin/bash + +set -ex + +echo +echo "==== ROCm Packages Installed ====" +dpkg -l | grep rocm || echo "No ROCm packages found." + +echo +echo "==== Install dependencies and aiter ====" +pip install --upgrade pandas zmq einops numpy==1.26.2 +pip uninstall -y aiter || true +pip install --upgrade "pybind11>=3.0.1" +pip install --upgrade "ninja>=1.11.1" +python3 setup.py develop + +# Read BUILD_TRITON env var, default to 1. If 1, install Triton; if 0, skip installation. +BUILD_TRITON=${BUILD_TRITON:-1} + +if [[ "$BUILD_TRITON" == "1" ]]; then + echo + echo "==== Install triton ====" + pip uninstall -y triton || true + git clone --depth=1 https://github.com/triton-lang/triton || true + cd triton + pip install -r python/requirements.txt + pip install filecheck + MAX_JOBS=64 pip --retries=5 install . + cd .. +else + echo + echo "[SKIP] Triton installation skipped because BUILD_TRITON=$BUILD_TRITON" +fi + +echo +echo "==== Show installed packages ====" +pip list diff --git a/.github/scripts/build_triton.sh b/.github/scripts/build_triton.sh deleted file mode 100755 index 730d81273a..0000000000 --- a/.github/scripts/build_triton.sh +++ /dev/null @@ -1,27 +0,0 @@ -#!/bin/bash - -set -ex - -echo -echo "==== ROCm Packages Installed ====" -dpkg -l | grep rocm || echo "No ROCm packages found." - -echo -echo "==== Install dependencies and aiter ====" -pip install --upgrade pandas zmq einops numpy==1.26.2 -pip uninstall -y aiter || true -pip install --upgrade "pybind11>=3.0.1" -python3 setup.py develop - -echo -echo "==== Install triton ====" -pip uninstall -y triton || true -git clone --depth=1 https://github.com/triton-lang/triton || true -cd triton -pip install -r python/requirements.txt -pip install filecheck -MAX_JOBS=64 pip install . - -echo -echo "==== Show installed packages ====" -pip list diff --git a/.github/workflows/aiter-test.yaml b/.github/workflows/aiter-test.yaml index f010eab246..52bc515414 100644 --- a/.github/workflows/aiter-test.yaml +++ b/.github/workflows/aiter-test.yaml @@ -104,14 +104,14 @@ jobs: --name aiter_test \ ${{ env.DOCKER_IMAGE }} - - name: Setup-Triton + - name: Setup Aiter run: | set -ex - echo "Setting up Triton..." + echo "Setting up Aiter..." docker exec \ -w /workspace \ aiter_test \ - ./.github/scripts/build_triton.sh + bash -c "BUILD_TRITON=0 ./.github/scripts/build_aiter_triton.sh" - name: Tests run: | @@ -190,7 +190,7 @@ jobs: docker exec \ -w /workspace \ aiter_test \ - ./.github/scripts/build_triton.sh + bash -c "BUILD_TRITON=0 ./.github/scripts/build_aiter_triton.sh" - name: Tests run: | diff --git a/.github/workflows/operators-tuning.yaml b/.github/workflows/operators-tuning.yaml index a9a75add45..aae4f26d42 100644 --- a/.github/workflows/operators-tuning.yaml +++ b/.github/workflows/operators-tuning.yaml @@ -54,14 +54,14 @@ jobs: --name operators_tuning_test \ rocm/pytorch:latest - - name: Setup-Triton + - name: Setup Aiter and Triton run: | set -ex - echo "Setting up Triton..." + echo "Setting up Aiter and Triton..." docker exec \ -w /workspace \ operators_tuning_test \ - ./.github/scripts/build_triton.sh + bash -c "BUILD_TRITON=0 ./.github/scripts/build_aiter_triton.sh" - name: Show Computing Units run: | diff --git a/.github/workflows/triton-test.yaml b/.github/workflows/triton-test.yaml index 41b9bb29b7..784953fba0 100644 --- a/.github/workflows/triton-test.yaml +++ b/.github/workflows/triton-test.yaml @@ -65,14 +65,14 @@ jobs: triton_test \ bash -c "pip install speedtest-cli && speedtest-cli --simple" || true - - name: Setup-Triton + - name: Setup Aiter and Triton run: | set -ex - echo "Setuping Triton..." + echo "Setting up Aiter and Triton..." docker exec \ -w /workspace \ triton_test \ - ./.github/scripts/build_triton.sh + ./.github/scripts/build_aiter_triton.sh - name: Install Pytest run: |