Skip to content

ci: add 12.1a to aarch64 jit-cache wheel arch targets #3907

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
yichengj0 wants to merge 3 commits into from
Closed

ci: add 12.1a to aarch64 jit-cache wheel arch targets #3907

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8f258df
ci: add 12.1a to aarch64 jit-cache wheel arch targets
yichengj0 Jul 10, 2026
ff797a9
docs: note the example arch list needs CUDA 13.0
yichengj0 Jul 10, 2026
fb7e206
Merge branch 'main' into wheel-arch-121a
dierksen Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/nightly-release.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ jobs:
- name: Build wheel in container
env:
DOCKER_IMAGE: ${{ matrix.arch == 'aarch64' && format('pytorch/manylinuxaarch64-builder:cuda{0}', matrix.cuda) || format('pytorch/manylinux2_28-builder:cuda{0}', matrix.cuda) }}
FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '12.9' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || (matrix.cuda < '13.0' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f' || (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f')) }}
FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '12.9' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || (matrix.cuda < '13.0' && (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f 12.1a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f') || (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f 12.1a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f')) }}
FLASHINFER_DEV_RELEASE_SUFFIX: ${{ needs.setup.outputs.dev_suffix }}
run: |
# Extract CUDA major and minor versions
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/release.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@ jobs:
- name: Build wheel in container
env:
DOCKER_IMAGE: ${{ matrix.arch == 'aarch64' && format('pytorch/manylinuxaarch64-builder:cuda{0}', matrix.cuda) || format('pytorch/manylinux2_28-builder:cuda{0}', matrix.cuda) }}
FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '12.9' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || (matrix.cuda < '13.0' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f' || (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f')) }}
FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '12.9' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || (matrix.cuda < '13.0' && (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f 12.1a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f') || (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f 12.1a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f')) }}
run: |
# Extract CUDA major and minor versions
CUDA_MAJOR=$(echo "${{ matrix.cuda }}" | cut -d'.' -f1)
Expand Down
3 changes: 2 additions & 1 deletion CLAUDE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -585,7 +585,8 @@ When ready to distribute:
```bash
# Build flashinfer-jit-cache package
cd flashinfer-jit-cache
export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 11.0a 12.0f"
# The full list below needs CUDA 13.0; trim it for older toolkits.
export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f 12.1a"
python -m build --no-isolation --wheel
```

Expand Down
3 changes: 2 additions & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,8 @@ python -m pip install dist/*.whl

```bash
# flashinfer-jit-cache (customize for your target GPUs)
export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f"
# The full list below needs CUDA 13.0; trim it for older toolkits.
export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f 12.1a"
cd flashinfer-jit-cache
python -m build --no-isolation --wheel
python -m pip install dist/*.whl
Expand Down
4 changes: 2 additions & 2 deletions docs/installation.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,11 +105,11 @@ You can follow the steps below to install FlashInfer from source code:
python -m build --no-isolation --wheel
python -m pip install dist/*.whl

Build ``flashinfer-jit-cache`` (customize ``FLASHINFER_CUDA_ARCH_LIST`` for your target GPUs):
Build ``flashinfer-jit-cache`` (customize ``FLASHINFER_CUDA_ARCH_LIST`` for your target GPUs; the full list below needs CUDA 13.0, trim it for older toolkits):

.. code-block:: bash

export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f"
export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f 12.1a"
cd flashinfer-jit-cache
python -m build --no-isolation --wheel
python -m pip install dist/*.whl
Expand Down
9 changes: 7 additions & 2 deletions scripts/task_test_jit_cache_package_build_import.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,17 +143,22 @@ arches = ["7.5", "8.0", "8.9", "9.0a"]
if cuda_ver is not None:
try:
major, minor = map(int, cuda_ver.split(".")[:2])
machine = (os.environ.get("ARCH") or platform.machine()).lower()
is_aarch64 = machine in ("aarch64", "arm64")
if (major, minor) >= (13, 0):
arches.append("10.0a")
arches.append("10.3a")
machine = (os.environ.get("ARCH") or platform.machine()).lower()
if machine in ("aarch64", "arm64"):
if is_aarch64:
arches.append("11.0a")
arches.append("12.0f")
if is_aarch64:
arches.append("12.1a")
elif (major, minor) >= (12, 9):
arches.append("10.0a")
arches.append("10.3a")
arches.append("12.0f")
if is_aarch64:
arches.append("12.1a")
Comment on lines 148 to +161

@gemini-code-assist gemini-code-assist Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic blocks for CUDA >= 13.0 and >= 12.9 are highly redundant, differing only by the conditional addition of 11.0a on aarch64 for CUDA >= 13.0. Consolidating these blocks reduces duplication and improves maintainability while preserving the exact order of the architectures.

Suggested change
if (major, minor) >= (13, 0):
arches.append("10.0a")
arches.append("10.3a")
machine = (os.environ.get("ARCH") or platform.machine()).lower()
if machine in ("aarch64", "arm64"):
if is_aarch64:
arches.append("11.0a")
arches.append("12.0f")
if is_aarch64:
arches.append("12.1a")
elif (major, minor) >= (12, 9):
arches.append("10.0a")
arches.append("10.3a")
arches.append("12.0f")
if is_aarch64:
arches.append("12.1a")
if (major, minor) >= (12, 9):
arches.append("10.0a")
arches.append("10.3a")
if (major, minor) >= (13, 0) and is_aarch64:
arches.append("11.0a")
arches.append("12.0f")
if is_aarch64:
arches.append("12.1a")

@yichengj0 yichengj0 Jul 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks — the suggestion is output-equivalent, but keeping the branches separate is intentional: the three-branch if/elif maps one-to-one onto the workflow ternary's branches in release.yml/nightly-release.yml (12.8 / 12.9 / 13.0), so each branch reads as that CUDA version's complete list and can be eyeball-diffed against the workflow string directly. Merging 12.9 and 13.0 would save three lines but lose that mapping, so we'll keep the current structure.

elif (major, minor) >= (12, 8):
arches.append("10.0a")
arches.append("12.0a")
Expand Down
Loading