GPU Performance Benchmark #7
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: GPU Performance Benchmark | |
| on: | |
| push: | |
| branches: [main, master] | |
| paths: | |
| - 'pallasbench/kernels/**' | |
| - 'pallasbench/baselines/**' | |
| workflow_dispatch: | |
| inputs: | |
| size: | |
| description: 'Benchmark size (SMALL, MEDIUM, LARGE)' | |
| required: true | |
| default: 'MEDIUM' | |
| type: choice | |
| options: [SMALL, MEDIUM, LARGE] | |
| levels: | |
| description: 'Levels to benchmark (space-separated)' | |
| required: true | |
| default: '1 2 3' | |
| alert_threshold: | |
| description: 'Alert threshold percentage (e.g., 150 for 1.5x regression)' | |
| required: false | |
| default: '150' | |
| env: | |
| BENCH_SIZE: ${{ inputs.size || 'MEDIUM' }} | |
| BENCH_LEVELS: ${{ inputs.levels || '1 2 3' }} | |
| BENCH_ALERT: ${{ inputs.alert_threshold || '150' }} | |
| jobs: | |
| benchmark: | |
| runs-on: [self-hosted, gpu, a100] | |
| name: GPU benchmark (${{ inputs.size || 'MEDIUM' }}) | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: Set up Python | |
| uses: actions/setup-python@v5 | |
| with: | |
| python-version: "3.11" | |
| - name: Install JAX (CUDA) | |
| run: | | |
| pip install -e ".[gpu,dev]" | |
| - name: Verify GPU | |
| run: | | |
| python -c "import jax; print(f'Backend: {jax.default_backend()}'); print(f'Devices: {jax.devices()}')" | |
| - name: Run benchmark | |
| run: | | |
| python scripts/run_benchmark.py \ | |
| --levels ${{ env.BENCH_LEVELS }} \ | |
| --size ${{ env.BENCH_SIZE }} \ | |
| --n-trials 100 \ | |
| --n-warmup 20 \ | |
| --output-dir results/ \ | |
| --pytest-benchmark-json results/benchmark.json | |
| - name: Store benchmark result | |
| uses: benchmark-action/github-action-benchmark@v1 | |
| with: | |
| name: PallasBench GPU (${{ env.BENCH_SIZE }}) | |
| tool: 'customSmallerIsBetter' | |
| output-file-path: results/benchmark.json | |
| github-token: ${{ secrets.GITHUB_TOKEN }} | |
| alert-threshold: '${{ env.BENCH_ALERT }}%' | |
| comment-on-alert: true | |
| fail-on-alert: true | |
| auto-push: true | |
| - name: Upload results | |
| uses: actions/upload-artifact@v4 | |
| if: always() | |
| with: | |
| name: gpu-benchmark-${{ env.BENCH_SIZE }} | |
| path: results/*.json |