From 26311733f131a9de23aa10622ac5126f52e1b5e3 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Tue, 18 Aug 2026 15:20:27 -0700 Subject: [PATCH] Add gfx950 ROCm CI for AOTI Add an opt-in, path-filtered ROCm job on the gfx950 runner. Build the HIP runtime and native runner, exercise targeted Python and C++ coverage, verify generated AOTI code and HIP-only linkage, and provide a generic ROCm LLM preset. With assistance from Claude Code and Codex. --- .ci/scripts/test-rocm-aoti.sh | 158 ++++++++++++++++++++++++++ .github/pytorch-probot.yml | 1 + .github/workflows/rocm.yml | 123 ++++++++++++++++++++ CMakePresets.json | 43 +++++++ backends/cuda/rocm.md | 14 +-- backends/cuda/tests/test_fused_moe.py | 6 +- 6 files changed, 334 insertions(+), 11 deletions(-) create mode 100644 .ci/scripts/test-rocm-aoti.sh create mode 100644 .github/workflows/rocm.yml diff --git a/.ci/scripts/test-rocm-aoti.sh b/.ci/scripts/test-rocm-aoti.sh new file mode 100644 index 00000000000..0f4ac9d826f --- /dev/null +++ b/.ci/scripts/test-rocm-aoti.sh @@ -0,0 +1,158 @@ +#!/usr/bin/env bash +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +set -euo pipefail + +ROCM_VERSION="${ROCM_VERSION:-7.1}" +ROCM_PATH="${ROCM_PATH:-/opt/rocm}" +PYTORCH_ROCM_INDEX="${PYTORCH_ROCM_INDEX:-https://download.pytorch.org/whl/test/rocm${ROCM_VERSION}}" +TORCHAO_ROCM_WHEEL_BASE="${TORCHAO_ROCM_WHEEL_BASE:-https://download.pytorch.org/whl/nightly/rocm${ROCM_VERSION}}" +ROCM_CI_TMP_ROOT="${RUNNER_TEMP:-/tmp}" +mkdir -p "${ROCM_CI_TMP_ROOT}" 2>/dev/null || ROCM_CI_TMP_ROOT=/tmp +ROCM_CI_TMPDIR="$(mktemp -d "${ROCM_CI_TMP_ROOT}/executorch-rocm-ci.XXXXXX")" +trap 'rm -rf "${ROCM_CI_TMPDIR}"' EXIT + +export ROCM_PATH +export HIP_VISIBLE_DEVICES=0 +export CUDA_VISIBLE_DEVICES=0 +export TORCHINDUCTOR_CACHE_DIR="${ROCM_CI_TMPDIR}/inductor-cache" +export TORCHINDUCTOR_COMPILE_THREADS=1 + +read -r TORCH_VERSION TORCHAO_VERSION < <( + python - <<'PY' +from install_requirements import TORCHAO_NIGHTLY_VERSION +from torch_pin import TORCH_VERSION + +print(TORCH_VERSION, TORCHAO_NIGHTLY_VERSION) +PY +) +# TorchAO ROCm wheels are not exposed by the per-version pip index. +TORCHAO_WHEEL="${TORCHAO_ROCM_WHEEL_BASE}/torchao-${TORCHAO_VERSION}" +TORCHAO_WHEEL+="%2Brocm${ROCM_VERSION}-cp310-abi3-manylinux_2_28_x86_64.whl" +python -m pip install "torch==${TORCH_VERSION}" \ + --index-url "${PYTORCH_ROCM_INDEX}" +python -m pip install -r requirements-dev.txt \ + "${TORCHAO_WHEEL}" +EXECUTORCH_BUILD_MINIMAL=1 \ + python -m pip install --editable . --no-build-isolation + +if ! command -v conda >/dev/null; then + echo "The ROCm CI image must provide conda for its C++ runtime libraries" + exit 1 +fi +conda install -y -c conda-forge 'libstdcxx-ng>=12' + +python - <<'PY' +import torch +import torchao +import triton + +assert torch.version.hip is not None, "PyTorch is not a ROCm build" +assert torch.version.cuda is None, "PyTorch unexpectedly reports a CUDA runtime" +assert torch.cuda.is_available(), "No AMD GPU is visible through PyTorch" +assert "+rocm" in torchao.__version__, "TorchAO is not a ROCm build" + +device = torch.cuda.get_device_properties(0) +arch = device.gcnArchName.split(":", 1)[0] +assert arch == "gfx950", f"Expected gfx950, got {device.gcnArchName}" +assert device.warp_size == 64, f"Expected wave64, got {device.warp_size}" +assert arch in torch.cuda.get_arch_list(), ( + f"{arch} is not supported by this PyTorch build: {torch.cuda.get_arch_list()}" +) + +print("PyTorch:", torch.__version__) +print("TorchAO:", torchao.__version__) +print("ROCm:", torch.version.hip) +print("Triton:", triton.__version__) +print("Device:", device.name, device.gcnArchName) +print("Architectures:", torch.cuda.get_arch_list()) +PY + +if command -v rocminfo >/dev/null; then + rocminfo | sed -n '1,160p' +fi +if command -v rocm-smi >/dev/null; then + rocm-smi --showproductname --showmeminfo vram --showuse --showtemp || true +elif command -v amd-smi >/dev/null; then + amd-smi static --gpu all || true +fi + +TORCH_CMAKE_PREFIX="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)')" +export CMAKE_PREFIX_PATH="${TORCH_CMAKE_PREFIX}" + +cmake --preset llm-release-rocm -DEXECUTORCH_BUILD_TESTS=ON +cmake --build cmake-out-rocm-llm \ + --target executor_runner test_cuda_allocator test_cuda_mutable_state \ + --parallel "$(nproc)" + +PYTHON_PREFIX="$(python -c 'import sys; print(sys.prefix)')" +export LD_LIBRARY_PATH="${PYTHON_PREFIX}/lib:${PWD}/cmake-out-rocm-llm/backends/cuda:${PWD}/cmake-out-rocm-llm/extension/cuda:${ROCM_PATH}/lib:${LD_LIBRARY_PATH:-}" + +ctest --test-dir cmake-out-rocm-llm \ + -R 'test_cuda_(allocator|mutable_state)' \ + -V + +# Targeted builds leave the shim in the build tree rather than install lib/. +for ROCM_CI_BINARY in \ + cmake-out-rocm-llm/backends/cuda/libaoti_cuda_shims.so \ + cmake-out-rocm-llm/executor_runner; do + ROCM_CI_LINKS="$(ldd "${ROCM_CI_BINARY}")" + printf '%s\n' "${ROCM_CI_LINKS}" + if grep -q 'not found' <<<"${ROCM_CI_LINKS}"; then + echo "Missing runtime dependency in ${ROCM_CI_BINARY}" + exit 1 + fi + if grep -q 'libcudart' <<<"${ROCM_CI_LINKS}"; then + echo "Unexpected CUDA runtime dependency in ${ROCM_CI_BINARY}" + exit 1 + fi + if ! grep -q 'libamdhip64' <<<"${ROCM_CI_LINKS}"; then + echo "HIP runtime dependency missing from ${ROCM_CI_BINARY}" + exit 1 + fi +done + +if ! grep -q '^EXECUTORCH_BUILD_CUDA:BOOL=OFF$' \ + cmake-out-rocm-llm/CMakeCache.txt; then + echo "ROCm CI unexpectedly enabled the CUDA build" + exit 1 +fi +if ! grep -q '^EXECUTORCH_BUILD_ROCM:BOOL=ON$' \ + cmake-out-rocm-llm/CMakeCache.txt; then + echo "ROCm CI did not enable the ROCm build" + exit 1 +fi +if ! grep -q 'find_dependency(hip CONFIG)' \ + cmake-out-rocm-llm/executorch-backend-dependencies.cmake; then + echo "ROCm package metadata does not declare its HIP dependency" + exit 1 +fi +if grep -q 'find_dependency(CUDAToolkit)' \ + cmake-out-rocm-llm/executorch-backend-dependencies.cmake; then + echo "ROCm package metadata unexpectedly depends on CUDAToolkit" + exit 1 +fi + +export EXECUTORCH_EXECUTOR_RUNNER="${PWD}/cmake-out-rocm-llm/executor_runner" +ROCM_EXCLUDED_TESTS=( + # This focused job does not install the optional flash-linear-attention package. + --ignore=backends/cuda/tests/test_chunk_gated_delta_rule.py + # Defer specialized SDPA suites until their CI cost is measured. + --ignore=backends/cuda/tests/test_tq4_sdpa.py + --ignore=backends/cuda/tests/test_triton_sdpa.py + --ignore=backends/cuda/tests/test_triton_sdpa_splitk.py + # This test process core-dumps on ROCm instead of reporting a failure. + --ignore=backends/cuda/tests/test_triton_sdpa_nan.py +) +python -m pytest -v -o 'addopts=' \ + "${ROCM_EXCLUDED_TESTS[@]}" \ + backends/cuda/tests \ + backends/cuda/passes/tests + +ROCM_POINTWISE_DIR="${ROCM_CI_TMPDIR}/pointwise" +python -m examples.cuda.scripts.export_amd_pointwise \ + --output-dir "${ROCM_POINTWISE_DIR}" diff --git a/.github/pytorch-probot.yml b/.github/pytorch-probot.yml index d5c027ad4b4..87f4832c7da 100644 --- a/.github/pytorch-probot.yml +++ b/.github/pytorch-probot.yml @@ -7,6 +7,7 @@ ciflow_push_tags: - ciflow/cuda-perf - ciflow/metal - ciflow/mlx +- ciflow/rocm - ciflow/nightly - ciflow/trunk - ciflow/binaries diff --git a/.github/workflows/rocm.yml b/.github/workflows/rocm.yml new file mode 100644 index 00000000000..193d592f2c8 --- /dev/null +++ b/.github/workflows/rocm.yml @@ -0,0 +1,123 @@ +# Test ExecuTorch AOTI ROCm support. +# Use gfx950 for functional coverage and reserve scarce gfx1100 RDNA runners +# for a future focused performance job. + +name: Test ROCm AOTI + +on: + push: + branches: + - main + - release/* + tags: + - ciflow/rocm/* + pull_request: + paths: + - .github/workflows/rocm.yml + - .github/workflows/_ci-run-decision.yml + - .github/workflows/_get-changed-files.yml + - .ci/scripts/test-rocm-aoti.sh + - CMakeLists.txt + - CMakePresets.json + - install_requirements.py + - torch_pin.py + - backends/aoti/** + - '!backends/aoti/**.md' + - backends/cuda/** + - '!backends/cuda/**.md' + - extension/cuda/** + - examples/cuda/** + - '!examples/cuda/**.md' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + cancel-in-progress: false + +permissions: + contents: read + +jobs: + changed-files: + name: Get changed files + uses: ./.github/workflows/_get-changed-files.yml + with: + include-push-diff: true + + run-decision: + name: CI run decision + uses: ./.github/workflows/_ci-run-decision.yml + + aoti-run-decision: + name: AOTI CI run decision + needs: [changed-files, run-decision] + runs-on: ubuntu-latest + outputs: + run-aoti: ${{ steps.decide.outputs.run-aoti }} + steps: + - name: Select AOTI runs + id: decide + env: + CHANGED_FILES: ${{ needs.changed-files.outputs.changed-files }} + FULL_RUN: ${{ needs.run-decision.outputs.is-full-run }} + run: | + set -efu + RUN_AOTI=false + + if [[ "$CHANGED_FILES" == "*" || "$FULL_RUN" == "true" ]]; then + RUN_AOTI=true + fi + + for path in $CHANGED_FILES; do + case "$path" in + backends/aoti/*.md | backends/cuda/*.md | examples/cuda/*.md) + continue + ;; + .github/workflows/rocm.yml | \ + .github/workflows/_ci-run-decision.yml | \ + .github/workflows/_get-changed-files.yml | \ + .ci/scripts/test-rocm-aoti.sh | \ + CMakeLists.txt | \ + CMakePresets.json | \ + install_requirements.py | \ + torch_pin.py | \ + backends/aoti/* | \ + backends/cuda/* | \ + extension/cuda/* | \ + examples/cuda/*) + RUN_AOTI=true + break + ;; + esac + done + + echo "run-aoti=$RUN_AOTI" >> "$GITHUB_OUTPUT" + + unittest-rocm-gfx950: + name: unittest-rocm-gfx950-rocm${{ matrix.rocm-version }} + needs: [aoti-run-decision] + if: needs.aoti-run-decision.outputs.run-aoti == 'true' + strategy: + fail-fast: false + matrix: + rocm-version: ["7.1"] + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + permissions: + id-token: write + contents: read + with: + timeout: 90 + no-sudo: true + runner: linux.rocm.gpu.gfx950.1 + gpu-arch-type: rocm + gpu-arch-version: ${{ matrix.rocm-version }} + use-custom-docker-registry: false + docker-image: pytorch/manylinux2_28-builder:rocm${{ matrix.rocm-version }} + submodules: recursive + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + script: | + set -eux + eval "$(conda shell.bash hook)" + conda create -n executorch-rocm-ci python=3.10 -y + conda activate executorch-rocm-ci + ROCM_VERSION=${{ matrix.rocm-version }} bash .ci/scripts/test-rocm-aoti.sh diff --git a/CMakePresets.json b/CMakePresets.json index 6ddea5fd69c..09685ffe65a 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -161,6 +161,25 @@ "list": ["Linux", "Windows"] } }, + { + "name": "llm-release-rocm", + "displayName": "LLM release build with ROCm", + "inherits": [ + "llm-release" + ], + "binaryDir": "${sourceDir}/cmake-out-rocm-llm", + "cacheVariables": { + "CMAKE_INSTALL_PREFIX": "${sourceDir}/cmake-out-rocm-llm", + "CMAKE_PREFIX_PATH": "$penv{ROCM_PATH};/opt/rocm;$penv{CMAKE_PREFIX_PATH}", + "EXECUTORCH_BUILD_CUDA": "OFF", + "EXECUTORCH_BUILD_ROCM": "ON" + }, + "condition": { + "lhs": "${hostSystemName}", + "type": "equals", + "rhs": "Linux" + } + }, { "name": "llm-release-metal", "displayName": "LLM release build with Metal", @@ -397,6 +416,16 @@ ], "jobs": 0 }, + { + "name": "llm-release-rocm-install", + "displayName": "Build and install LLM extension release artifacts (ROCm)", + "configurePreset": "llm-release-rocm", + "configuration": "Release", + "targets": [ + "install" + ], + "jobs": 0 + }, { "name": "llm-release-metal-install", "displayName": "Build and install LLM extension release artifacts (Metal)", @@ -500,6 +529,20 @@ } ] }, + { + "name": "llm-release-rocm", + "displayName": "Configure, build and install ExecuTorch LLM extension with ROCm enabled", + "steps": [ + { + "type": "configure", + "name": "llm-release-rocm" + }, + { + "type": "build", + "name": "llm-release-rocm-install" + } + ] + }, { "name": "llm-release-metal", "displayName": "Configure, build and install ExecuTorch LLM extension with Metal enabled", diff --git a/backends/cuda/rocm.md b/backends/cuda/rocm.md index 2344e8ef99e..dbab11c85e0 100644 --- a/backends/cuda/rocm.md +++ b/backends/cuda/rocm.md @@ -7,7 +7,7 @@ GPUs as device type `cuda`, so names such as `CudaBackend` and `EXECUTORCH_BUILD_ROCM` is off by default, is never auto-enabled, and is mutually exclusive with `EXECUTORCH_BUILD_CUDA`. It requires `EXECUTORCH_BUILD_EXTENSION_TENSOR`. Execution has been validated only on -MI300X (`gfx942`). +MI300X (`gfx942`); CI is configured to exercise `gfx950`. ## Requirements and limitations @@ -18,23 +18,19 @@ MI300X (`gfx942`). - Python pybindings cannot allocate ROCm device memory; use a native runner such as `executor_runner`. - Installed CMake consumers must be able to find the HIP package. -- Model runners do not yet link the ROCm backend, and there is no ROCm CI. +- Model runners do not yet claim ROCm support. The gfx950 CI job covers the + backend build, native runtime, AOTI export and execution, and Triton W4 tests. ## Build ```bash -cmake -S . -B cmake-out-rocm \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_PREFIX_PATH="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)');/opt/rocm" \ - -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ - -DEXECUTORCH_BUILD_ROCM=ON -cmake --build cmake-out-rocm --target aoti_cuda_backend aoti_cuda_shims -j +cmake --workflow --preset llm-release-rocm ``` Verify that the build uses HIP and not the NVIDIA runtime: ```bash -ldd cmake-out-rocm/backends/cuda/libaoti_cuda_shims.so | grep -E 'amdhip64|cudart' +ldd cmake-out-rocm-llm/lib/libaoti_cuda_shims.so | grep -E 'amdhip64|cudart' ``` A `libcudart` dependency is a configuration error. diff --git a/backends/cuda/tests/test_fused_moe.py b/backends/cuda/tests/test_fused_moe.py index 324fd88907d..1a93607c34f 100644 --- a/backends/cuda/tests/test_fused_moe.py +++ b/backends/cuda/tests/test_fused_moe.py @@ -25,7 +25,6 @@ import torch import torch.nn as nn import torch.nn.functional as F - from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner from executorch.backends.cuda.triton.kernels.fused_moe import ( @@ -43,7 +42,10 @@ from torch.export import export EXECUTORCH_ROOT = os.path.normpath(os.path.join(os.path.dirname(__file__), "../../..")) -RUNNER_PATH = os.path.join(EXECUTORCH_ROOT, "cmake-out", "executor_runner") +RUNNER_PATH = os.environ.get( + "EXECUTORCH_EXECUTOR_RUNNER", + os.path.join(EXECUTORCH_ROOT, "cmake-out", "executor_runner"), +) # Test configurations: (seed, M, hidden, intermediate, num_experts, top_k, group_size) TEST_CONFIGS = [