[Triton/Gluon] Add config-aware repr to the attention kernels #10188
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: Triton Test | |
| on: | |
| push: | |
| branches: [main] | |
| pull_request: | |
| # labeled: re-run when adding ci:triton-300x to start extra MI300X jobs | |
| types: [opened, synchronize, reopened, ready_for_review, labeled] | |
| branches: [main] | |
| paths: | |
| - "aiter/ops/triton/**" | |
| - "op_tests/triton_tests/**" | |
| - "op_tests/op_benchmarks/triton/**" | |
| - ".github/scripts/build_aiter_triton.sh" | |
| - ".github/scripts/download_triton_wheel.sh" | |
| - ".github/scripts/install_triton.sh" | |
| - ".github/scripts/select_triton_tests.py" | |
| - ".github/scripts/split_tests.sh" | |
| - ".github/scripts/verify_triton_pin.py" | |
| - ".github/requirements/triton-test.txt" | |
| # Docs-only changes (README.md, configs/CLAUDE.md, conv/DESIGN.md, ...) | |
| # don't affect test results — skip the test suites when a PR touches | |
| # nothing else. | |
| - "!**/*.md" | |
| - '!docs/**' | |
| - '!LICENSE' | |
| - '!.gitignore' | |
| - '!.github/workflows/**' | |
| - '!.github/scripts/sglang_downstream.py' | |
| # Re-include Triton workflow files after the broad workflow exclusion. | |
| - ".github/workflows/triton-test.yaml" | |
| - ".github/workflows/prepare-triton-wheel.yaml" | |
| - ".github/workflows/ci-config.yaml" | |
| workflow_dispatch: | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.ref }} | |
| cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} | |
| jobs: | |
| ci_config: | |
| if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} | |
| uses: ./.github/workflows/ci-config.yaml | |
| check-signal: | |
| if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} | |
| runs-on: ubuntu-latest | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| - name: Wait for Checks workflow | |
| run: ./.github/scripts/check_signal.sh | |
| env: | |
| GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} | |
| GITHUB_SHA: ${{ github.sha }} | |
| # Step 1: split triton tests into 8 shards, output triton_shard_0.list ... triton_shard_7.list | |
| split_triton_tests: | |
| if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} | |
| runs-on: ubuntu-latest | |
| needs: [check-signal] | |
| outputs: | |
| shard_count: 8 | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| - name: Split Triton Tests (8 shards) | |
| run: ./.github/scripts/split_tests.sh --shards 8 --test-type triton | |
| - name: Upload test shard lists as artifact | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: triton_shards | |
| path: triton_shard_*.list | |
| retention-days: 7 | |
| prepare_triton_wheel: | |
| if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} | |
| needs: [check-signal, ci_config] | |
| uses: ./.github/workflows/prepare-triton-wheel.yaml | |
| with: | |
| docker-image: ${{ needs.ci_config.outputs.pytorch_py312_image }} | |
| artifact-name: ${{ needs.ci_config.outputs.triton_wheel_artifact_name }} | |
| retention-days: 3 | |
| # Step 2: MI35X matrix jobs | |
| triton: | |
| if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} | |
| name: Triton Tests (MI35X) / Shard ${{ matrix.shard }} | |
| runs-on: linux-aiter-mi35x-1 | |
| timeout-minutes: 60 | |
| needs: [split_triton_tests, prepare_triton_wheel, check-signal, ci_config] | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| shard: [0, 1, 2, 3, 4, 5, 6, 7] | |
| env: | |
| DOCKER_IMAGE: ${{ needs.ci_config.outputs.pytorch_py312_image }} | |
| TRITON_WHEEL_ARTIFACT_NAME: ${{ needs.ci_config.outputs.triton_wheel_artifact_name }} | |
| TRITON_TEST: "op_tests/triton_tests/" | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| with: | |
| fetch-depth: 1 | |
| submodules: 'recursive' | |
| - name: Network preflight before Triton artifact downloads | |
| if: ${{ startsWith(runner.name, 'linux-aiter-mi35x-1-') }} | |
| continue-on-error: true | |
| timeout-minutes: 5 | |
| run: | | |
| set +e | |
| echo "=== Runner info ===" | |
| echo "RUNNER_NAME=${RUNNER_NAME}" | |
| echo "RUNNER_OS=${RUNNER_OS}" | |
| echo "hostname=$(hostname -f 2>/dev/null || hostname)" | |
| date -u | |
| echo "=== DNS check ===" | |
| getent hosts github.com || true | |
| getent hosts api.github.com || true | |
| getent hosts productionresultssa4.blob.core.windows.net || true | |
| getent hosts productionresultssa8.blob.core.windows.net || true | |
| echo "=== GitHub and Azure Blob connectivity check ===" | |
| for url in \ | |
| https://api.github.com/rate_limit \ | |
| https://github.com/ROCm/aiter \ | |
| https://productionresultssa4.blob.core.windows.net/ \ | |
| https://productionresultssa8.blob.core.windows.net/ | |
| do | |
| curl -L -o /dev/null -sS \ | |
| --connect-timeout 10 \ | |
| --max-time 30 \ | |
| -w "url=${url} code=%{http_code} dns=%{time_namelookup}s connect=%{time_connect}s tls=%{time_appconnect}s ttfb=%{time_starttransfer}s total=%{time_total}s speed=%{speed_download}B/s\n" \ | |
| "$url" || true | |
| done | |
| - name: Download test shard lists | |
| id: download_triton_shards | |
| uses: actions/download-artifact@v4 | |
| timeout-minutes: 15 | |
| with: | |
| name: triton_shards | |
| - name: Download Triton wheel artifact | |
| id: download_triton_wheel | |
| uses: actions/download-artifact@v4 | |
| timeout-minutes: 30 | |
| continue-on-error: true | |
| with: | |
| name: ${{ env.TRITON_WHEEL_ARTIFACT_NAME }} | |
| path: triton_wheels | |
| - name: Speedtest after Triton artifact download failure | |
| if: ${{ always() && startsWith(runner.name, 'linux-aiter-mi35x-1-') && (steps.download_triton_shards.outcome == 'failure' || steps.download_triton_wheel.outcome == 'failure') }} | |
| continue-on-error: true | |
| timeout-minutes: 5 | |
| run: | | |
| set +e | |
| echo "=== Runner info ===" | |
| echo "RUNNER_NAME=${RUNNER_NAME}" | |
| echo "RUNNER_OS=${RUNNER_OS}" | |
| echo "hostname=$(hostname -f 2>/dev/null || hostname)" | |
| date -u | |
| echo "=== speedtest-cli ===" | |
| python3 -m pip install --user --disable-pip-version-check speedtest-cli || true | |
| speedtest_bin="$(python3 -m site --user-base)/bin/speedtest-cli" | |
| timeout 120 "$speedtest_bin" --simple --secure || true | |
| - name: List test shard files | |
| run: | | |
| ls -l triton_shard_*.list | |
| - name: Docker login | |
| if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.head.repo.fork }} | |
| env: | |
| DOCKER_USERNAME: ${{ secrets[needs.ci_config.outputs.docker_username_secret] }} | |
| DOCKER_PASSWORD: ${{ secrets[needs.ci_config.outputs.docker_password_secret] }} | |
| run: | | |
| for attempt in 1 2 3; do | |
| if echo "$DOCKER_PASSWORD" | docker login -u "$DOCKER_USERNAME" --password-stdin; then | |
| echo "Docker login succeeded on attempt ${attempt}" | |
| exit 0 | |
| fi | |
| echo "Docker login attempt ${attempt} failed" | |
| if [ "${attempt}" != 3 ]; then | |
| sleep 10 | |
| fi | |
| done | |
| echo "Docker login failed after 3 attempts, continuing anyway" | |
| exit 0 | |
| - name: Export test file list for this shard as env | |
| id: set_shard_files | |
| run: | | |
| TRITON_TEST=$(cat triton_shard_${{ matrix.shard }}.list) | |
| echo "$TRITON_TEST" | |
| echo "TRITON_TEST=$TRITON_TEST" >> $GITHUB_ENV | |
| - name: Run the container | |
| run: | | |
| set -ex | |
| echo "Starting container: triton_test" | |
| if [ -f "/etc/podinfo/gha-render-devices" ]; then | |
| DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices) | |
| else | |
| DEVICE_FLAG="--device /dev/dri" | |
| fi | |
| docker run -dt \ | |
| --device=/dev/kfd $DEVICE_FLAG \ | |
| --shm-size=16G \ | |
| -e TRITON_WHEEL_DIR=/workspace/triton_wheels \ | |
| --group-add $(getent group render | cut -d: -f3) \ | |
| --group-add $(getent group video | cut -d: -f3) \ | |
| -v "${{ github.workspace }}:/workspace" \ | |
| -w /workspace \ | |
| --name triton_test \ | |
| ${{ env.DOCKER_IMAGE }} | |
| - name: Setup Aiter and Triton | |
| run: | | |
| set -ex | |
| echo "Setting up Aiter and Triton..." | |
| docker exec -w /workspace triton_test ./.github/scripts/build_aiter_triton.sh | |
| # TODO: Uncomment [docker exec -w /workspace triton_test cat "${ENV_FILE}" >> "${GITHUB_ENV}"] | |
| # command to enable test selection. | |
| # - name: Triton Test Selection Script | |
| # # main branch should always run the full test suite. | |
| # if: ${{ github.ref != 'refs/heads/main' }} | |
| # run: | | |
| # set -ex | |
| # git fetch --no-tags origin --depth=1 \ | |
| # "pull/${{ github.event.pull_request.number }}/merge:pr-merge" | |
| # git fetch --no-tags origin --depth=1 \ | |
| # "refs/heads/${{ github.event.pull_request.base.ref }}:target" | |
| # ENV_FILE=$(docker exec -w /workspace triton_test \ | |
| # mktemp /workspace/github_env.XXXXXXXXXX.tmp) | |
| # docker exec -w /workspace triton_test \ | |
| # git config --global --add safe.directory /workspace | |
| # docker exec -w /workspace triton_test \ | |
| # python .github/scripts/select_triton_tests.py \ | |
| # --source pr-merge --target target \ | |
| # --env-var TRITON_TEST --env-file "${ENV_FILE}" | |
| # # docker exec -w /workspace triton_test cat "${ENV_FILE}" >> "${GITHUB_ENV}" | |
| # docker exec -w /workspace triton_test rm "${ENV_FILE}" | |
| - name: Triton Tests | |
| run: | | |
| set -ex | |
| echo "Running Triton Tests..." | |
| docker exec -w /workspace triton_test mkdir -p test-reports | |
| # MI35X: skip the MHA-PE backward test that fails an accuracy check on gfx950 with the coming Triton release. | |
| docker exec -w /workspace triton_test pytest -v ${TRITON_TEST} \ | |
| --deselect "op_tests/triton_tests/attention/test_mha_with_pe.py::test_mha_backward_with_pe" \ | |
| --junitxml=test-reports/triton.xml | |
| - name: Upload test logs | |
| uses: actions/upload-artifact@v4 | |
| if: success() | |
| with: | |
| name: triton-test-shard-${{ matrix.shard }} | |
| path: test-reports/triton.xml | |
| retention-days: 7 | |
| - name: Cleanup container | |
| if: always() | |
| run: | | |
| docker rm -f triton_test || true | |
| # Step 2b: MI300X matrix jobs (opt-in via ci:triton-300x on PRs, always on main) | |
| triton-mi300x: | |
| if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') && (github.ref == 'refs/heads/main' || (github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'ci:triton-300x'))) }} | |
| name: Triton Tests (MI300X) / Shard ${{ matrix.shard }} | |
| runs-on: linux-aiter-oci-mi300x-1 | |
| timeout-minutes: 60 | |
| needs: [split_triton_tests, prepare_triton_wheel, check-signal, ci_config] | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| shard: [0, 1, 2, 3, 4, 5, 6, 7] | |
| env: | |
| DOCKER_IMAGE: ${{ needs.ci_config.outputs.pytorch_py312_image }} | |
| TRITON_WHEEL_ARTIFACT_NAME: ${{ needs.ci_config.outputs.triton_wheel_artifact_name }} | |
| TRITON_TEST: "op_tests/triton_tests/" | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| with: | |
| fetch-depth: 1 | |
| submodules: 'recursive' | |
| - name: Download test shard lists | |
| uses: actions/download-artifact@v4 | |
| timeout-minutes: 15 | |
| with: | |
| name: triton_shards | |
| - name: Download Triton wheel artifact | |
| id: download_triton_wheel | |
| uses: actions/download-artifact@v4 | |
| timeout-minutes: 30 | |
| continue-on-error: true | |
| with: | |
| name: ${{ env.TRITON_WHEEL_ARTIFACT_NAME }} | |
| path: triton_wheels | |
| - name: List test shard files | |
| run: | | |
| ls -l triton_shard_*.list | |
| - name: Docker login | |
| if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.head.repo.fork }} | |
| env: | |
| DOCKER_USERNAME: ${{ secrets[needs.ci_config.outputs.docker_username_secret] }} | |
| DOCKER_PASSWORD: ${{ secrets[needs.ci_config.outputs.docker_password_secret] }} | |
| run: | | |
| for attempt in 1 2 3; do | |
| if echo "$DOCKER_PASSWORD" | docker login -u "$DOCKER_USERNAME" --password-stdin; then | |
| echo "Docker login succeeded on attempt ${attempt}" | |
| exit 0 | |
| fi | |
| echo "Docker login attempt ${attempt} failed" | |
| if [ "${attempt}" != 3 ]; then | |
| sleep 10 | |
| fi | |
| done | |
| echo "Docker login failed after 3 attempts, continuing anyway" | |
| exit 0 | |
| - name: Export test file list for this shard as env | |
| id: set_shard_files | |
| run: | | |
| TRITON_TEST=$(cat triton_shard_${{ matrix.shard }}.list) | |
| echo "$TRITON_TEST" | |
| echo "TRITON_TEST=$TRITON_TEST" >> $GITHUB_ENV | |
| - name: Run the container | |
| run: | | |
| set -ex | |
| echo "Starting container: triton_test" | |
| if [ -f "/etc/podinfo/gha-render-devices" ]; then | |
| DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices) | |
| else | |
| DEVICE_FLAG="--device /dev/dri" | |
| fi | |
| docker run -dt \ | |
| --device=/dev/kfd $DEVICE_FLAG \ | |
| --shm-size=16G \ | |
| -e TRITON_WHEEL_DIR=/workspace/triton_wheels \ | |
| --group-add $(getent group render | cut -d: -f3) \ | |
| --group-add $(getent group video | cut -d: -f3) \ | |
| -v "${{ github.workspace }}:/workspace" \ | |
| -w /workspace \ | |
| --name triton_test \ | |
| ${{ env.DOCKER_IMAGE }} | |
| - name: Setup Aiter and Triton | |
| run: | | |
| set -ex | |
| echo "Setting up Aiter and Triton..." | |
| docker exec -w /workspace triton_test ./.github/scripts/build_aiter_triton.sh | |
| - name: Triton Tests | |
| run: | | |
| set -ex | |
| echo "Running Triton Tests..." | |
| docker exec -w /workspace triton_test mkdir -p test-reports | |
| docker exec -w /workspace triton_test pytest -v ${TRITON_TEST} --junitxml=test-reports/triton.xml | |
| - name: Upload test logs | |
| uses: actions/upload-artifact@v4 | |
| if: success() | |
| with: | |
| name: triton-test-mi300x-shard-${{ matrix.shard }} | |
| path: test-reports/triton.xml | |
| retention-days: 7 | |
| - name: Cleanup container | |
| if: always() | |
| run: | | |
| docker rm -f triton_test || true | |
| triton-mi300x-test-finish: | |
| if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') && (github.ref == 'refs/heads/main' || (github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'ci:triton-300x'))) }} | |
| name: Triton MI300X Test Results | |
| runs-on: ubuntu-latest | |
| needs: [triton-mi300x] | |
| steps: | |
| - name: Download all MI300X test reports | |
| uses: actions/download-artifact@v4 | |
| timeout-minutes: 15 | |
| with: | |
| pattern: triton-test-mi300x-shard-* | |
| path: . | |
| - name: Check Triton MI300X Test Results | |
| run: | | |
| set -ex | |
| echo "Checking Triton MI300X Test Results..." | |
| all_passed=true | |
| for shard in {0..7}; do | |
| if [ ! -f triton-test-mi300x-shard-${shard}/triton.xml ]; then | |
| echo "MI300X test report for shard ${shard} not found." | |
| all_passed=false | |
| break | |
| fi | |
| done | |
| if [ "$all_passed" = true ]; then | |
| echo "All MI300X tests passed." | |
| else | |
| echo "MI300X test failures or errors detected." | |
| exit 1 | |
| fi | |
| triton-test-finish: | |
| if: ${{ (!github.event.pull_request || !github.event.pull_request.draft) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} | |
| name: Triton Test Results | |
| runs-on: ubuntu-latest | |
| needs: [triton] | |
| steps: | |
| - name: Download all test reports | |
| uses: actions/download-artifact@v4 | |
| timeout-minutes: 15 | |
| with: | |
| pattern: triton-test-shard-* | |
| path: . | |
| - name: Check Triton Test Results | |
| run: | | |
| set -ex | |
| echo "Checking Triton Test Results..." | |
| all_passed=true | |
| for shard in {0..7}; do | |
| if [ ! -f triton-test-shard-${shard}/triton.xml ]; then | |
| echo "Test report for shard ${shard} not found." | |
| all_passed=false | |
| break | |
| fi | |
| done | |
| if [ "$all_passed" = true ]; then | |
| echo "All tests passed." | |
| else | |
| echo "Test failures or errors detected." | |
| exit 1 | |
| fi |