Skip to content
Merged
Changes from all commits
Commits
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
22 changes: 7 additions & 15 deletions flashinfer/compilation_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,28 +36,20 @@ def _normalize_cuda_arch(major: int, minor: int) -> tuple[int, str]:
tuple with the correct architecture suffix for nvcc.

SM 9.x -> 'a' suffix (e.g. compute_90a)
SM 12.x -> always normalized to SM 120 with 'f' suffix (e.g. compute_120f)
when the installed CUDA toolchain supports it (CUDA >= 13.0),
otherwise 'a'. This covers both SM 12.0 and SM 12.1 (DGX Spark).
SM 12.x -> always normalized to SM 120 with 'f' suffix (e.g. compute_120f).
This covers both SM 12.0 and SM 12.1 (DGX Spark) when the installed CUDA toolchain supports it (CUDA >= 12.9).
SM 10+ -> 'a' suffix (e.g. compute_100a)
SM < 9 -> no suffix
"""
if major == 9:
return (major, str(minor) + "a")
elif major == 12:
try:
from flashinfer.jit.cpp_ext import is_cuda_version_at_least
from flashinfer.jit.cpp_ext import is_cuda_version_at_least

if is_cuda_version_at_least("13.0"):
return (major, "0f")
except (ImportError, RuntimeError, ValueError):
logger.debug(
"Could not determine CUDA version; "
"falling back to 'a' suffix for SM %d.%d",
major,
minor,
)
return (major, "0a")
if is_cuda_version_at_least("12.9"):
Comment thread
kahyunnam marked this conversation as resolved.
return (major, "0f")
else:
raise RuntimeError("SM 12.x requires CUDA >= 12.9")
elif major >= 10:
return (major, str(minor) + "a")
return (major, str(minor))
Expand Down
Loading