Skip to content

Add colfax_cutlass backend to flash_attention operator #2330

Add colfax_cutlass backend to flash_attention operator

Add colfax_cutlass backend to flash_attention operator #2330

Workflow file for this run

name: TorchBench PR Test on A10G
on:
pull_request:
workflow_dispatch:
push:
branches:
- main
env:
CONDA_ENV: "torchbench"
DOCKER_IMAGE: "ghcr.io/pytorch/torchbench:latest"
SETUP_SCRIPT: "/workspace/setup_instance.sh"
HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
jobs:
pr-test:
# AWS A10G GPU instance label: linux.g5.4xlarge.nvidia.gpu
# OS version: Amazon Linux 2
runs-on: [self-hosted, linux.g5.4xlarge.nvidia.gpu]
timeout-minutes: 1440 # 24 hours
environment: docker-s3-upload
steps:
- name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
uses: pytorch/test-infra/.github/actions/setup-ssh@main
with:
github-secret: ${{ secrets.TORCHBENCH_ACCESS_TOKEN }}
- name: Checkout TorchBench
uses: actions/checkout@v3
with:
path: benchmark
- name: Pull docker image
uses: pytorch/test-infra/.github/actions/pull-docker-image@main
with:
docker-image: ${{ env.DOCKER_IMAGE }}
- name: Install NVIDIA Driver, docker runtime, set GPU_FLAG
id: install-nvidia-driver
uses: pytorch/test-infra/.github/actions/setup-nvidia@main
- name: Install and Test TorchBench
run: |
container_name=$(docker run \
-e CONDA_ENV="${CONDA_ENV}" \
-e SETUP_SCRIPT="${SETUP_SCRIPT}" \
-e HUGGING_FACE_HUB_TOKEN="${HUGGING_FACE_HUB_TOKEN}" \
--tty \
--detach \
--shm-size=32gb \
-v "${PWD}/benchmark:/benchmark" \
--gpus all \
-w / \
"${{ env.DOCKER_IMAGE }}" \
tail -f /dev/null
)
echo "Container name: ${container_name}"
docker exec -t -w "/" "${container_name}" bash -c "sudo chown -R runner /benchmark; sudo chgrp -R runner /benchmark"
docker exec -t -w "/benchmark" "${container_name}" bash /benchmark/scripts/torchbench_install.sh
docker exec -t -w "/benchmark" "${container_name}" bash /benchmark/scripts/torchbench_test.sh
- name: Teardown Linux
uses: pytorch/test-infra/.github/actions/teardown-linux@main
if: always()
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
cancel-in-progress: true