Skip to content

fix: harden setup.py CUDA build guards and wheel fallback messaging #1004

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

Open
Chessing234 wants to merge 3 commits into
base: main
Choose a base branch
from
+33 −9
Open

fix: harden setup.py CUDA build guards and wheel fallback messaging #1004

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
41bc5dc
fix: guard bare_metal_version when CUDA_HOME is missing
Chessing234 Jul 31, 2026
f88fe1c
fix: fail clearly when guessing CUDA wheels on CPU-only torch
Chessing234 Jul 31, 2026
fafe5fb
docs: mention Colab/Kaggle when nvcc is missing during CUDA build
Chessing234 Jul 31, 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
42 changes: 33 additions & 9 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,9 +116,11 @@ def check_if_cuda_home_none(global_option: str) -> None:
# warn instead of error because user could be downloading prebuilt wheels, so nvcc won't be necessary
# in that case.
warnings.warn(
f"{global_option} was requested, but nvcc was not found. Are you sure your environment has nvcc available? "
f"{global_option} was requested, but nvcc was not found. Are you sure your environment has nvcc available? "
"If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, "
"only images whose names contain 'devel' will provide nvcc."
"only images whose names contain 'devel' will provide nvcc. "
"Cloud notebooks (Colab/Kaggle) often lack nvcc unless you opt into MAMBA_KEEP_CUDA_BUILD "
"and compile selective_scan_cuda locally."
)


Expand Down Expand Up @@ -164,6 +166,7 @@ def append_nvcc_threads(nvcc_extra_args):
check_if_cuda_home_none(PACKAGE_NAME)
# Check, if CUDA11 is installed for compute capability 8.0

bare_metal_version = None
if CUDA_HOME is not None:
_, bare_metal_version = get_cuda_bare_metal_version(CUDA_HOME)
if bare_metal_version < Version("11.6"):
Expand All @@ -172,10 +175,17 @@ def append_nvcc_threads(nvcc_extra_args):
"Note: make sure nvcc has a supported version by running nvcc -V."
)

if torch.version.cuda is None:
raise RuntimeError(
f"{PACKAGE_NAME} CUDA build requires a CUDA-enabled PyTorch build, "
f"but torch.version.cuda is None (torch {torch.__version__})."
)
torch_cuda_version = parse(torch.version.cuda)
cuda_version_for_arch = bare_metal_version if bare_metal_version is not None else torch_cuda_version

# If system CUDA and PyTorch CUDA have different major versions,
# clear TORCH_CUDA_ARCH_LIST to prevent cpp_extension from erroring
torch_cuda_version = parse(torch.version.cuda)
if bare_metal_version.major != torch_cuda_version.major:
if bare_metal_version is not None and bare_metal_version.major != torch_cuda_version.major:
os.environ["TORCH_CUDA_ARCH_LIST"] = ""

cc_flag.append("-gencode")
Expand All @@ -184,15 +194,15 @@ def append_nvcc_threads(nvcc_extra_args):
cc_flag.append("arch=compute_80,code=sm_80")
cc_flag.append("-gencode")
cc_flag.append("arch=compute_87,code=sm_87")
if bare_metal_version >= Version("11.8"):
if cuda_version_for_arch >= Version("11.8"):
cc_flag.append("-gencode")
cc_flag.append("arch=compute_90,code=sm_90")
if bare_metal_version >= Version("12.8"):
if cuda_version_for_arch >= Version("12.8"):
cc_flag.append("-gencode")
cc_flag.append("arch=compute_100,code=sm_100")
cc_flag.append("-gencode")
cc_flag.append("arch=compute_120,code=sm_120")
if bare_metal_version >= Version("13.0"):
if cuda_version_for_arch >= Version("13.0"):
cc_flag.append("-gencode")
cc_flag.append("arch=compute_103,code=sm_103")
cc_flag.append("-gencode")
Expand Down Expand Up @@ -287,6 +297,11 @@ def get_wheel_url():
else:
# We're using the CUDA version used to build torch, not the one currently installed
# _, cuda_version_raw = get_cuda_bare_metal_version(CUDA_HOME)
if torch.version.cuda is None:
raise ValueError(
"Cannot guess a CUDA wheel URL without torch.version.cuda set; "
"install a CUDA-enabled PyTorch build or set MAMBA_FORCE_BUILD=TRUE."
)
torch_cuda_version = parse(torch.version.cuda)
# For CUDA 11, we only compile for CUDA 11.8, and for CUDA 12 we only compile for CUDA 12.3
# to save CI time. Minor versions should be compatible.
Expand Down Expand Up @@ -355,8 +370,17 @@ def run(self):
print("Raw wheel path", wheel_path)
shutil.move(wheel_filename, wheel_path)
except urllib.error.HTTPError:
print("Precompiled wheel not found. Building from source...")
# If the wheel could not be downloaded, build from source
print(
"Precompiled wheel not found for "
f"torch {torch.__version__}, Python {sys.version_info.major}.{sys.version_info.minor}, "
f"platform {get_platform()}. Building selective_scan_cuda from source..."
)
print(
"Source compilation can take 15–30+ minutes with little or no pip output. "
"Prebuilt wheels are published on GitHub releases for select torch/CUDA/Python "
"combinations; see README Installation. If you only need Mamba-2/3, omit "
"MAMBA_KEEP_CUDA_BUILD to skip this CUDA extension."
)
super().run()

setup(
Expand Down