Skip to content

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

Open
Chessing234 wants to merge 3 commits into
state-spaces:mainfrom
Chessing234:fix/setup-cuda-home-guards
Open

fix: harden setup.py CUDA build guards and wheel fallback messaging#1004
Chessing234 wants to merge 3 commits into
state-spaces:mainfrom
Chessing234:fix/setup-cuda-home-guards

Conversation

@Chessing234

Copy link
Copy Markdown
Contributor

Summary

  • Fix UnboundLocalError when MAMBA_KEEP_CUDA_BUILD=TRUE but CUDA_HOME/nvcc is missing by falling back to torch.version.cuda for arch flags.
  • Fail clearly when guessing CUDA wheel URLs on CPU-only PyTorch instead of opaque packaging errors.
  • Print torch/Python/platform context when prebuilt wheels 404, noting that selective_scan_cuda source builds are slow and optional for Mamba-2/3.
  • Extend the missing-nvcc warning with Colab/Kaggle guidance.

Test plan

  • Reviewed setup.py control flow for CUDA_HOME is None and torch.version.cuda is None paths
  • CI: tagged release wheel build (maintainer)

Made with Cursor

MAMBA_KEEP_CUDA_BUILD referenced bare_metal_version even when nvcc was
unavailable, causing UnboundLocalError. Fall back to torch.version.cuda
for arch flags and only compare CUDA majors when nvcc is present.
get_wheel_url() called parse(None) when torch.version.cuda was unset,
producing an opaque packaging error. Raise a direct message instead.
Extend the CUDA_HOME warning so cloud notebook users know missing nvcc is
expected unless they explicitly requested selective_scan_cuda.
@Chessing234
Chessing234 force-pushed the fix/setup-cuda-home-guards branch from e76cde4 to fafe5fb Compare August 5, 2026 10:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant