build(kernel-cache): add sm_80 (A100/A800) to default arches (#75)

Signed-off-by: jasonfan <36730318+jason-fxz@users.noreply.github.com>
This commit is contained in:
jasonfan
2026-08-22 23:38:42 -07:00
parent f0abe587a1
commit 184a4f114d
2 changed files with 12 additions and 8 deletions
+11 -7
View File
@@ -100,15 +100,19 @@ def _build_jit_cache() -> None:
"yes", "yes",
"on", "on",
} }
# Multi-arch fatbin: build a SASS cubin for every target arch so the wheel runs driver-only # Multi-arch fatbin: one SASS cubin per listed arch, plus the PTX of the highest one.
# on any of them (no per-GPU JIT / nvcc). tvm-ffi reads TVM_FFI_CUDA_ARCH_LIST; freetoken's # A GPU whose arch is not listed and is below the highest one has no usable image.
# _cuda_cflags adds the top arch's PTX for forward-compat to newer GPUs. Default covers Ampere # 8.0 -> A100, A800, A30 (Ampere, datacenter)
# consumer (8.6), Ada / 40xx (8.9), Hopper (9.0), Blackwell datacenter (10.0) + consumer / # 8.6 -> RTX 30 series, A10, A40 (Ampere, consumer / workstation)
# 50xx (12.0). Override the set with FREETOKEN_KERNEL_CACHE_ARCHES (space-separated maj.min), # 8.9 -> RTX 40 series, L4, L40, RTX 6000 Ada (Ada Lovelace)
# or TVM_FFI_CUDA_ARCH_LIST directly. Needs an nvcc that supports every listed arch. # 9.0 -> H100, H800, H20 (Hopper)
# 10.0 -> B200, B100, GB200 (Blackwell, datacenter)
# 12.0 -> RTX 50 series, RTX PRO 6000 Blackwell (Blackwell, consumer / workstation)
# Override 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: if "TVM_FFI_CUDA_ARCH_LIST" not in os.environ:
os.environ["TVM_FFI_CUDA_ARCH_LIST"] = os.getenv( 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", "8.0 8.6 8.9 9.0 10.0 12.0"
) )
compile_and_package_kernels( compile_and_package_kernels(
out_dir=out_dir, out_dir=out_dir,
+1 -1
View File
@@ -60,7 +60,7 @@ warn_arch_override() {
if [[ -n "${!var:-}" ]]; then if [[ -n "${!var:-}" ]]; then
warn "############################################################" warn "############################################################"
warn "$var='${!var}' is set in this shell and OVERRIDES the" 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 (8.0 8.6 8.9 9.0 10.0 12.0, see"
warn "freetoken-kernel-cache/build_backend.py). The kernel-cache" warn "freetoken-kernel-cache/build_backend.py). The kernel-cache"
warn "wheel will only carry SASS for the listed archs — do NOT" warn "wheel will only carry SASS for the listed archs — do NOT"
warn "release it unless the narrowing is intentional." warn "release it unless the narrowing is intentional."