Skip to content

[NVIDIA] Thor & Spark Support #2028

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

Merged
merged 20 commits into from
Nov 13, 2025
+17 −7
Merged

[NVIDIA] Thor & Spark Support #2028

Show file tree
Hide file tree
Changes from 19 commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
705d15a
Thor & Spark Support
johnnynunez Nov 3, 2025
5eea497
Merge branch 'flashinfer-ai:main' into main
johnnynunez Nov 8, 2025
16ed2c3
Merge branch 'flashinfer-ai:main' into main
johnnynunez Nov 11, 2025
4958688
fix xqa compilation
yzh119 Nov 11, 2025
9e5a259
Merge branch 'main' of https://github.com/johnnynunez/flashinfer into…
yzh119 Nov 11, 2025
6f35f3a
Thor & Spark Support
johnnynunez Nov 11, 2025
3b6d0f6
Update release.yml
johnnynunez Nov 11, 2025
0769cf4
Update task_test_jit_cache_package_build_import.sh
johnnynunez Nov 11, 2025
b030292
revert
johnnynunez Nov 11, 2025
ab28a10
Merge remote-tracking branch 'origin/main'
johnnynunez Nov 11, 2025
563d7fd
revert
johnnynunez Nov 11, 2025
dd428d8
Update task_test_jit_cache_package_build_import.sh
johnnynunez Nov 11, 2025
77108b5
Merge branch 'main' of https://github.com/johnnynunez/flashinfer into…
yzh119 Nov 11, 2025
43a0d3a
fix xqa thor compilation
yzh119 Nov 12, 2025
afeb8dd
thor is only in cuda 13
johnnynunez Nov 12, 2025
a4609fa
Merge remote-tracking branch 'origin/main'
johnnynunez Nov 12, 2025
cc1d2cc
thor's __CUDA_ARCH__ is 1010 instead of 1100
yzh119 Nov 12, 2025
e3836d2
Merge branch 'main' of https://github.com/johnnynunez/flashinfer into…
yzh119 Nov 12, 2025
ad31661
Merge branch 'flashinfer-ai:main' into main
johnnynunez Nov 12, 2025
bdf2ed8
revert 1010 to 1100
yzh119 Nov 13, 2025
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 @@ -145,7 +145,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.8' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0a' }}
FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '13.0' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0f 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.8' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0a' }}
FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '13.0' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0f 12.0f' }}
run: |
# Extract CUDA major and minor versions
CUDA_MAJOR=$(echo "${{ matrix.cuda }}" | cut -d'.' -f1)
Expand Down
2 changes: 1 addition & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ python -m pip install dist/*.whl

`flashinfer-jit-cache` (customize `FLASHINFER_CUDA_ARCH_LIST` for your target GPUs):
```bash
export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 10.0a 10.3a 12.0a"
export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 10.0a 10.3a 11.0f 12.0f"
cd flashinfer-jit-cache
python -m build --no-isolation --wheel
python -m pip install dist/*.whl
Expand Down
2 changes: 1 addition & 1 deletion csrc/xqa/mha.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ __constant__ constexpr uint32_t cacheVTileSeqLen = 32;
constexpr uint32_t preferedKHeadPartBytes = 64;
__constant__ constexpr uint32_t cacheVTileSeqLen = 32;
#elif __CUDA_ARCH__ == 800 || __CUDA_ARCH__ == 870 || __CUDA_ARCH__ == 900 || \
__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030
__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030 || __CUDA_ARCH__ == 1010
constexpr uint32_t preferedKHeadPartBytes = 128;
__constant__ constexpr uint32_t cacheVTileSeqLen = 64;
#else
Expand Down
3 changes: 2 additions & 1 deletion csrc/xqa/utils.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,8 @@ __constant__ constexpr float kE4M3_MAX = 448.F;
constexpr uint32_t kMAX_SMEM_SIZE = (99u << 10);
#elif __CUDA_ARCH__ == 800 || __CUDA_ARCH__ == 870
constexpr uint32_t kMAX_SMEM_SIZE = (163u << 10);
#elif __CUDA_ARCH__ == 900 || __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030
#elif __CUDA_ARCH__ == 900 || __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030 || \
__CUDA_ARCH__ == 1010
constexpr uint32_t kMAX_SMEM_SIZE = (227u << 10);
#endif
#endif
Expand Down
2 changes: 1 addition & 1 deletion docs/installation.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ You can follow the steps below to install FlashInfer from source code:

.. code-block:: bash

export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 10.0a 10.3a 12.0a"
export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 10.0a 10.3a 11.0f 12.0f"
cd flashinfer-jit-cache
python -m build --no-isolation --wheel
python -m pip install dist/*.whl
Expand Down
11 changes: 10 additions & 1 deletion 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 @@ -43,7 +43,16 @@ arches = ["7.5", "8.0", "8.9", "9.0a"]
if cuda_ver is not None:
try:
major, minor = map(int, cuda_ver.split(".")[:2])
if (major, minor) >= (12, 8):
if (major, minor) >= (13, 0):
arches.append("10.0a")
arches.append("10.3a")
arches.append("11.0f")
arches.append("12.0f")
elif (major, minor) >= (12, 9):
arches.append("10.0a")
arches.append("10.3a")
arches.append("12.0f")
elif (major, minor) >= (12, 8):
arches.append("10.0a")
arches.append("12.0a")
except Exception:
Expand Down