Enable FLASH_ATTN backend with upstream flash-attn CK on ROCm for decode

.github/workflows/build-rocm-wheels.yml

@@ -2,12 +2,11 @@ name: Build ROCm Wheels
 
 on:
   push:
-    tags:
-      - 'v*'
     branches:
-      - main
-      - 'matthias.awq_gemv'
-      - 'release/*'
+      - gfx11
+  pull_request:
+    branches:
+      - gfx11
   workflow_dispatch:
     inputs:
       pytorch_index:
@@ -16,17 +15,20 @@ on:
         required: false
       rocm_arch:
         description: 'ROCm architecture (e.g., gfx1151, gfx1150;gfx1151)'
-        default: 'gfx1150;gfx1151;gfx1152;gfx1153'
+        default: 'gfx1150;gfx1151'
         required: false
 
 env:
   PYTORCH_INDEX_URL: ${{ github.event.inputs.pytorch_index || 'https://rocm.nightlies.amd.com/v2/gfx1151' }}
-  PYTORCH_ROCM_ARCH: ${{ github.event.inputs.rocm_arch || 'gfx1150;gfx1151;gfx1152;gfx1153' }}
+  PYTORCH_ROCM_ARCH: ${{ github.event.inputs.rocm_arch || 'gfx1150;gfx1151' }}
 
 jobs:
   build-wheel:
     runs-on: ubuntu-latest
-
+    permissions:
+      id-token: write
+      contents: read
+
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -37,8 +39,8 @@ jobs:
       - name: Fetch upstream tags for versioning
         run: |
           git remote add upstream https://github.com/vllm-project/vllm.git || true
-          git fetch upstream --tags
-          echo "Git describe: $(git describe --tags --always)"
+          git fetch upstream --tags || echo "Warning: could not fetch upstream tags"
+          git describe --tags --always || echo "Warning: git describe failed"
 
       - name: Free disk space
         uses: jlumbroso/free-disk-space@main
@@ -184,194 +186,27 @@ jobs:
           path: dist/*.whl
           retention-days: 90
 
-      - name: Generate wheel index
-        run: |
-          python scripts/generate-wheel-index.py \
-            --wheels-dir dist \
-            --output-dir gh-pages \
-            --base-url "https://${{ github.repository_owner }}.github.io/${{ github.event.repository.name }}/wheels"
-
-      - name: Upload index artifact
-        uses: actions/upload-artifact@v4
+      - name: Configure AWS credentials via OIDC
+        if: github.event_name == 'push' && github.repository_owner == 'ROCm'
+        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0
         with:
-          name: wheel-index
-          path: gh-pages/
-          retention-days: 90
+          role-to-assume: arn:aws:iam::317668459450:role/therock-aig-embd-gfx11-wheels-s3-oidc
+          aws-region: us-east-2
 
-  create-release:
-    needs: build-wheel
-    runs-on: ubuntu-latest
-    permissions:
-      contents: write
-    outputs:
-      release_tag: ${{ steps.create-tag.outputs.tag }}
-
-    steps:
-      - name: Download wheel artifact
-        uses: actions/download-artifact@v4
-        with:
-          name: vllm-rocm-wheel
-          path: dist/
-
-      - name: Determine release tag
-        id: create-tag
+      - name: Upload wheel to S3
+        if: github.event_name == 'push' && github.repository_owner == 'ROCm'
         run: |
-          # shellcheck disable=SC2193  # github.ref is substituted by Actions
-          if [[ "${{ github.ref }}" == refs/tags/* ]]; then
-            TAG="${{ github.ref_name }}"
-          else
-            WHEEL=$(find dist -name '*.whl' -print -quit)
-            BASENAME=$(basename "$WHEEL" .whl)
-            # Extract version from wheel name: vllm-<version>-cp312-...
-            VERSION=$(echo "$BASENAME" | sed 's/^vllm-//; s/-cp[0-9].*//')
-            TAG="wheels/rocm-${VERSION}"
-          fi
-          echo "tag=$TAG" >> "$GITHUB_OUTPUT"
-          echo "Release tag: $TAG"
-
-      - name: Create GitHub Release
-        uses: softprops/action-gh-release@v2
-        with:
-          tag_name: ${{ steps.create-tag.outputs.tag }}
-          target_commitish: ${{ github.sha }}
-          files: dist/*.whl
-          prerelease: ${{ !startsWith(github.ref, 'refs/tags/') }}
-          body: |
-            ## vLLM ROCm Wheel
-
-            Built for ROCm architectures: `${{ env.PYTORCH_ROCM_ARCH }}`
-
-            ### Installation
-            ```bash
-            pip install torch --extra-index-url https://rocm.nightlies.amd.com/v2/gfx1151
-            pip install vllm --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: vllm-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 }}
-        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 = """<!DOCTYPE html>
-          <html>
-            <head><meta name="pypi:repository-version" content="1.0"><title>{package_name}</title></head>
-            <body>
-          {items}
-            </body>
-          </html>
-          """
-
-          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("_", "-")
-
-          # Collect wheels from all GitHub Releases
-          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 = {}  # pkg_name -> [(whl_name, download_url)]
-          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))
-
-          # Include new wheels from this build if not already found in releases
-          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'    <a href="{url}">{whl_name}</a><br/>')
-              (pkg_dir / "index.html").write_text(
-                  PACKAGE_INDEX_TEMPLATE.format(package_name=pkg_name, items="\n".join(links))
-              )
-
-          top_index = f"""<!DOCTYPE html>
-          <html>
-            <head><meta name="pypi:repository-version" content="1.0"><title>ROCm Wheels</title></head>
-            <body>
-              <h1>ROCm Wheels</h1>
-              <p>Pre-built wheels for AMD ROCm (gfx11 / Strix Halo / RDNA3+)</p>
-              <p><em>Generated: {datetime.now().isoformat()}</em></p>
-              <h2>Installation</h2>
-              <pre>pip install torch --extra-index-url https://rocm.nightlies.amd.com/v2/gfx1151
-          pip install vllm --extra-index-url {BASE_URL}</pre>
-              <h2>Available Packages</h2>
-          """
-          for p in sorted(wheels):
-              top_index += f'    <a href="{p}/">{p}</a><br/>\n'
-          top_index += "  </body>\n</html>"
-          (out / "index.html").write_text(top_index)
-          print("Index generated successfully")
+          pip install boto3
+          python - <<'PYEOF'
+          import boto3, glob
+
+          s3 = boto3.client("s3")
+          bucket = "aig-embd-gfx11-wheels"
+
+          for whl in glob.glob("dist/*.whl"):
+              key = whl.split("/")[-1]
+              print(f"Uploading {key} to s3://{bucket}/{key}")
+              s3.upload_file(whl, bucket, key)
+              print(f"  Done")
           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

vllm/v1/attention/backends/fa_utils.py

@@ -31,7 +31,96 @@
     get_scheduler_metadata = xpu_ops.get_scheduler_metadata  # type: ignore[assignment]
 elif current_platform.is_rocm():
     try:
-        from flash_attn import flash_attn_varlen_func  # type: ignore[no-redef]
+        import torch
+        from flash_attn.flash_attn_interface import (
+            _wrapped_flash_attn_varlen_forward,
+        )
+
+        def flash_attn_varlen_func(  # type: ignore[no-redef,misc]
+            q: torch.Tensor,
+            k: torch.Tensor,
+            v: torch.Tensor,
+            cu_seqlens_q: torch.Tensor,
+            max_seqlen_q: int,
+            max_seqlen_k: int,
+            softmax_scale: float,
+            causal: bool = True,
+            window_size: "list[int] | None" = None,
+            block_table: "torch.Tensor | None" = None,
+            softcap: float = 0.0,
+            alibi_slopes: "torch.Tensor | None" = None,
+            out: "torch.Tensor | None" = None,
+            seqused_k: "torch.Tensor | None" = None,
+            cu_seqlens_k: "torch.Tensor | None" = None,
+            return_softmax_lse: bool = False,
+            # vLLM-specific kwargs not supported by upstream, ignored:
+            scheduler_metadata: Any = None,
+            fa_version: Any = None,
+            q_descale: Any = None,
+            k_descale: Any = None,
+            v_descale: Any = None,
+            num_splits: Any = None,
+            s_aux: Any = None,
+        ) -> "torch.Tensor | tuple[torch.Tensor, torch.Tensor]":
+            """Wrapper that adapts vLLM's calling convention to upstream
+            flash-attn 2.x API on ROCm.
+
+            The upstream _wrapped_flash_attn_varlen_forward accepts both
+            cu_seqlens_k and seqused_k. When using paged KV cache
+            (block_table is set), seqused_k gives per-sequence K lengths
+            and cu_seqlens_k can be a dummy tensor.
+            """
+            if window_size is None:
+                window_size_left = -1
+                window_size_right = -1
+            else:
+                window_size_left = window_size[0]
+                window_size_right = window_size[1]
+
+            # Build cu_seqlens_k if not provided.
+            # For paged KV cache, cu_seqlens_k must be the cumulative
+            # sum of actual sequence lengths (seqused_k), NOT zeros.
+            if cu_seqlens_k is None:
+                if seqused_k is not None:
+                    cu_seqlens_k = torch.zeros(
+                        seqused_k.shape[0] + 1,
+                        dtype=torch.int32,
+                        device=q.device,
+                    )
+                    cu_seqlens_k[1:] = seqused_k.cumsum(0)
+                else:
+                    batch_size = cu_seqlens_q.shape[0] - 1
+                    cu_seqlens_k = torch.zeros(
+                        batch_size + 1, dtype=torch.int32, device=q.device
+                    )
+
+            out_padded, softmax_lse, _, _ = _wrapped_flash_attn_varlen_forward(
+                q,
+                k,
+                v,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                max_seqlen_q,
+                max_seqlen_k,
+                0.0,  # dropout_p
+                softmax_scale,
+                causal,
+                window_size_left=window_size_left,
+                window_size_right=window_size_right,
+                softcap=softcap,
+                alibi_slopes=alibi_slopes,
+                return_softmax=False,
+                block_table=block_table,
+                seqused_k=seqused_k,
+            )
+            if out is not None:
+                out.copy_(out_padded)
+                result = out
+            else:
+                result = out_padded
+            if return_softmax_lse:
+                return result, softmax_lse
+            return result
 
         # Mark that upstream flash-attn is available on ROCm
         _ROCM_FLASH_ATTN_AVAILABLE = True
@@ -59,8 +148,8 @@ def get_flash_attn_version(
     if current_platform.is_xpu():
         return 2
     if current_platform.is_rocm():
-        # ROCm doesn't use vllm_flash_attn; return None to skip fa_version arg
-        return None
+        # ROCm uses upstream flash-attn which is FA2-compatible
+        return 2 if _ROCM_FLASH_ATTN_AVAILABLE else None
     try:
         from vllm.vllm_flash_attn.flash_attn_interface import (
             fa_version_unsupported_reason,

vllm/v1/attention/backends/flash_attn.py

@@ -43,6 +43,7 @@
 from vllm.config.cache import CacheDType
 from vllm.distributed.parallel_state import get_dcp_group
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
 from vllm.utils.math_utils import cdiv, round_up
 from vllm.v1.attention.backend import (
@@ -70,6 +71,9 @@ class FlashAttentionBackend(AttentionBackend):
 
     @staticmethod
     def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        if current_platform.is_rocm():
+            # CK flash-attn requires block_size divisible by 128
+            return [MultipleOf(128)]
         vllm_config = get_current_vllm_config()
         model_config = vllm_config.model_config
         cache_config = vllm_config.cache_config