tritonbench.yml

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

          deactivate
          . "${SETUP_SCRIPT}"
          $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 }}