TritonBench #523
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: TritonBench | |
| on: | |
| schedule: | |
| # Run every 12 hours | |
| - cron: '0 */12 * * *' | |
| workflow_dispatch: | |
| inputs: | |
| tritonbench_branch: | |
| description: TritonBench branch (main) | |
| required: true | |
| type: string | |
| default: main | |
| benchmarks: | |
| description: | | |
| A comma-separated list of benchmarks from tritonbench/benchmarks (optional, default to run nightly) | |
| required: false | |
| type: string | |
| runners: | |
| description: | | |
| A comma-separated list of runners from .github/scripts/genenerate_tritonbench_matrix.py to run the benchmark (optional, default to run b200) | |
| required: true | |
| type: string | |
| default: b200 | |
| triton_channel: | |
| description: | | |
| Triton channel to run: all (default, runs both triton-main and meta-triton), triton-main, or meta-triton | |
| required: false | |
| type: choice | |
| default: all | |
| options: | |
| - all | |
| - triton-main | |
| - meta-triton | |
| triton_repo: | |
| description: Triton repo to build from (optional) | |
| required: false | |
| type: string | |
| default: '' | |
| triton_commit: | |
| description: Triton commit or ref to build from (optional) | |
| required: false | |
| type: string | |
| default: '' | |
| benchmark_parameters: | |
| description: Extra parameters to pass to run-benchmark.sh (optional) | |
| required: false | |
| type: string | |
| default: '' | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} | |
| cancel-in-progress: true | |
| jobs: | |
| set-parameters: | |
| runs-on: ubuntu-latest | |
| outputs: | |
| benchmark_matrix: ${{ steps.set-parameters.outputs.benchmark_matrix }} | |
| steps: | |
| - name: Checkout repository | |
| uses: actions/checkout@v4 | |
| - uses: actions/setup-python@v5 | |
| with: | |
| python-version: '3.12' | |
| - name: Set parameters | |
| id: set-parameters | |
| shell: bash | |
| env: | |
| BENCHMARKS: ${{ inputs.benchmarks || '' }} | |
| RUNNERS: ${{ inputs.runners || '' }} | |
| TRITON_CHANNEL: ${{ inputs.triton_channel || 'all' }} | |
| run: | | |
| set -eux | |
| if [ "${TRITON_CHANNEL}" == "all" ]; then | |
| TRITON="triton-main,meta-triton" | |
| else | |
| TRITON="${TRITON_CHANNEL}" | |
| fi | |
| # The generated matrix is grouped by benchmark and runner | |
| python .github/scripts/generate_tritonbench_matrix.py \ | |
| --benchmarks "${BENCHMARKS}" \ | |
| --runners "${RUNNERS}" \ | |
| --triton "${TRITON}" | |
| benchmarks: | |
| name: Run TritonBench benchmarks | |
| needs: set-parameters | |
| if: ${{ !github.event.pull_request.head.repo.fork && github.repository_owner == 'pytorch' }} | |
| strategy: | |
| matrix: ${{ fromJson(needs.set-parameters.outputs.benchmark_matrix) }} | |
| fail-fast: false | |
| runs-on: ${{ matrix.runner }} | |
| env: | |
| WORKSPACE_DIR: "/workspace" | |
| SETUP_SCRIPT: "/workspace/setup-instance.sh" | |
| UV_VENV_DIR: "/workspace/uv_venvs" | |
| CONDA_ENV: ${{ matrix.triton_channel }} | |
| TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN }} | |
| JOB_NAME: tritonbench-${{ matrix.runner }}-benchmark-${{ matrix.triton_channel }}-periodic-${{ matrix.benchmarks }} | |
| RUNNER_TYPE: ${{ matrix.runner }} | |
| environment: pytorch-x-vllm | |
| permissions: | |
| id-token: write | |
| contents: read | |
| steps: | |
| - name: Checkout repository | |
| uses: actions/checkout@v4 | |
| - name: Install system dependencies | |
| shell: bash | |
| run: | | |
| sudo apt-get update | |
| sudo apt-get install -y libnuma-dev numactl | |
| - name: Checkout TritonBench repository | |
| uses: actions/checkout@v4 | |
| with: | |
| repository: meta-pytorch/tritonbench | |
| path: triton-benchmarks/tritonbench | |
| ref: ${{ inputs.tritonbench_branch }} | |
| submodules: recursive | |
| fetch-depth: 0 | |
| - name: Checkout Triton | |
| uses: actions/checkout@v4 | |
| with: | |
| repository: ${{ inputs.triton_repo || (matrix.triton_channel == 'meta-triton' && 'facebookexperimental/triton' || 'triton-lang/triton') }} | |
| path: triton-benchmarks/triton | |
| ref: ${{ inputs.triton_commit || 'main' }} | |
| submodules: recursive | |
| - uses: actions/setup-python@v5 | |
| # Amazon Linux fails on this step | |
| continue-on-error: true | |
| with: | |
| python-version: '3.12' | |
| cache: 'pip' | |
| - name: Check if the device is supported | |
| shell: bash | |
| run: | | |
| set -eux | |
| if command -v nvidia-smi; then | |
| DEVICE_NAME=cuda | |
| nvidia-smi | |
| elif command -v rocm-smi; then | |
| DEVICE_NAME=rocm | |
| rocm-smi | |
| else | |
| DEVICE_NAME=cpu | |
| lscpu | |
| fi | |
| echo "DEVICE_NAME=$DEVICE_NAME" >> $GITHUB_ENV | |
| - name: Set GPU name and type | |
| shell: bash | |
| run: | | |
| set -eux | |
| if [[ "${DEVICE_NAME}" == "cuda" ]]; then | |
| DEVICE_TYPE=$(nvidia-smi -i 0 --query-gpu=name --format=csv,noheader | awk '{print $2}') | |
| CUDA_HOME="/usr/local/cuda" | |
| echo "CUDA_HOME=$CUDA_HOME" >> $GITHUB_ENV | |
| elif [[ "${DEVICE_NAME}" == "rocm" ]]; then | |
| DEVICE_TYPE=$(rocminfo | grep "Marketing Name" | tail -n1 | awk -F':' '{print $2}' | xargs) | |
| elif [[ "${DEVICE_NAME}" == "cpu" ]]; then | |
| DEVICE_TYPE=$(lscpu | grep 'Model name' | cut -f 2 -d ":" | awk '{$1=$1}1' | cut -f 2 -d " ") | |
| fi | |
| echo "DEVICE_TYPE=$DEVICE_TYPE" >> $GITHUB_ENV | |
| - name: Install TritonBench and Build Triton | |
| working-directory: triton-benchmarks/tritonbench | |
| env: | |
| MAX_JOBS: 16 | |
| run: | | |
| set -eux | |
| # If running in nightly environment | |
| # Check out the "nightly" commit | |
| if [ -z ${{ inputs.triton_repo }} ] && \ | |
| [ -z ${{ inputs.triton_commit }} ]; then | |
| . ./.ci/triton/triton_install_utils.sh | |
| checkout_triton main \ | |
| "${GITHUB_WORKSPACE}/triton-benchmarks/triton" \ | |
| "1" | |
| fi | |
| bash ./.ci/tritonbench/setup-env.sh --cuda \ | |
| --custom-triton ${GITHUB_WORKSPACE}/triton-benchmarks/triton | |
| . "${SETUP_SCRIPT}" | |
| . "${UV_VENV_DIR}/${CONDA_ENV}/bin/activate" | |
| $HOME/.local/bin/uv pip install --group ci | |
| - name: Run TritonBench | |
| working-directory: triton-benchmarks/tritonbench | |
| run: | | |
| set -eux | |
| . "${SETUP_SCRIPT}" | |
| # install deps (if there is any) | |
| if [ -f benchmarks/${{ matrix.benchmarks }}/requirements.txt ]; then | |
| uv pip install -r benchmarks/${{ matrix.benchmarks }}/requirements.txt | |
| fi | |
| # Run TritonBench on the first available CPU core | |
| # Single CPU core is needed to stabilize the benchmark results | |
| first_available_core=$(taskset -pc $$| sed -n 's/.*: \([0-9][0-9]*\).*/\1/p') | |
| if [ ${{ matrix.benchmarks }} == "tlx" ]; then | |
| bash .ci/tritonbench/run-benchmark.sh ${{ matrix.benchmarks }} --conda-env ${{ env.CONDA_ENV }} ${{ inputs.benchmark_parameters }} | |
| else | |
| taskset -c ${first_available_core} bash .ci/tritonbench/run-benchmark.sh ${{ matrix.benchmarks }} --conda-env ${{ env.CONDA_ENV }} ${{ inputs.benchmark_parameters }} | |
| fi | |
| mv .benchmarks results-${{ env.CONDA_ENV }} | |
| # post-process result.json | |
| latest_result_json=$(find ./results-${CONDA_ENV} -name "result.json" | sort -r | head -n 1) | |
| python3 ./.ci/test_infra/oss_ci_benchmark_v3.py --json ${latest_result_json} \ | |
| --add-github-env --output ${latest_result_json} | |
| - name: Authenticate with AWS | |
| # AWS CUDA runners already have access to the bucket via its runner IAM role | |
| if: env.DEVICE_NAME == 'rocm' || contains(env.DEVICE_TYPE, 'B200') | |
| uses: aws-actions/configure-aws-credentials@ececac1a45f3b08a01d2dd070d28d111c5fe6722 # v4.1.0 | |
| with: | |
| role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_upload-benchmark-results | |
| # The max duration enforced by the server side | |
| role-duration-seconds: 18000 | |
| aws-region: us-east-1 | |
| # Keep a copy of the benchmark results on GitHub for reference | |
| - uses: actions/upload-artifact@v4 | |
| if: always() | |
| with: | |
| name: ${{ env.JOB_NAME }} | |
| path: triton-benchmarks/tritonbench/results-${{ env.CONDA_ENV }} | |
| retention-days: 30 | |
| - name: Upload result to Scribe | |
| if: ${{ env.TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN != '' }} | |
| working-directory: triton-benchmarks/tritonbench | |
| run: | | |
| . "${SETUP_SCRIPT}" | |
| latest_result_json=$(find ./results-${CONDA_ENV} -name "result.json" | sort -r | head -n 1) | |
| python3 ./.ci/upload/scribe.py --json ${latest_result_json} | |
| - name: Rewrite Tritonbench result json to ClickHouse style | |
| working-directory: triton-benchmarks/tritonbench | |
| run: | | |
| . "${SETUP_SCRIPT}" | |
| latest_result_json=$(find ./results-${CONDA_ENV} -name "result.json" | sort -r | head -n 1) | |
| python3 ./.ci/test_infra/oss_ci_benchmark_v3.py --json ${latest_result_json} \ | |
| --output clickhouse-results/result-${CONDA_ENV}.json | |
| - name: Upload result to ClickHouse | |
| uses: pytorch/test-infra/.github/actions/upload-benchmark-results@main | |
| with: | |
| benchmark-results-dir: triton-benchmarks/tritonbench/clickhouse-results | |
| dry-run: false | |
| schema-version: v3 | |
| github-token: ${{ secrets.GITHUB_TOKEN }} |