Skip to content

Commit 78440aa

Browse files
xuzhao9xuzhao9
authored
[tritonbench] add bisect workflow (#137)
1 parent 5124544 commit 78440aa

2 files changed

Lines changed: 222 additions & 3 deletions

File tree

‎.github/scripts/test_generate_tritonbench_matrix.py‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,10 @@
88
def test_generate_benchmark_matrix():
99
# All combinations, no duplication
1010
benchmarks = []
11+
triton_channels = []
1112
runners = []
1213
output = json.dumps(
13-
generate_benchmark_matrix(benchmarks, runners), indent=2
14+
generate_benchmark_matrix(benchmarks, triton_channels, runners), indent=2
1415
)
1516
assert_expected_inline(
1617
output,
@@ -19,14 +20,34 @@ def test_generate_benchmark_matrix():
1920
"include": [
2021
{
2122
"runner": "linux.dgx.b200",
22-
"triton_channel: "triton-main",
23+
"triton_channel": "triton-main",
2324
"benchmarks": "nightly"
2425
},
2526
{
2627
"runner": "linux.dgx.b200",
27-
"triton_channel: "meta-triton",
28+
"triton_channel": "meta-triton",
2829
"benchmarks": "nightly"
2930
}
3031
]
3132
}""",
3233
)
34+
35+
benchmarks = ["bisect"]
36+
triton_channels = ["triton-main"]
37+
runners = ["b200"]
38+
output = json.dumps(
39+
generate_benchmark_matrix(benchmarks, triton_channels, runners), indent=2
40+
)
41+
assert_expected_inline(
42+
output,
43+
"""\
44+
{
45+
"include": [
46+
{
47+
"runner": "linux.dgx.b200",
48+
"triton_channel": "triton-main",
49+
"benchmarks": "bisect"
50+
}
51+
]
52+
}""",
53+
)

‎.github/workflows/tritonbench-bisect.yml‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
name: TritonBench Bisect
2+
3+
on:
4+
workflow_dispatch:
5+
tritonbench_branch:
6+
description: TritonBench branch (main)
7+
required: true
8+
type: string
9+
default: main
10+
runners:
11+
description: |
12+
Hardware
13+
required: true
14+
type: choice
15+
options:
16+
- 'b200'
17+
bisect_type:
18+
type: choice
19+
default: 'performance'
20+
options:
21+
- 'performance'
22+
- 'functional'
23+
triton_channel:
24+
type: choice
25+
default: 'triton-main'
26+
options:
27+
- 'triton-main'
28+
- 'meta-triton'
29+
description: 'Triton channel to bisect'
30+
repro_cmdline:
31+
required: True
32+
type: string
33+
description: |
34+
The command line to reproduce the regression
35+
good_commit:
36+
required: True
37+
type: string
38+
description: |
39+
Last good commit (no regression)
40+
bad_commit:
41+
required: True
42+
type: string
43+
description: |
44+
First bad commit (has regression)
45+
regression_threshold:
46+
type: number
47+
default: 10
48+
description: |
49+
Performance regression threshold in %
50+
51+
52+
concurrency:
53+
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }}
54+
cancel-in-progress: true
55+
56+
57+
jobs:
58+
set-parameters:
59+
runs-on: ubuntu-latest
60+
outputs:
61+
benchmark_matrix: ${{ steps.set-parameters.outputs.benchmark_matrix }}
62+
steps:
63+
- name: Checkout repository
64+
uses: actions/checkout@v4
65+
66+
- uses: actions/setup-python@v5
67+
with:
68+
python-version: '3.12'
69+
70+
- name: Set parameters
71+
id: set-parameters
72+
shell: bash
73+
env:
74+
RUNNERS: ${{ inputs.runners || '' }}
75+
run: |
76+
set -eux
77+
78+
# The generated matrix is grouped by benchmark and runner
79+
python .github/scripts/generate_tritonbench_matrix.py \
80+
--runners "${RUNNERS}" --benchmark bisect --triton ${{ inputs.triton_channel }}
81+
82+
bisect:
83+
name: Run TritonBench Bisect
84+
needs: set-parameters
85+
if: ${{ !github.event.pull_request.head.repo.fork && github.repository_owner == 'pytorch' }}
86+
strategy:
87+
matrix: ${{ fromJson(needs.set-parameters.outputs.benchmark_matrix) }}
88+
fail-fast: false
89+
runs-on: ${{ matrix.runner }}
90+
env:
91+
WORKSPACE_DIR: "/workspace"
92+
SETUP_SCRIPT: "/workspace/setup-instance.sh"
93+
UV_VENV_DIR: "/workspace/uv_venvs"
94+
CONDA_ENV: ${{ matrix.triton_channel }}
95+
TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN }}
96+
JOB_NAME: tritonbench-${{ matrix.runner }}-bisect-${{ matrix.triton_channel }}
97+
RUNNER_TYPE: ${{ matrix.runner }}
98+
GOOD_COMMIT: ${{ inputs.good_commit }}
99+
BAD_COMMIT: ${{ inputs.bad_commit }}
100+
REGRESSION_THRESHOLD: ${{ inputs.regression_threshold }}
101+
REPRO_CMDLINE: ${{ inputs.repro_cmdline }}
102+
FUNCTIONAL: ${{ inputs.bisect_type == 'functional' && '1' || '0' }}
103+
environment: pytorch-x-vllm
104+
permissions:
105+
id-token: write
106+
contents: read
107+
steps:
108+
- name: Checkout repository
109+
uses: actions/checkout@v4
110+
111+
- name: Install system dependencies
112+
shell: bash
113+
run: |
114+
sudo apt-get update
115+
sudo apt-get install -y libnuma-dev numactl
116+
117+
- name: Checkout TritonBench repository
118+
uses: actions/checkout@v4
119+
with:
120+
repository: meta-pytorch/tritonbench
121+
path: triton-benchmarks/tritonbench
122+
ref: main
123+
submodules: recursive
124+
fetch-depth: 0
125+
126+
- uses: actions/setup-python@v5
127+
# Amazon Linux fails on this step
128+
continue-on-error: true
129+
with:
130+
python-version: '3.12'
131+
cache: 'pip'
132+
133+
- name: Check if the device is supported
134+
shell: bash
135+
run: |
136+
set -eux
137+
138+
if command -v nvidia-smi; then
139+
DEVICE_NAME=cuda
140+
nvidia-smi
141+
elif command -v rocm-smi; then
142+
DEVICE_NAME=rocm
143+
rocm-smi
144+
else
145+
DEVICE_NAME=cpu
146+
lscpu
147+
fi
148+
echo "DEVICE_NAME=$DEVICE_NAME" >> $GITHUB_ENV
149+
150+
- name: Set GPU name and type
151+
shell: bash
152+
run: |
153+
set -eux
154+
155+
if [[ "${DEVICE_NAME}" == "cuda" ]]; then
156+
DEVICE_TYPE=$(nvidia-smi -i 0 --query-gpu=name --format=csv,noheader | awk '{print $2}')
157+
CUDA_HOME="/usr/local/cuda"
158+
echo "CUDA_HOME=$CUDA_HOME" >> $GITHUB_ENV
159+
elif [[ "${DEVICE_NAME}" == "rocm" ]]; then
160+
DEVICE_TYPE=$(rocminfo | grep "Marketing Name" | tail -n1 | awk -F':' '{print $2}' | xargs)
161+
elif [[ "${DEVICE_NAME}" == "cpu" ]]; then
162+
DEVICE_TYPE=$(lscpu | grep 'Model name' | cut -f 2 -d ":" | awk '{$1=$1}1' | cut -f 2 -d " ")
163+
fi
164+
echo "DEVICE_TYPE=$DEVICE_TYPE" >> $GITHUB_ENV
165+
166+
- name: Install TritonBench
167+
working-directory: triton-benchmarks/tritonbench
168+
run: |
169+
set -eux
170+
# Use MAX_JOBS=16 to avoid OOM compiling Triton
171+
if [ "${CONDA_ENV}" == "triton-main" ]; then
172+
CMD_SUFFIX="--triton-main"
173+
elif [ "${CONDA_ENV}" == "meta-triton" ]; then
174+
CMD_SUFFIX="--meta-triton"
175+
else
176+
echo "unknown conda env: ${CONDA_ENV}"
177+
exit 1
178+
fi
179+
MAX_JOBS=16 bash ./.ci/tritonbench/setup-env.sh --cuda ${CMD_SUFFIX}
180+
181+
- name: Run TritonBench Bisect
182+
working-directory: triton-benchmarks/tritonbench
183+
run: |
184+
set -eux
185+
186+
# Run TritonBench on the first available CPU core
187+
# Single CPU core is needed to stabilize the benchmark results
188+
first_available_core=$(taskset -pc $$| sed -n 's/.*: \([0-9][0-9]*\).*/\1/p')
189+
taskset -c ${first_available_core} bash .ci/bisect/run.sh
190+
mv bisect_logs bisect-results-${{ env.CONDA_ENV }}
191+
192+
# Keep a copy of the benchmark results on GitHub for reference
193+
- uses: actions/upload-artifact@v4
194+
if: always()
195+
with:
196+
name: ${{ env.JOB_NAME }}
197+
path: triton-benchmarks/tritonbench/bisect-results-${{ env.CONDA_ENV }}
198+
retention-days: 30

0 commit comments

Comments
 (0)