From 35668da6db1e23111f1fbcc8a552455927a70f2a Mon Sep 17 00:00:00 2001 From: Lukas Rakauskas Date: Sat, 22 Aug 2026 11:56:49 +0300 Subject: [PATCH] fix(cuda): support Turing GPUs --- freetoken-kernel-cache/build_backend.py | 8 ++++---- python/freetoken/kernel/backend.py | 21 ++++++++++++++++++--- python/freetoken/kernel/triton/attention.py | 6 ++++++ python/freetoken/kernel/utils.py | 2 +- scripts/build-release-wheels.sh | 2 +- tests/kernels/test_backend.py | 18 ++++++++++++++++++ tests/kernels/test_triton_attention.py | 3 +++ 7 files changed, 51 insertions(+), 9 deletions(-) create mode 100644 tests/kernels/test_backend.py diff --git a/freetoken-kernel-cache/build_backend.py b/freetoken-kernel-cache/build_backend.py index ff2f09571..d03545418 100644 --- a/freetoken-kernel-cache/build_backend.py +++ b/freetoken-kernel-cache/build_backend.py @@ -102,13 +102,13 @@ def _build_jit_cache() -> None: } # Multi-arch fatbin: build a SASS cubin for every target arch so the wheel runs driver-only # on any of them (no per-GPU JIT / nvcc). tvm-ffi reads TVM_FFI_CUDA_ARCH_LIST; freetoken's - # _cuda_cflags adds the top arch's PTX for forward-compat to newer GPUs. Default covers Ampere - # consumer (8.6), Ada / 40xx (8.9), Hopper (9.0), Blackwell datacenter (10.0) + consumer / - # 50xx (12.0). Override the set with FREETOKEN_KERNEL_CACHE_ARCHES (space-separated maj.min), + # _cuda_cflags adds the top arch's PTX for forward-compat to newer GPUs. Default covers Turing + # (7.5), Ampere consumer (8.6), Ada / 40xx (8.9), Hopper (9.0), Blackwell datacenter (10.0) + # + consumer / 50xx (12.0). Override the set with FREETOKEN_KERNEL_CACHE_ARCHES (space-separated maj.min), # or TVM_FFI_CUDA_ARCH_LIST directly. Needs an nvcc that supports every listed arch. if "TVM_FFI_CUDA_ARCH_LIST" not in os.environ: os.environ["TVM_FFI_CUDA_ARCH_LIST"] = os.getenv( - "FREETOKEN_KERNEL_CACHE_ARCHES", "8.6 8.9 9.0 10.0 12.0" + "FREETOKEN_KERNEL_CACHE_ARCHES", "7.5 8.6 8.9 9.0 10.0 12.0" ) compile_and_package_kernels( out_dir=out_dir, diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 3037ad8d7..0a6027767 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -21,14 +21,29 @@ def _importable(name: str) -> bool: return False +def _supports_sm80() -> bool: + """Whether this GPU can run the optional native kernel wheels. + + The sglang and triton-kernels wheels used by FreeToken carry Ampere-and-newer + kernels. On older GPUs, selecting them by importability causes a late CUDA + launch failure instead of using FreeToken's fallbacks. + """ + try: + import torch + + return not torch.cuda.is_available() or torch.cuda.get_device_capability() >= (8, 0) + except Exception: + return True + + @functools.cache def is_flashinfer_installed() -> bool: - return _importable("flashinfer") + return _supports_sm80() and _importable("flashinfer") @functools.cache def is_sgl_kernel_installed() -> bool: - return _importable("sgl_kernel") + return _supports_sm80() and _importable("sgl_kernel") @functools.cache @@ -39,7 +54,7 @@ def is_triton_kernels_installed() -> bool: source tree and has no Windows wheel. It is also not one of the six ops ``freetoken.kernel.triton`` reimplements, so its call-site carries its own fallback. """ - return _importable("triton_kernels") + return _supports_sm80() and _importable("triton_kernels") @functools.cache diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c2358d84f..ef4e0fe62 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -28,6 +28,12 @@ def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[i A100/H100); shrink only where it does not. ``smem_optin == 0`` (unknown) conservatively selects the small tiles, i.e. the prior consumer-safe behavior. """ + # Turing exposes only 64 KiB of opt-in shared memory. Triton needs more than + # the q/k/v tile estimate, so use a smaller tile instead of compiling a kernel + # that the device cannot launch. + if 0 < smem_optin <= 65536: + return (32, 16) if head_dim <= 256 else (16, 16) + budget = smem_optin * 0.8 # headroom for scores/acc/alignment/triton scratch def fits(block_m: int, block_n: int) -> bool: diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b59..fb4f448b4 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -24,7 +24,7 @@ def _cuda_cflags(extra: List[str]) -> List[str]: """CUDA nvcc flags for a kernel build. During the multi-arch AOT cache build, - `TVM_FFI_CUDA_ARCH_LIST` (e.g. "8.6 8.9 9.0 10.0 12.0") makes tvm-ffi emit a SASS cubin + `TVM_FFI_CUDA_ARCH_LIST` (e.g. "7.5 8.6 8.9 9.0 10.0 12.0") makes tvm-ffi emit a SASS cubin (`-gencode ...code=sm_XX`) for each listed arch — but NO PTX. We add the PTX of the HIGHEST listed arch so a GPU newer than any listed one (no matching SASS) still runs via the driver's PTX→SASS JIT (driver-only, no CUDA toolkit). One top PTX suffices: the loader always diff --git a/scripts/build-release-wheels.sh b/scripts/build-release-wheels.sh index b23a71f11..739b193ad 100755 --- a/scripts/build-release-wheels.sh +++ b/scripts/build-release-wheels.sh @@ -60,7 +60,7 @@ warn_arch_override() { if [[ -n "${!var:-}" ]]; then warn "############################################################" warn "$var='${!var}' is set in this shell and OVERRIDES the" - warn "default multi-arch list (8.6 8.9 9.0 10.0 12.0, see" + warn "default multi-arch list (7.5 8.6 8.9 9.0 10.0 12.0, see" warn "freetoken-kernel-cache/build_backend.py). The kernel-cache" warn "wheel will only carry SASS for the listed archs — do NOT" warn "release it unless the narrowing is intentional." diff --git a/tests/kernels/test_backend.py b/tests/kernels/test_backend.py new file mode 100644 index 000000000..e929fecad --- /dev/null +++ b/tests/kernels/test_backend.py @@ -0,0 +1,18 @@ +def test_native_wheels_are_disabled_on_turing(monkeypatch): + import torch + from freetoken.kernel import backend + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (7, 5)) + backend.is_flashinfer_installed.cache_clear() + backend.is_sgl_kernel_installed.cache_clear() + backend.is_triton_kernels_installed.cache_clear() + + try: + assert not backend.is_flashinfer_installed() + assert not backend.is_sgl_kernel_installed() + assert not backend.is_triton_kernels_installed() + finally: + backend.is_flashinfer_installed.cache_clear() + backend.is_sgl_kernel_installed.cache_clear() + backend.is_triton_kernels_installed.cache_clear() diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index 6f4afca9e..5f04118f0 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -600,6 +600,9 @@ def test_extend_triton_attention_with_sinks_matches_reference(use_split_inputs: # consumer opt-in smem (sm_89 ~99KB): shrink once head_dim >= 256 (256, 101376, (64, 32)), (512, 101376, (16, 16)), + # Turing exposes 64 KiB; use a tile that fits Triton's larger allocation. + (256, 65536, (32, 16)), + (512, 65536, (16, 16)), # unknown budget -> conservative small tiles (prior consumer-safe behavior) (256, 0, (64, 32)), (512, 0, (16, 16)),