[CI] flash-mla: Add the sgl-flash-mla release workflow (#45651)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

.github/workflows/release-whl-flashmla.yml

@@ -0,0 +1,252 @@
+name: Release sgl-flash-mla
+
+on:
+  workflow_dispatch:
+    inputs:
+      version:
+        description: "Wheel version (e.g. 1.0.0, 1.0.1rc0)"
+        type: string
+        required: true
+      branch:
+        description: "FlashMLA branch to build from (default: sgl)"
+        type: string
+        required: false
+        default: 'sgl'
+      validation_only:
+        description: "Build and test without publishing wheels"
+        type: boolean
+        required: false
+        default: false
+      skip_tests:
+        description: "Skip the GPU test suite (publishes untested wheels)"
+        type: boolean
+        required: false
+        default: false
+
+concurrency:
+  group: release-sgl-flash-mla-${{ github.ref }}
+  cancel-in-progress: true
+
+env:
+  # The wheel pins this; it is also what the test jobs install.
+  # The extensions use PyTorch's stable C++ ABI and would load against any 2.10+,
+  # but a release is tested against one version.
+  TORCH_VER: "2.14.1"
+
+jobs:
+  build:
+    if: github.repository == 'sgl-project/sglang'
+    strategy:
+      matrix:
+        # One wheel per architecture. No Python matrix: the wheel is abi3
+        # (cp310-abi3) and serves CPython 3.10+, so 3.12 only drives the build.
+        python-version: ["3.12"]
+        cuda-version: ["13.4"]
+        arch: [x86_64, aarch64]
+        include:
+          - arch: x86_64
+            runner: x64-kernel-build-node
+          - arch: aarch64
+            runner: arm-kernel-build-node
+    runs-on: ${{ matrix.runner }}
+    steps:
+      - name: Clean workspace (remove root-owned files from prior runs)
+        run: |
+          docker run --rm -v "${{ github.workspace }}:/workspace" alpine:3 \
+            sh -c 'rm -rf /workspace/..?* /workspace/.[!.]* /workspace/*' || true
+
+      - uses: actions/checkout@v4
+
+      - name: Checkout FlashMLA
+        uses: actions/checkout@v4
+        with:
+          repository: sgl-project/FlashMLA
+          ref: ${{ inputs.branch || 'sgl' }}
+          path: FlashMLA
+          submodules: recursive
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Set wheel version
+        run: |
+          echo -n "${{ inputs.version }}" > FlashMLA/flash_mla/VERSION
+          cat FlashMLA/flash_mla/VERSION
+
+      - name: Build wheel
+        run: |
+          chmod +x ./scripts/build_sgl_flash_mla.sh ./scripts/rename_sgl_flash_mla_whl.sh
+          ./scripts/build_sgl_flash_mla.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" "${{ github.workspace }}/FlashMLA" "${{ matrix.arch }}"
+
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: flashmla-wheel-${{ matrix.arch }}
+          path: FlashMLA/dist/*.whl
+
+      - name: Upload PyPI artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: flashmla-pypi-${{ matrix.arch }}
+          path: FlashMLA/dist-pypi/*.whl
+
+  test:
+    needs: build
+    if: ${{ !inputs.skip_tests }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          # SM90 runs only the frozen kernels: the current ones are sm_100a/sm_103a.
+          # The frozen ones are built for sm_90a/sm_100a/sm_103a, so SM100 runs both.
+          - arch_label: sm90
+            runner: 8-gpu-h200
+            wheel_arch: x86_64
+            tests: legacy
+            timeout: 90
+          - arch_label: sm100
+            runner: 8-gpu-b200
+            wheel_arch: x86_64
+            tests: current legacy
+            timeout: 180
+          - arch_label: sm103-aarch64
+            runner: 4-gpu-gb300
+            wheel_arch: aarch64
+            tests: current legacy
+            timeout: 180
+    runs-on: ${{ matrix.runner }}
+    timeout-minutes: ${{ matrix.timeout }}
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Checkout FlashMLA
+        uses: actions/checkout@v4
+        with:
+          repository: sgl-project/FlashMLA
+          ref: ${{ inputs.branch || 'sgl' }}
+          path: FlashMLA
+          submodules: recursive
+
+      # The tests need only tests/. The source package would shadow the installed
+      # wheel (its directory has no built extensions), so it is removed.
+      - name: Remove the FlashMLA source package
+        run: rm -rf FlashMLA/flash_mla
+
+      - name: Download wheel
+        uses: actions/download-artifact@v4
+        with:
+          path: dist/
+          merge-multiple: true
+          pattern: flashmla-wheel-${{ matrix.wheel_arch }}
+
+      - name: Install wheel and deps
+        run: |
+          python3 -m pip install --upgrade pip
+          python3 -m pip install "torch==${TORCH_VER}" --index-url https://download.pytorch.org/whl/cu130 --force-reinstall
+          python3 -m pip install numpy einops rich pytest
+          # --force-reinstall: a base image may already carry this version, in
+          # which case pip skips the wheel and the job tests a stale extension.
+          python3 -m pip install dist/*.whl --force-reinstall --no-deps
+          # Both extensions register through STABLE_TORCH_LIBRARY at dlopen time,
+          # so a missing symbol surfaces here rather than at the first launch.
+          python3 -c "
+          import flash_mla
+          from flash_mla._loader import get_ops
+          print('flash_mla:', flash_mla.__file__)
+          import os
+          assert not flash_mla.__file__.startswith(os.environ['GITHUB_WORKSPACE']), 'not the installed wheel'
+          print('version:', flash_mla.__version__)
+          print('has_legacy:', flash_mla.has_legacy())
+          assert get_ops(use_legacy=False).sparse_decode_fwd is not None
+          assert get_ops(use_legacy=True).dense_decode_fwd is not None
+          "
+
+      - name: Run tests (current kernels)
+        if: contains(matrix.tests, 'current')
+        working-directory: FlashMLA
+        env:
+          PYTHONPATH: ${{ github.workspace }}/FlashMLA/tests
+        run: |
+          python3 -c "import os, flash_mla; assert not flash_mla.__file__.startswith(os.environ['GITHUB_WORKSPACE']), flash_mla.__file__"
+          python3 tests/test-sparse-prefill.py
+          python3 tests/test-sparse-decode.py
+
+      - name: Run tests (frozen kernels)
+        if: contains(matrix.tests, 'legacy')
+        working-directory: FlashMLA
+        env:
+          # conftest.py only fixes sys.path under pytest; these are scripts.
+          PYTHONPATH: ${{ github.workspace }}/FlashMLA/tests
+        run: |
+          python3 -c "import os, flash_mla; assert not flash_mla.__file__.startswith(os.environ['GITHUB_WORKSPACE']), flash_mla.__file__"
+          python3 -m pytest tests/legacy/ -q
+          # Dense MLA decoding is SM90-only.
+          if [ "${{ matrix.arch_label }}" = sm90 ]; then
+            python3 tests/legacy/test_flash_mla_dense_decoding.py
+          fi
+          python3 tests/legacy/test_flash_mla_sparse_prefill.py
+          python3 tests/legacy/test_flash_mla_sparse_decoding.py
+          python3 tests/legacy/test_flash_mla_sparse_decoding_v32_bit_equal.py
+
+  release:
+    if: |
+      !cancelled() && !inputs.validation_only &&
+      needs.build.result == 'success' &&
+      (needs.test.result == 'success' || needs.test.result == 'skipped')
+    needs: [build, test]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Download artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: dist/
+          merge-multiple: true
+          pattern: flashmla-wheel-*
+
+      - name: Download PyPI artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: dist-pypi/
+          merge-multiple: true
+          pattern: flashmla-pypi-*
+
+      - name: Record that the GPU test suite was skipped
+        if: ${{ inputs.skip_tests }}
+        run: |
+          echo "Releasing v${{ inputs.version }} with the GPU test suite skipped." \
+            >> "$GITHUB_STEP_SUMMARY"
+
+      - name: Upload to PyPI
+        run: |
+          pip install --upgrade twine "packaging>=24.2"
+          python3 -m twine upload --skip-existing dist-pypi/* -u __token__ -p ${{ secrets.SGL_FLASH_MLA_PYPI_TOKEN }}
+
+      - name: Release
+        uses: softprops/action-gh-release@v2
+        with:
+          tag_name: v${{ inputs.version }}
+          repository: sgl-project/whl
+          token: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }}
+          files: |
+            dist/*
+
+      - name: Clone wheel index
+        run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl
+        env:
+          WHL_TOKEN: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }}
+
+      - name: Update wheel index
+        run: python3 scripts/update_flashmla_whl_index.py --cuda 130
+
+      - name: Push wheel index
+        run: |
+          cd sgl-whl
+          git config --local user.name "sglang-bot"
+          git config --local user.email "sglangbot@gmail.com"
+          git add -A
+          git commit -m "update sgl-flash-mla whl index for v${{ inputs.version }}${{ inputs.skip_tests && ' (GPU tests skipped)' || '' }}"
+          git push

docker/sgl-flash-mla.Dockerfile

@@ -0,0 +1,37 @@
+ARG BASE_IMG=pytorch/manylinux2_28-builder
+# FlashMLA's kernels use `cvt` with `.scaled::n2::ue8m0`, which needs nvcc >= 13.1,
+# and older toolkits emit suboptimal instructions for them. Build with 13.4.
+ARG CUDA_VERSION=13.4
+
+FROM ${BASE_IMG}:cuda${CUDA_VERSION}
+
+ARG ARCH=x86_64
+ARG CUDA_VERSION=13.4
+ARG PYTHON_VERSION=3.12
+ARG PYTHON_TAG=cp312-cp312
+ARG TORCH_VER=2.14.1
+ARG PIP_DEFAULT_INDEX=https://pypi.python.org/simple
+ARG PYTORCH_MIRROR=download.pytorch.org
+
+ENV PYTHON_ROOT_PATH=/opt/python/${PYTHON_TAG}
+ENV PATH=${PYTHON_ROOT_PATH}/bin:${PATH}
+
+RUN yum install -y --nogpgcheck git wget tar gcc gcc-c++ make \
+ && yum clean all && rm -rf /var/cache/yum
+
+# The build links -lcuda, which on a machine with no driver resolves against the
+# toolkit's stub.
+RUN set -eux; \
+    if [ "${ARCH}" = "aarch64" ]; then _LIB=sbsa; else _LIB="${ARCH}"; fi; \
+    mkdir -p /usr/lib/${ARCH}-linux-gnu/; \
+    ln -sf /usr/local/cuda-${CUDA_VERSION}/targets/${_LIB}-linux/lib/stubs/libcuda.so /usr/lib/${ARCH}-linux-gnu/libcuda.so
+
+# torch is needed at build time (setup.py imports it) and is what the wheel pins.
+# The CUDA minor version of this torch does not have to match CUDA_VERSION above:
+# the extensions link libcudart.so.13, whose SONAME carries only the major, and
+# use no API newer than 13.0.
+RUN --mount=type=cache,id=sgl-flash-mla-pip,target=/root/.cache/pip \
+    set -eux; \
+    ${PYTHON_ROOT_PATH}/bin/pip install torch==${TORCH_VER} --index-url https://${PYTORCH_MIRROR}/whl/cu130; \
+    ${PYTHON_ROOT_PATH}/bin/pip install --index-url ${PIP_DEFAULT_INDEX} \
+        ninja setuptools wheel build numpy

scripts/build_sgl_flash_mla.sh

@@ -0,0 +1,130 @@
+#!/bin/bash
+# Build the sgl-flash-mla wheel inside a CUDA-versioned container.
+#
+# Usage: build_sgl_flash_mla.sh <PYTHON_VERSION> <CUDA_VERSION> <FLASHMLA_SRC> [ARCH]
+#   PYTHON_VERSION: e.g. 3.12 -- only picks the interpreter that drives the build.
+#                   The wheel is abi3 and serves CPython 3.10+ whichever is used.
+#   CUDA_VERSION:   13.4
+#   FLASHMLA_SRC:   path to a checkout of sgl-project/FlashMLA (submodules included)
+#   ARCH:           x86_64 (default) or aarch64
+#
+# Writes:
+#   <FLASHMLA_SRC>/dist/      — wheel tagged +cu130 and manylinux2014
+#   <FLASHMLA_SRC>/dist-pypi/ — the same wheel with +cu130 stripped
+#                               (PyPI rejects local-version segments)
+#
+# On the two version numbers: the kernels are compiled by CUDA 13.4, because they
+# use `cvt` with `.scaled::n2::ue8m0` (nvcc >= 13.1) and older toolkits emit
+# suboptimal instructions for them. The +cu130 tag describes the *runtime* the
+# wheel is published for: the extensions link libcudart.so.13, whose SONAME
+# carries only the major version, and call no cudart API newer than 13.0, so a
+# 13.4-built wheel loads against the cu130 torch SGLang ships.
+set -ex
+
+if [ $# -lt 3 ]; then
+  echo "Usage: $0 <PYTHON_VERSION> <CUDA_VERSION> <FLASHMLA_SRC> [ARCH]"
+  exit 1
+fi
+
+PYTHON_VERSION="$1"
+CUDA_VERSION="$2"
+FLASHMLA_SRC="$(cd "$3" && pwd)"
+ARCH="${4:-$(uname -i)}"
+
+case "${CUDA_VERSION}" in
+  13.4) CU_TAG=cu130 ;;
+  *)
+    echo "Unsupported CUDA_VERSION: ${CUDA_VERSION} (the kernels require 13.4)" >&2
+    exit 1
+    ;;
+esac
+
+if [ "${ARCH}" = "aarch64" ]; then
+  BASE_IMG="pytorch/manylinuxaarch64-builder"
+else
+  BASE_IMG="pytorch/manylinux2_28-builder"
+fi
+
+PY_TAG="cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}"
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+DOCKERFILE="${REPO_ROOT}/docker/sgl-flash-mla.Dockerfile"
+RENAME_SCRIPT="${SCRIPT_DIR}/rename_sgl_flash_mla_whl.sh"
+
+DEPS_TAG="sgl-flash-mla-deps:cuda${CUDA_VERSION}-${PY_TAG}-${ARCH}"
+
+echo "----------------------------------------"
+echo "PYTHON_VERSION: ${PYTHON_VERSION}"
+echo "CUDA_VERSION:   ${CUDA_VERSION}"
+echo "CU_TAG:         ${CU_TAG}"
+echo "ARCH:           ${ARCH}"
+echo "BASE_IMG:       ${BASE_IMG}"
+echo "FLASHMLA_SRC:   ${FLASHMLA_SRC}"
+echo "----------------------------------------"
+
+docker build \
+  -f "${DOCKERFILE}" "$(dirname "${DOCKERFILE}")" \
+  --build-arg BASE_IMG="${BASE_IMG}" \
+  --build-arg CUDA_VERSION="${CUDA_VERSION}" \
+  --build-arg ARCH="${ARCH}" \
+  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" \
+  --build-arg PYTHON_TAG="${PY_TAG}" \
+  -t "${DEPS_TAG}" \
+  --network=host
+
+mkdir -p "${FLASHMLA_SRC}/dist"
+
+# 1) Build. --no-build-isolation because setup.py imports torch, which the image
+#    provides; FlashMLA declares no PEP 518 build requirements.
+#    NVCC_THREADS defaults to 32 in setup.py; MAX_JOBS is what actually
+#    parallelises, and without ninja on PATH torch silently ignores it.
+docker run --rm \
+  --network=host \
+  -v "${FLASHMLA_SRC}:/flashmla" \
+  -w /flashmla \
+  -e MAX_JOBS="${MAX_JOBS:-$(nproc)}" \
+  "${DEPS_TAG}" \
+  bash -c 'set -eux; python3 setup.py bdist_wheel'
+
+# 2) Retag inside the same image, which has a working pip/wheel and can rewrite
+#    the root-owned files the build container just wrote.
+docker run --rm \
+  -v "${FLASHMLA_SRC}:/flashmla" \
+  -v "${RENAME_SCRIPT}:/rename_sgl_flash_mla_whl.sh:ro" \
+  -w /flashmla \
+  "${DEPS_TAG}" \
+  bash /rename_sgl_flash_mla_whl.sh dist "${CU_TAG}" "${ARCH}"
+
+# 3) Sibling dist-pypi/ with the local version stripped.
+docker run --rm \
+  -v "${FLASHMLA_SRC}:/flashmla" \
+  -w /flashmla \
+  "${DEPS_TAG}" \
+  bash -c '
+set -eux
+mkdir -p dist-pypi
+for w in dist/*.whl; do
+  tmp=$(mktemp -d)
+  python3 -m wheel unpack "$w" --dest "$tmp"
+  unpacked=$(find "$tmp" -mindepth 1 -maxdepth 1 -type d | head -1)
+  info=$(find "$unpacked" -maxdepth 1 -type d -name "*.dist-info" | head -1)
+  meta="$info/METADATA"
+  orig=$(grep "^Version:" "$meta" | head -1 | sed "s/^Version:[[:space:]]*//")
+  new=$(echo "$orig" | sed "s/+cu[0-9]\+$//")
+  if [ "$orig" != "$new" ]; then
+    sed -i "s/^Version:.*/Version: ${new}/" "$meta"
+    old_base=$(basename "$info")
+    new_base="${old_base/${orig}/${new}}"
+    mv "$info" "$(dirname "$info")/${new_base}"
+  fi
+  python3 -m wheel pack "$unpacked" --dest-dir dist-pypi
+  rm -rf "$tmp"
+done
+ls -lh dist-pypi/
+'
+
+echo "Wheels in ${FLASHMLA_SRC}/dist:"
+ls -lh "${FLASHMLA_SRC}/dist"/*.whl 2>/dev/null || true
+echo "PyPI-ready wheels in ${FLASHMLA_SRC}/dist-pypi:"
+ls -lh "${FLASHMLA_SRC}/dist-pypi"/*.whl 2>/dev/null || true

scripts/rename_sgl_flash_mla_whl.sh

@@ -0,0 +1,75 @@
+#!/usr/bin/env bash
+# Retag a freshly-built sgl-flash-mla wheel for the sgl-whl index.
+#
+# setuptools tags the extension wheel `cp310-abi3-linux_x86_64`. `linux_x86_64`
+# is not installable from an index -- pip and PyPI both reject it -- so the
+# platform becomes manylinux2014_<arch>, and the version gains the +cuXXX local
+# segment the index keys on.
+#
+# The abi3 tag is left alone: the extensions register through
+# STABLE_TORCH_LIBRARY and need no CPython API beyond the Limited API, so one
+# wheel really does serve CPython 3.10+.
+#
+# Input:  dist/sgl_flash_mla-<VERSION>-cp310-abi3-linux_<ARCH>.whl
+# Output: dist/sgl_flash_mla-<VERSION>+<CU_TAG>-cp310-abi3-manylinux2014_<ARCH>.whl
+#
+# Usage: rename_sgl_flash_mla_whl.sh <WHEEL_DIR> <CU_TAG> <ARCH>
+set -ex
+
+if [ $# -lt 3 ]; then
+  echo "Usage: $0 <WHEEL_DIR> <CU_TAG> <ARCH>"
+  exit 1
+fi
+
+WHEEL_DIR="$1"
+CU_TAG="$2"
+ARCH="$3"
+PLAT_TAG="manylinux2014_${ARCH}"
+
+PYTHON="${PYTHON:-python3}"
+"${PYTHON}" -m pip install --quiet wheel
+
+shopt -s nullglob
+wheel_files=("${WHEEL_DIR}"/sgl_flash_mla-*.whl)
+if [ ${#wheel_files[@]} -eq 0 ]; then
+  echo "No sgl_flash_mla wheel found under ${WHEEL_DIR}" >&2
+  exit 1
+fi
+
+for wheel in "${wheel_files[@]}"; do
+  TMPDIR=$(mktemp -d)
+  trap 'rm -rf -- "$TMPDIR"' ERR
+
+  "${PYTHON}" -m wheel unpack "$wheel" --dest "$TMPDIR"
+  UNPACKED=$(find "$TMPDIR" -mindepth 1 -maxdepth 1 -type d | head -1)
+  DIST_INFO=$(find "$UNPACKED" -maxdepth 1 -type d -name "*.dist-info" | head -1)
+  WHEEL_META="${DIST_INFO}/WHEEL"
+  METADATA_FILE="${DIST_INFO}/METADATA"
+
+  # Only the platform part of the tag changes; the interpreter/ABI part stays abi3.
+  sed -i -E "s/^(Tag: [^-]+-[^-]+)-linux_${ARCH}$/\1-${PLAT_TAG}/" "$WHEEL_META"
+  grep -q "^Tag: .*-${PLAT_TAG}$" "$WHEEL_META" || {
+    echo "Failed to retag ${wheel}; WHEEL says:" >&2
+    cat "$WHEEL_META" >&2
+    exit 1
+  }
+
+  ORIG_VERSION=$(grep '^Version:' "$METADATA_FILE" | head -1 | sed 's/^Version:[[:space:]]*//')
+  if [[ "$ORIG_VERSION" == *"+${CU_TAG}"* ]]; then
+    NEW_VERSION="$ORIG_VERSION"
+  else
+    NEW_VERSION="${ORIG_VERSION}+${CU_TAG}"
+    sed -i "s/^Version:.*/Version: ${NEW_VERSION}/" "$METADATA_FILE"
+    OLD_BASE=$(basename "$DIST_INFO")
+    NEW_BASE="${OLD_BASE/${ORIG_VERSION}/${NEW_VERSION}}"
+    mv "$DIST_INFO" "${UNPACKED}/${NEW_BASE}"
+  fi
+
+  rm -f "$wheel"
+  "${PYTHON}" -m wheel pack "$UNPACKED" --dest-dir "$WHEEL_DIR"
+  rm -rf "$TMPDIR"
+  trap - ERR
+done
+
+echo "Renamed wheels in ${WHEEL_DIR}:"
+ls -lh "${WHEEL_DIR}"/*.whl

scripts/update_flashmla_whl_index.py

@@ -0,0 +1,44 @@
+# Generates a PEP 503 simple index for sgl-flash-mla wheels under
+# sgl-whl/cu<version>/sgl-flash-mla/index.html. Mirrors update_deepgemm_whl_index.py
+# so consumers can `pip install sgl-flash-mla --extra-index-url https://...whl/cu130`.
+
+import argparse
+import hashlib
+import pathlib
+import re
+
+SUPPORTED_CUDA_VERSIONS = ["130"]
+
+
+def update_wheel_index(cuda_version, wheel_dir):
+    index_dir = pathlib.Path(f"sgl-whl/cu{cuda_version}/sgl-flash-mla")
+    index_dir.mkdir(exist_ok=True, parents=True)
+    base_url = "https://github.com/sgl-project/whl/releases/download"
+
+    suffix = f"+cu{cuda_version}"
+    for path in sorted(pathlib.Path(wheel_dir).glob("*.whl")):
+        if suffix not in path.name:
+            continue
+        with open(path, "rb") as f:
+            sha256 = hashlib.sha256(f.read()).hexdigest()
+        match = re.match(r"sgl_flash_mla-([0-9][^-+]*)(?:\+cu[0-9]+)?-", path.name)
+        if not match:
+            continue
+        ver = match.group(1)
+        full_url = f"{base_url}/v{ver}/{path.name}#sha256={sha256}"
+        with (index_dir / "index.html").open("a") as f:
+            f.write(f'<a href="{full_url}">{path.name}</a><br>\n')
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--cuda", type=str, required=True, choices=SUPPORTED_CUDA_VERSIONS
+    )
+    parser.add_argument("--wheel-dir", type=str, default="dist")
+    args = parser.parse_args()
+    update_wheel_index(args.cuda, args.wheel_dir)
+
+
+if __name__ == "__main__":
+    main()