From d50db16f86a216da7f9b8e73827a2cac8460afa4 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 30 Apr 2026 00:37:17 +0200 Subject: [PATCH 1/8] Auto-build flash-attn wheels on push, upload to S3 Replace the GitHub Releases / gh-pages publishing path with a direct upload to s3://aig-embd-gfx11-wheels/simple/flash-attn/ (the same PEP 503 index used by build-rocm-wheels.yml). Each push to gfx11 runs a check job that resolves the upstream Dao-AILab/flash-attention `main` HEAD and queries S3; the build job is skipped when a wheel matching the upstream short SHA already exists. The wheel version is derived from `git describe` against the latest v2.* tag, e.g. `2.8.4.dev472+gb995b246` for 472 commits past v2.8.3, and becomes plain `2.8.3` (or `2.8.4`) again once upstream lands a new tag. Changes: - Switch source from the v2.* release-list to upstream `main` because the latest release (v2.8.3, Aug 2025) is too old to include the gfx11 improvements we want. - Drop schedule and workflow_dispatch triggers (the workflow file is not on the default branch, so neither would actually fire). - Drop the create-release and publish-to-gh-pages jobs. - Drop the FLASH_ATTN_LOCAL_VERSION suffix; the SHA in the version string is enough to identify the build. Signed-off-by: Matthias Gehre --- .github/workflows/build-flash-attn-wheels.yml | 321 +++++++----------- 1 file changed, 123 insertions(+), 198 deletions(-) diff --git a/.github/workflows/build-flash-attn-wheels.yml b/.github/workflows/build-flash-attn-wheels.yml index ec0af6e5e65b..2f36f82ae0be 100644 --- a/.github/workflows/build-flash-attn-wheels.yml +++ b/.github/workflows/build-flash-attn-wheels.yml @@ -1,6 +1,7 @@ name: Build Flash Attn ROCm Wheels -# This workflow builds flash_attn wheels with the Triton backend for ROCm. +# This workflow builds flash_attn wheels with the Triton backend for ROCm and +# uploads them to s3://aig-embd-gfx11-wheels/simple/flash-attn/ (PEP 503 index). # # IMPORTANT: Architecture Independence # ------------------------------------ @@ -11,46 +12,86 @@ name: Build Flash Attn ROCm Wheels # 2. The Triton kernels are compiled JIT (just-in-time) at runtime # 3. The resulting wheel is architecture-agnostic # -# The PyTorch index is only used during build because setup.py imports torch -# to detect IS_HIP_EXTENSION. The actual GPU target (gfx1150, gfx1151, etc.) -# is determined at runtime when Triton compiles the attention kernels. +# Versioning: the workflow tracks the `main` branch of +# Dao-AILab/flash-attention. The wheel version is derived from `git describe` +# against the latest v2.* tag, e.g. +# on a tag → 2.8.3 +# N commits past → 2.8.4.dev+g +# Once upstream tags v2.8.4, the dev-version override falls away automatically. # -# This means users can install this wheel and use it on any gfx11 target -# that Triton supports, regardless of which PyTorch index was used during build. +# The check job skips the build when a wheel matching the upstream HEAD SHA +# is already present in the S3 index. on: - workflow_dispatch: - inputs: - pytorch_index: - description: 'PyTorch index URL (only needed for build-time torch import)' - default: 'https://rocm.nightlies.amd.com/v2-staging/gfx1151' - required: false - flash_attn_repo: - description: 'Flash Attention repository URL' - default: 'https://github.com/ROCm/flash-attention' - required: false - flash_attn_ref: - description: 'Flash Attention git ref (branch/tag/commit)' - default: 'main' - required: false - local_version: - description: 'Local version suffix for wheel (e.g., rocm.triton.gfx11)' - default: 'rocm.triton.gfx11' - required: false + push: + branches: + - gfx11 env: - PYTORCH_INDEX_URL: ${{ github.event.inputs.pytorch_index || 'https://rocm.nightlies.amd.com/v2-staging/gfx1151' }} - FLASH_ATTN_REPO: ${{ github.event.inputs.flash_attn_repo || 'https://github.com/ROCm/flash-attention' }} - FLASH_ATTN_REF: ${{ github.event.inputs.flash_attn_ref || 'main' }} - # Enable Triton backend - this skips CK compilation and makes the wheel arch-agnostic - FLASH_ATTENTION_TRITON_AMD_ENABLE: "TRUE" - # Custom version suffix to indicate this is a ROCm Triton build - FLASH_ATTN_LOCAL_VERSION: ${{ github.event.inputs.local_version || 'rocm.triton.gfx11' }} + PYTORCH_INDEX_URL: 'https://rocm.nightlies.amd.com/v2-staging/gfx1151' + FLASH_ATTN_REPO: 'Dao-AILab/flash-attention' + FLASH_ATTN_BRANCH: 'main' + S3_BUCKET: aig-embd-gfx11-wheels + S3_ROLE_ARN: arn:aws:iam::317668459450:role/therock-aig-embd-gfx11-wheels-s3-oidc + AWS_REGION: us-east-1 jobs: + check: + runs-on: ubuntu-latest + permissions: + id-token: write + contents: read + outputs: + sha: ${{ steps.resolve.outputs.sha }} + should_build: ${{ steps.decide.outputs.should_build }} + steps: + - name: Resolve flash-attention HEAD + id: resolve + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + REPO: ${{ env.FLASH_ATTN_REPO }} + BRANCH: ${{ env.FLASH_ATTN_BRANCH }} + run: | + set -euo pipefail + SHA=$(gh api "repos/${REPO}/commits/${BRANCH}" -q '.sha' | cut -c1-8) + echo "Latest ${REPO}@${BRANCH} = ${SHA}" + echo "sha=$SHA" >> "$GITHUB_OUTPUT" + + - name: Configure AWS credentials via OIDC + uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0 + with: + role-to-assume: ${{ env.S3_ROLE_ARN }} + aws-region: ${{ env.AWS_REGION }} + + - name: Decide whether to build + id: decide + env: + SHA: ${{ steps.resolve.outputs.sha }} + run: | + set -euo pipefail + # `git describe` puts the SHA after `g` in the version string, so + # the wheel filename contains `+g${SHA}-` (or, on a tag, no SHA at + # all -- but then upstream has cut a release and we just want the + # next release-tracking workflow to fire). + MATCH=$(aws s3api list-objects-v2 \ + --bucket "$S3_BUCKET" \ + --prefix "simple/flash-attn/flash_attn-" \ + --query "Contents[?contains(Key, '+g${SHA}-')].Key" \ + --output text) + if [ -n "$MATCH" ] && [ "$MATCH" != "None" ]; then + echo "Already in S3: $MATCH" + SHOULD=false + else + echo "No wheel matching +g${SHA}- in S3; building" + SHOULD=true + fi + echo "should_build=$SHOULD" >> "$GITHUB_OUTPUT" + build-wheel: + needs: check + if: needs.check.outputs.should_build == 'true' runs-on: ubuntu-latest - + steps: - name: Checkout code uses: actions/checkout@v4 @@ -58,11 +99,38 @@ jobs: fetch-depth: 1 - name: Clone flash-attention repository + env: + REPO: ${{ env.FLASH_ATTN_REPO }} + BRANCH: ${{ env.FLASH_ATTN_BRANCH }} run: | - git clone --depth 1 --branch ${{ env.FLASH_ATTN_REF }} ${{ env.FLASH_ATTN_REPO }} flash-attention + # Need full history (or at least back to the latest v2.* tag) so + # `git describe` can compute commit count + short SHA. + git clone --branch "$BRANCH" "https://github.com/${REPO}" flash-attention cd flash-attention - echo "Flash Attention commit: $(git rev-parse HEAD)" - echo "Flash Attention ref: ${{ env.FLASH_ATTN_REF }}" + echo "Flash Attention HEAD: $(git rev-parse HEAD)" + git describe --tags --long --match 'v2.*' + + - name: Compute wheel version from git describe + working-directory: flash-attention + run: | + set -euo pipefail + RAW=$(git describe --tags --long --match 'v2.*' --abbrev=8) + echo "git describe: $RAW" + # RAW looks like: v2.8.3-42-gabc12345 + TAG=${RAW#v}; TAG=${TAG%-*-*} # 2.8.3 + COUNT=$(echo "$RAW" | awk -F- '{print $(NF-1)}') # 42 + SHA=$(echo "$RAW" | awk -F- '{print $NF}') # gabc12345 + if [ "$COUNT" -eq 0 ]; then + VERSION="$TAG" + else + NEXT=$(echo "$TAG" | awk -F. -v OFS=. '{ $NF=$NF+1; print }') + VERSION="${NEXT}.dev${COUNT}+${SHA}" + fi + echo "Computed VERSION=$VERSION" + # Upstream's setup.py reads __version__ from flash_attn/__init__.py; + # rewrite it so bdist_wheel produces the right filename. + sed -i "s/^__version__ = .*/__version__ = \"${VERSION}\"/" flash_attn/__init__.py + grep '^__version__' flash_attn/__init__.py - name: Free disk space uses: jlumbroso/free-disk-space@main @@ -95,10 +163,7 @@ jobs: - name: Install PyTorch from ROCm nightly run: | - # Install PyTorch with ROCm support (needed for setup.py to detect HIP) pip install torch --index-url ${{ env.PYTORCH_INDEX_URL }} - - # Show what's installed pip list | grep -i -E "torch" || true python -c "import torch; print(f'PyTorch version: {torch.__version__}'); print(f'HIP: {torch.version.hip}')" @@ -109,13 +174,8 @@ jobs: - name: Build wheel env: FLASH_ATTENTION_TRITON_AMD_ENABLE: "TRUE" - FLASH_ATTN_LOCAL_VERSION: ${{ env.FLASH_ATTN_LOCAL_VERSION }} working-directory: flash-attention run: | - echo "=== Environment ===" - echo "FLASH_ATTENTION_TRITON_AMD_ENABLE=$FLASH_ATTENTION_TRITON_AMD_ENABLE" - echo "FLASH_ATTN_LOCAL_VERSION=$FLASH_ATTN_LOCAL_VERSION" - echo "" echo "=== Building wheel ===" python setup.py bdist_wheel --dist-dir=dist ls -la dist/ @@ -140,172 +200,37 @@ jobs: path: flash-attention/dist/*.whl retention-days: 90 - create-release: + upload-wheel: needs: build-wheel runs-on: ubuntu-latest + if: github.repository_owner == 'ROCm' permissions: - contents: write - outputs: - release_tag: ${{ steps.create-tag.outputs.tag }} + id-token: write + contents: read steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + sparse-checkout: .github/workflows/scripts + fetch-depth: 1 + - name: Download wheel artifact uses: actions/download-artifact@v4 with: name: flash-attn-rocm-wheel path: dist/ - - name: Determine release tag - id: create-tag - run: | - WHEEL=$(find dist -name '*.whl' -print -quit) - BASENAME=$(basename "$WHEEL" .whl) - VERSION=$(echo "$BASENAME" | sed 's/^flash_attn-//; s/-cp[0-9].*//') - TAG="wheels/flash-attn-${VERSION}" - echo "tag=$TAG" >> "$GITHUB_OUTPUT" - echo "Release tag: $TAG" - - - name: Create GitHub Release - uses: softprops/action-gh-release@v2 + - name: Configure AWS credentials via OIDC + uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0 with: - tag_name: ${{ steps.create-tag.outputs.tag }} - target_commitish: ${{ github.sha }} - files: dist/*.whl - prerelease: true - body: | - ## Flash Attention ROCm Wheel (Triton backend) - - Built with `FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE` (architecture-agnostic). + role-to-assume: ${{ env.S3_ROLE_ARN }} + aws-region: ${{ env.AWS_REGION }} - ### Installation - ```bash - pip install torch --extra-index-url https://rocm.nightlies.amd.com/v2-staging/gfx1151 - pip install flash-attn --extra-index-url https://${{ github.repository_owner }}.github.io/${{ github.event.repository.name }}/wheels - ``` - - publish-to-gh-pages: - needs: create-release - runs-on: ubuntu-latest - permissions: - contents: write - pages: write - - steps: - - name: Download wheel artifact - uses: actions/download-artifact@v4 - with: - name: flash-attn-rocm-wheel - path: new-wheels/ - - - name: Generate index pointing to GitHub Releases - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - REPO: ${{ github.repository }} - BASE_URL: "https://${{ github.repository_owner }}.github.io/${{ github.event.repository.name }}/wheels" - RELEASE_TAG: ${{ needs.create-release.outputs.release_tag }} + - name: Upload wheel to S3 run: | - pip install packaging - - cat > /tmp/generate_index.py << 'PYEOF' - import os, json, subprocess - from pathlib import Path - from datetime import datetime - from urllib.parse import quote - - REPO = os.environ["REPO"] - BASE_URL = os.environ["BASE_URL"] - RELEASE_TAG = os.environ.get("RELEASE_TAG", "latest") - - PACKAGE_INDEX_TEMPLATE = """ - - {package_name} - - {items} - - - """ - - def parse_wheel_name(filename): - base = filename.removesuffix(".whl") - parts = base.split("-") - for i, part in enumerate(parts): - if part[0].isdigit(): - return "-".join(parts[:i]).replace("_", "-") - return parts[0].replace("_", "-") - - result = subprocess.run( - ["gh", "release", "list", "--repo", REPO, "--limit", "100", "--json", "tagName"], - capture_output=True, text=True - ) - releases = json.loads(result.stdout) if result.returncode == 0 else [] - print(f"Found {len(releases)} releases") - - wheels = {} - for rel in releases: - tag = rel["tagName"] - result = subprocess.run( - ["gh", "release", "view", tag, "--repo", REPO, "--json", "assets"], - capture_output=True, text=True - ) - if result.returncode != 0: - continue - assets = json.loads(result.stdout).get("assets", []) - for asset in assets: - name = asset["name"] - url = asset["url"] - if name.endswith(".whl"): - pkg = parse_wheel_name(name) - wheels.setdefault(pkg, []).append((name, url)) - - new_dir = Path("new-wheels") - if new_dir.exists(): - for whl in new_dir.glob("*.whl"): - pkg = parse_wheel_name(whl.name) - release_names = {w[0] for w in wheels.get(pkg, [])} - if whl.name not in release_names: - url = f"https://github.com/{REPO}/releases/download/{quote(RELEASE_TAG, safe='')}/{quote(whl.name, safe='')}" - wheels.setdefault(pkg, []).append((whl.name, url)) - - print(f"Total: {sum(len(v) for v in wheels.values())} wheels in {len(wheels)} packages") - - out = Path("gh-pages/wheels") - out.mkdir(parents=True, exist_ok=True) - - for pkg_name, whl_list in wheels.items(): - pkg_dir = out / pkg_name - pkg_dir.mkdir(exist_ok=True) - links = [] - for whl_name, url in sorted(whl_list, key=lambda x: x[0], reverse=True): - links.append(f' {whl_name}
') - (pkg_dir / "index.html").write_text( - PACKAGE_INDEX_TEMPLATE.format(package_name=pkg_name, items="\n".join(links)) - ) - - top_index = f""" - - ROCm Wheels - -

ROCm Wheels

-

Pre-built wheels for AMD ROCm (gfx11 / Strix Halo / RDNA3+)

-

Generated: {datetime.now().isoformat()}

-

Installation

- 
pip install torch --extra-index-url https://rocm.nightlies.amd.com/v2-staging/gfx1151
-          pip install vllm --extra-index-url {BASE_URL}
-

Available Packages

- """ - for p in sorted(wheels): - top_index += f' {p}
\n' - top_index += " \n" - (out / "index.html").write_text(top_index) - print("Index generated successfully") - PYEOF - - python /tmp/generate_index.py - - - name: Deploy to gh-pages - uses: peaceiris/actions-gh-pages@v4 - with: - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: ./gh-pages/wheels - destination_dir: wheels - keep_files: false + pip install boto3 + python .github/workflows/scripts/upload_wheel_s3.py \ + --bucket "$S3_BUCKET" \ + --package flash-attn \ + --wheel-dir dist/ From ce06bc59d73a13eb7094f33135c87ecb5929647e Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 30 Apr 2026 10:19:15 +0200 Subject: [PATCH 2/8] flash-attn CI: run on PRs to gfx11, upload only on push Adds pull_request trigger so PRs targeting gfx11 exercise the build, and gates the upload-wheel job on github.event_name == 'push' so PR runs validate the build without populating the S3 index. Signed-off-by: Matthias Gehre --- .github/workflows/build-flash-attn-wheels.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build-flash-attn-wheels.yml b/.github/workflows/build-flash-attn-wheels.yml index 2f36f82ae0be..b70ac0cb635e 100644 --- a/.github/workflows/build-flash-attn-wheels.yml +++ b/.github/workflows/build-flash-attn-wheels.yml @@ -26,6 +26,9 @@ on: push: branches: - gfx11 + pull_request: + branches: + - gfx11 env: PYTORCH_INDEX_URL: 'https://rocm.nightlies.amd.com/v2-staging/gfx1151' @@ -203,7 +206,7 @@ jobs: upload-wheel: needs: build-wheel runs-on: ubuntu-latest - if: github.repository_owner == 'ROCm' + if: github.event_name == 'push' && github.repository_owner == 'ROCm' permissions: id-token: write contents: read From 8826b2a8eca51e39abfb44590053883cf92027fc Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Mon, 4 May 2026 18:42:51 +0200 Subject: [PATCH 3/8] Pin flash-attn build to pre-aiter commit bbe25ba Upstream commit 3f94643 ("[AMD] Migrate to Triton Backend to Aiter") introduced a hard triton==3.5.1 pin and moved the AMD Triton backend out of the flash_attn package into aiter. This breaks ROCm users: the triton pin downgrades their ROCm triton, and the wheel is no longer self-contained. Pin to bbe25ba (the parent commit) which still bundles flash_attn_triton_amd/ and has no triton version constraint. Signed-off-by: Matthias Gehre --- .github/workflows/build-flash-attn-wheels.yml | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/.github/workflows/build-flash-attn-wheels.yml b/.github/workflows/build-flash-attn-wheels.yml index b70ac0cb635e..13e9a3e8a037 100644 --- a/.github/workflows/build-flash-attn-wheels.yml +++ b/.github/workflows/build-flash-attn-wheels.yml @@ -12,14 +12,15 @@ name: Build Flash Attn ROCm Wheels # 2. The Triton kernels are compiled JIT (just-in-time) at runtime # 3. The resulting wheel is architecture-agnostic # -# Versioning: the workflow tracks the `main` branch of -# Dao-AILab/flash-attention. The wheel version is derived from `git describe` -# against the latest v2.* tag, e.g. +# Versioning: the workflow tracks a pinned commit of +# Dao-AILab/flash-attention (before the aiter migration that introduced a +# hard triton==3.5.1 pin and removed the bundled AMD Triton backend). +# The wheel version is derived from `git describe` against the latest v2.* +# tag, e.g. # on a tag → 2.8.3 # N commits past → 2.8.4.dev+g -# Once upstream tags v2.8.4, the dev-version override falls away automatically. # -# The check job skips the build when a wheel matching the upstream HEAD SHA +# The check job skips the build when a wheel matching the pinned SHA # is already present in the S3 index. on: @@ -33,7 +34,7 @@ on: env: PYTORCH_INDEX_URL: 'https://rocm.nightlies.amd.com/v2-staging/gfx1151' FLASH_ATTN_REPO: 'Dao-AILab/flash-attention' - FLASH_ATTN_BRANCH: 'main' + FLASH_ATTN_REF: 'bbe25bacff9f337568b66e69950e549df9839068' S3_BUCKET: aig-embd-gfx11-wheels S3_ROLE_ARN: arn:aws:iam::317668459450:role/therock-aig-embd-gfx11-wheels-s3-oidc AWS_REGION: us-east-1 @@ -48,16 +49,14 @@ jobs: sha: ${{ steps.resolve.outputs.sha }} should_build: ${{ steps.decide.outputs.should_build }} steps: - - name: Resolve flash-attention HEAD + - name: Resolve flash-attention ref id: resolve env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - REPO: ${{ env.FLASH_ATTN_REPO }} - BRANCH: ${{ env.FLASH_ATTN_BRANCH }} + REF: ${{ env.FLASH_ATTN_REF }} run: | set -euo pipefail - SHA=$(gh api "repos/${REPO}/commits/${BRANCH}" -q '.sha' | cut -c1-8) - echo "Latest ${REPO}@${BRANCH} = ${SHA}" + SHA=${REF:0:8} + echo "Pinned ref = ${SHA}" echo "sha=$SHA" >> "$GITHUB_OUTPUT" - name: Configure AWS credentials via OIDC @@ -104,13 +103,14 @@ jobs: - name: Clone flash-attention repository env: REPO: ${{ env.FLASH_ATTN_REPO }} - BRANCH: ${{ env.FLASH_ATTN_BRANCH }} + REF: ${{ env.FLASH_ATTN_REF }} run: | # Need full history (or at least back to the latest v2.* tag) so # `git describe` can compute commit count + short SHA. - git clone --branch "$BRANCH" "https://github.com/${REPO}" flash-attention + git clone "https://github.com/${REPO}" flash-attention cd flash-attention - echo "Flash Attention HEAD: $(git rev-parse HEAD)" + git checkout "$REF" + echo "Flash Attention at: $(git rev-parse HEAD)" git describe --tags --long --match 'v2.*' - name: Compute wheel version from git describe From e7fbe5025d0915341f42905dacd67f23e3ae40c4 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Mon, 4 May 2026 19:03:18 +0200 Subject: [PATCH 4/8] Allow flash-attn wheel upload on PR events Remove the push-only guard so the upload-wheel job runs on both push and pull_request events, enabling upload from this PR. Signed-off-by: Matthias Gehre --- .github/workflows/build-flash-attn-wheels.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build-flash-attn-wheels.yml b/.github/workflows/build-flash-attn-wheels.yml index 13e9a3e8a037..2101a80e131e 100644 --- a/.github/workflows/build-flash-attn-wheels.yml +++ b/.github/workflows/build-flash-attn-wheels.yml @@ -206,7 +206,7 @@ jobs: upload-wheel: needs: build-wheel runs-on: ubuntu-latest - if: github.event_name == 'push' && github.repository_owner == 'ROCm' + if: github.repository_owner == 'ROCm' permissions: id-token: write contents: read From aa85bffea6e91aff51da157c5d831e89a63f9110 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 21 May 2026 00:43:15 -0600 Subject: [PATCH 5/8] Build AITER wheels alongside flash-attn Extend the workflow with check/build/upload jobs for AITER, mirroring the flash-attn pipeline: * check-aiter resolves AITER_REF (default branch matthias.gfx11) to a SHA and short-circuits when a matching wheel is already in s3://aig-embd-gfx11-wheels/simple/amd-aiter/. * build-aiter-wheel runs on ubuntu-latest, recursive-clones aiter, derives a PEP 440 version via git describe + SETUPTOOLS_SCM_PRETEND_VERSION (aiter uses setuptools_scm), installs torch from the same ROCm nightly index, and builds bdist_wheel with PREBUILD_KERNELS=0 so no HIP toolchain is required on the runner. * upload-aiter-wheel is wired up but gated `if: false` so nothing lands in S3 until validated; flipping to the standard `github.repository_owner == 'ROCm'` guard enables it. GPU_ARCHS=gfx1151 is set in env now even though aiter/setup.py ignores it at PREBUILD_KERNELS=0, so switching to kernels-on later is a single env-var change (plus moving the build job into a ROCm container). Signed-off-by: Matthias Gehre --- .github/workflows/build-flash-attn-wheels.yml | 255 +++++++++++++++++- 1 file changed, 244 insertions(+), 11 deletions(-) diff --git a/.github/workflows/build-flash-attn-wheels.yml b/.github/workflows/build-flash-attn-wheels.yml index 2101a80e131e..6a8a92917743 100644 --- a/.github/workflows/build-flash-attn-wheels.yml +++ b/.github/workflows/build-flash-attn-wheels.yml @@ -1,7 +1,10 @@ -name: Build Flash Attn ROCm Wheels +name: Build Flash Attn and AITER ROCm Wheels -# This workflow builds flash_attn wheels with the Triton backend for ROCm and -# uploads them to s3://aig-embd-gfx11-wheels/simple/flash-attn/ (PEP 503 index). +# This workflow builds two wheels and (optionally) uploads them to +# s3://aig-embd-gfx11-wheels/ (PEP 503 index): +# +# * flash_attn (Triton backend) -> simple/flash-attn/ +# * amd-aiter (kernels off) -> simple/amd-aiter/ # # IMPORTANT: Architecture Independence # ------------------------------------ @@ -12,15 +15,18 @@ name: Build Flash Attn ROCm Wheels # 2. The Triton kernels are compiled JIT (just-in-time) at runtime # 3. The resulting wheel is architecture-agnostic # -# Versioning: the workflow tracks a pinned commit of -# Dao-AILab/flash-attention (before the aiter migration that introduced a -# hard triton==3.5.1 pin and removed the bundled AMD Triton backend). -# The wheel version is derived from `git describe` against the latest v2.* -# tag, e.g. -# on a tag → 2.8.3 -# N commits past → 2.8.4.dev+g +# AITER is currently built with PREBUILD_KERNELS=0, which short-circuits all +# HIP compilation in aiter/setup.py. GPU_ARCHS is still set so flipping to +# PREBUILD_KERNELS=1 (kernels baked in for gfx1151) is a one-line change -- +# at that point the AITER build job will additionally need ROCm available +# on the runner (e.g. via a `rocm/dev-ubuntu-22.04` container). +# +# Versioning: each wheel version is derived from `git describe` of the +# upstream repo, e.g. +# on a tag -> 2.8.3 +# N commits past -> 2.8.4.dev+g # -# The check job skips the build when a wheel matching the pinned SHA +# Each check job skips the build when a wheel matching the pinned SHA # is already present in the S3 index. on: @@ -35,6 +41,14 @@ env: PYTORCH_INDEX_URL: 'https://rocm.nightlies.amd.com/v2-staging/gfx1151' FLASH_ATTN_REPO: 'Dao-AILab/flash-attention' FLASH_ATTN_REF: 'bbe25bacff9f337568b66e69950e549df9839068' + AITER_REPO: 'ROCm/aiter' + # Branch (or tag/SHA) of AITER to build. The check-aiter job resolves + # this to a SHA via `git ls-remote`. + AITER_REF: 'matthias.gfx11' + # Unused while AITER_PREBUILD_KERNELS=0; kept so flipping to 1 needs + # no other changes. + AITER_GPU_ARCHS: 'gfx1151' + AITER_PREBUILD_KERNELS: '0' S3_BUCKET: aig-embd-gfx11-wheels S3_ROLE_ARN: arn:aws:iam::317668459450:role/therock-aig-embd-gfx11-wheels-s3-oidc AWS_REGION: us-east-1 @@ -237,3 +251,222 @@ jobs: --bucket "$S3_BUCKET" \ --package flash-attn \ --wheel-dir dist/ + + check-aiter: + runs-on: ubuntu-latest + permissions: + id-token: write + contents: read + outputs: + sha: ${{ steps.resolve.outputs.sha }} + full_sha: ${{ steps.resolve.outputs.full_sha }} + should_build: ${{ steps.decide.outputs.should_build }} + steps: + - name: Resolve AITER ref to commit SHA + id: resolve + env: + REPO: ${{ env.AITER_REPO }} + REF: ${{ env.AITER_REF }} + run: | + set -euo pipefail + # AITER_REF may be a branch, tag, or SHA. `git ls-remote` works + # for branches/tags; for a raw SHA we just truncate. + if FULL=$(git ls-remote "https://github.com/${REPO}" "${REF}" \ + | awk 'NR==1 {print $1}') && [ -n "$FULL" ]; then + echo "Resolved ${REF} -> ${FULL}" + else + FULL=$REF + echo "Treating ${REF} as raw SHA" + fi + SHA=${FULL:0:8} + echo "full_sha=$FULL" >> "$GITHUB_OUTPUT" + echo "sha=$SHA" >> "$GITHUB_OUTPUT" + + - name: Configure AWS credentials via OIDC + uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0 + with: + role-to-assume: ${{ env.S3_ROLE_ARN }} + aws-region: ${{ env.AWS_REGION }} + + - name: Decide whether to build + id: decide + env: + SHA: ${{ steps.resolve.outputs.sha }} + run: | + set -euo pipefail + # Wheel filename contains `+g${SHA}-` (see version derivation in + # build-aiter-wheel). Package name `amd-aiter` becomes + # `amd_aiter-` in the filename. + MATCH=$(aws s3api list-objects-v2 \ + --bucket "$S3_BUCKET" \ + --prefix "simple/amd-aiter/amd_aiter-" \ + --query "Contents[?contains(Key, '+g${SHA}-')].Key" \ + --output text) + if [ -n "$MATCH" ] && [ "$MATCH" != "None" ]; then + echo "Already in S3: $MATCH" + SHOULD=false + else + echo "No wheel matching +g${SHA}- in S3; building" + SHOULD=true + fi + echo "should_build=$SHOULD" >> "$GITHUB_OUTPUT" + + build-aiter-wheel: + needs: check-aiter + if: needs.check-aiter.outputs.should_build == 'true' + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 1 + + - name: Clone AITER repository + env: + REPO: ${{ env.AITER_REPO }} + FULL_SHA: ${{ needs.check-aiter.outputs.full_sha }} + run: | + set -euo pipefail + # Recursive submodules pull in composable_kernel; only required + # once we enable PREBUILD_KERNELS=1, but cheap enough to do now + # so the build environment matches the future kernels-on path. + git clone --recursive "https://github.com/${REPO}" aiter + cd aiter + git checkout "$FULL_SHA" + git submodule update --init --recursive + echo "AITER at: $(git rev-parse HEAD)" + git describe --tags --long --match 'v*' || true + + - name: Compute wheel version from git describe + working-directory: aiter + run: | + set -euo pipefail + RAW=$(git describe --tags --long --match 'v*' --abbrev=8) + echo "git describe: $RAW" + # RAW looks like: v0.1.14-rc0-8-gc387fb45 (or v2.8.3-42-gabc12345) + TAG=${RAW#v}; TAG=${TAG%-*-*} # 0.1.14-rc0 + COUNT=$(echo "$RAW" | awk -F- '{print $(NF-1)}') # 8 + SHA=$(echo "$RAW" | awk -F- '{print $NF}') # gc387fb45 + if [ "$COUNT" -eq 0 ]; then + VERSION="$TAG" + else + # awk numeric-coerces the last dotted field, so e.g. `14-rc0` + # becomes 15 (rc suffix dropped). Good enough for a dev build. + NEXT=$(echo "$TAG" | awk -F. -v OFS=. '{ $NF=$NF+1; print }') + VERSION="${NEXT}.dev${COUNT}+${SHA}" + fi + echo "Computed VERSION=$VERSION" + # AITER's setup.py uses `use_scm_version=True`; the pretend-version + # env var overrides setuptools_scm so the wheel filename is + # predictable and matches the check-aiter S3 lookup. + echo "SETUPTOOLS_SCM_PRETEND_VERSION=$VERSION" >> "$GITHUB_ENV" + + - name: Free disk space + uses: jlumbroso/free-disk-space@main + with: + tool-cache: false + android: true + dotnet: true + haskell: true + large-packages: true + docker-images: true + swap-storage: true + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Cache pip packages + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: pip-aiter-${{ runner.os }}-py312 + restore-keys: | + pip-aiter-${{ runner.os }}- + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build wheel setuptools setuptools_scm ninja packaging psutil pyyaml + + - name: Install PyTorch from ROCm nightly + run: | + pip install torch --index-url ${{ env.PYTORCH_INDEX_URL }} + pip list | grep -i -E "torch" || true + python -c "import torch; print(f'PyTorch version: {torch.__version__}'); print(f'HIP: {torch.version.hip}')" + + - name: Install AITER runtime dependencies + working-directory: aiter + run: | + pip install -r requirements.txt + + - name: Build wheel + env: + PREBUILD_KERNELS: ${{ env.AITER_PREBUILD_KERNELS }} + GPU_ARCHS: ${{ env.AITER_GPU_ARCHS }} + working-directory: aiter + run: | + echo "=== Building wheel (PREBUILD_KERNELS=$PREBUILD_KERNELS, GPU_ARCHS=$GPU_ARCHS) ===" + python setup.py bdist_wheel --dist-dir=dist + ls -la dist/ + + - name: Rename wheel for manylinux compliance + working-directory: aiter + run: | + cd dist + for wheel in *.whl; do + if [[ "$wheel" == *"linux"* ]] && [[ "$wheel" != *"manylinux"* ]]; then + new_name="${wheel/linux/manylinux_2_31}" + mv "$wheel" "$new_name" + echo "Renamed $wheel to $new_name" + fi + done + ls -la + + - name: Upload wheel artifact + uses: actions/upload-artifact@v4 + with: + name: aiter-rocm-wheel + path: aiter/dist/*.whl + retention-days: 90 + + upload-aiter-wheel: + needs: build-aiter-wheel + runs-on: ubuntu-latest + # Upload is intentionally disabled while the AITER build job is being + # validated. To enable, replace `if: false` with the same guard the + # flash-attn upload uses: + # if: github.repository_owner == 'ROCm' + if: false + permissions: + id-token: write + contents: read + + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + sparse-checkout: .github/workflows/scripts + fetch-depth: 1 + + - name: Download wheel artifact + uses: actions/download-artifact@v4 + with: + name: aiter-rocm-wheel + path: dist/ + + - name: Configure AWS credentials via OIDC + uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0 + with: + role-to-assume: ${{ env.S3_ROLE_ARN }} + aws-region: ${{ env.AWS_REGION }} + + - name: Upload wheel to S3 + run: | + pip install boto3 + python .github/workflows/scripts/upload_wheel_s3.py \ + --bucket "$S3_BUCKET" \ + --package amd-aiter \ + --wheel-dir dist/ From a0238fc4fb0b4066335184eb03d0062933fe2e8d Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 21 May 2026 00:46:55 -0600 Subject: [PATCH 6/8] Rename flash-attn jobs to disambiguate from AITER `check` -> `check-flash-attn`, `build-wheel` -> `build-flash-attn-wheel`. Signed-off-by: Matthias Gehre --- .github/workflows/build-flash-attn-wheels.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/build-flash-attn-wheels.yml b/.github/workflows/build-flash-attn-wheels.yml index 6a8a92917743..528621a27314 100644 --- a/.github/workflows/build-flash-attn-wheels.yml +++ b/.github/workflows/build-flash-attn-wheels.yml @@ -54,7 +54,7 @@ env: AWS_REGION: us-east-1 jobs: - check: + check-flash-attn: runs-on: ubuntu-latest permissions: id-token: write @@ -103,9 +103,9 @@ jobs: fi echo "should_build=$SHOULD" >> "$GITHUB_OUTPUT" - build-wheel: - needs: check - if: needs.check.outputs.should_build == 'true' + build-flash-attn-wheel: + needs: check-flash-attn + if: needs.check-flash-attn.outputs.should_build == 'true' runs-on: ubuntu-latest steps: @@ -218,7 +218,7 @@ jobs: retention-days: 90 upload-wheel: - needs: build-wheel + needs: build-flash-attn-wheel runs-on: ubuntu-latest if: github.repository_owner == 'ROCm' permissions: From 34d7e1f0d0e86a3bed98720b6288bdc6819bb12c Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 21 May 2026 00:57:25 -0600 Subject: [PATCH 7/8] AITER: prebuild only module_gemm_w4a16 Switch the AITER build from PREBUILD_KERNELS=0 (no kernels) to PREBUILD_ONLY=module_gemm_w4a16, which restricts HIP compilation to that single module while leaving the rest of the package as JIT. NOTE: this enables HIP compilation, so `ubuntu-latest` without ROCm will no longer satisfy the build. The build-aiter-wheel job will need a ROCm container (e.g. `rocm/dev-ubuntu-22.04`) or an in-job ROCm install before the next run can succeed. This commit changes the env var only; the runner upgrade is a follow-up. Signed-off-by: Matthias Gehre --- .github/workflows/build-flash-attn-wheels.yml | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/.github/workflows/build-flash-attn-wheels.yml b/.github/workflows/build-flash-attn-wheels.yml index 528621a27314..0db41d7818a1 100644 --- a/.github/workflows/build-flash-attn-wheels.yml +++ b/.github/workflows/build-flash-attn-wheels.yml @@ -15,11 +15,11 @@ name: Build Flash Attn and AITER ROCm Wheels # 2. The Triton kernels are compiled JIT (just-in-time) at runtime # 3. The resulting wheel is architecture-agnostic # -# AITER is currently built with PREBUILD_KERNELS=0, which short-circuits all -# HIP compilation in aiter/setup.py. GPU_ARCHS is still set so flipping to -# PREBUILD_KERNELS=1 (kernels baked in for gfx1151) is a one-line change -- -# at that point the AITER build job will additionally need ROCm available -# on the runner (e.g. via a `rocm/dev-ubuntu-22.04` container). +# AITER is currently built with PREBUILD_ONLY=module_gemm_w4a16, which +# restricts the prebuild path in aiter/setup.py to a single module. This +# still triggers HIP compilation for that module, so the build job needs +# ROCm (hipcc + headers) on the runner -- e.g. via a `rocm/dev-ubuntu-22.04` +# container. GPU_ARCHS=gfx1151 selects the target arch. # # Versioning: each wheel version is derived from `git describe` of the # upstream repo, e.g. @@ -45,10 +45,10 @@ env: # Branch (or tag/SHA) of AITER to build. The check-aiter job resolves # this to a SHA via `git ls-remote`. AITER_REF: 'matthias.gfx11' - # Unused while AITER_PREBUILD_KERNELS=0; kept so flipping to 1 needs - # no other changes. AITER_GPU_ARCHS: 'gfx1151' - AITER_PREBUILD_KERNELS: '0' + # Allowlist of modules to prebuild (comma-separated). Overrides the + # PREBUILD_KERNELS modes and limits HIP compilation to just these. + AITER_PREBUILD_ONLY: 'module_gemm_w4a16' S3_BUCKET: aig-embd-gfx11-wheels S3_ROLE_ARN: arn:aws:iam::317668459450:role/therock-aig-embd-gfx11-wheels-s3-oidc AWS_REGION: us-east-1 @@ -404,11 +404,11 @@ jobs: - name: Build wheel env: - PREBUILD_KERNELS: ${{ env.AITER_PREBUILD_KERNELS }} + PREBUILD_ONLY: ${{ env.AITER_PREBUILD_ONLY }} GPU_ARCHS: ${{ env.AITER_GPU_ARCHS }} working-directory: aiter run: | - echo "=== Building wheel (PREBUILD_KERNELS=$PREBUILD_KERNELS, GPU_ARCHS=$GPU_ARCHS) ===" + echo "=== Building wheel (PREBUILD_ONLY=$PREBUILD_ONLY, GPU_ARCHS=$GPU_ARCHS) ===" python setup.py bdist_wheel --dist-dir=dist ls -la dist/ From c6377f4540ce7f9c34e86d3addf9724a92bed20a Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 21 May 2026 01:05:19 -0600 Subject: [PATCH 8/8] AITER: install ROCm SDK via TheRock wheels for HIP compilation PREBUILD_ONLY=module_gemm_w4a16 needs hipcc + ROCm headers on the runner. Instead of switching to a `rocm/dev-ubuntu-22.04` container, pip-install the multi-arch ROCm wheels (libraries + devel + the gfx1151 device package) from rocm.nightlies.amd.com, run `rocm-sdk init`, and export ROCM_PATH / ROCM_HOME / HIP_PATH + PATH so setup.py and torch.utils.cpp_extension find hipcc. Signed-off-by: Matthias Gehre --- .github/workflows/build-flash-attn-wheels.yml | 25 ++++++++++++++++--- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build-flash-attn-wheels.yml b/.github/workflows/build-flash-attn-wheels.yml index 0db41d7818a1..7bc4919bc9dd 100644 --- a/.github/workflows/build-flash-attn-wheels.yml +++ b/.github/workflows/build-flash-attn-wheels.yml @@ -16,10 +16,11 @@ name: Build Flash Attn and AITER ROCm Wheels # 3. The resulting wheel is architecture-agnostic # # AITER is currently built with PREBUILD_ONLY=module_gemm_w4a16, which -# restricts the prebuild path in aiter/setup.py to a single module. This -# still triggers HIP compilation for that module, so the build job needs -# ROCm (hipcc + headers) on the runner -- e.g. via a `rocm/dev-ubuntu-22.04` -# container. GPU_ARCHS=gfx1151 selects the target arch. +# restricts the prebuild path in aiter/setup.py to a single module. HIP +# compilation for that module is enabled by pip-installing TheRock ROCm +# wheels (libraries + devel + the gfx1151 device package) and running +# `rocm-sdk init` -- no `rocm/dev-ubuntu-22.04` container required. +# GPU_ARCHS=gfx1151 selects the target arch. # # Versioning: each wheel version is derived from `git describe` of the # upstream repo, e.g. @@ -397,6 +398,22 @@ jobs: pip list | grep -i -E "torch" || true python -c "import torch; print(f'PyTorch version: {torch.__version__}'); print(f'HIP: {torch.version.hip}')" + - name: Install ROCm SDK (TheRock multi-arch wheels) + run: | + set -euo pipefail + pip install \ + --index-url https://rocm.nightlies.amd.com/whl-multi-arch/ \ + "rocm[libraries,device-gfx1151,devel]" + rocm-sdk init + ROCM_PATH=$(rocm-sdk path --root) + echo "Resolved ROCM_PATH=$ROCM_PATH" + # Export for subsequent steps (setup.py / hipcc / torch.utils.cpp_extension). + echo "ROCM_PATH=$ROCM_PATH" >> "$GITHUB_ENV" + echo "ROCM_HOME=$ROCM_PATH" >> "$GITHUB_ENV" + echo "HIP_PATH=$ROCM_PATH" >> "$GITHUB_ENV" + echo "$ROCM_PATH/bin" >> "$GITHUB_PATH" + "$ROCM_PATH/bin/hipcc" --version || true + - name: Install AITER runtime dependencies working-directory: aiter run: |