Make CUDA selective scan install explicit opt-in (#977)

* Made CUDA selective scan optional w/ MAMBA_KEEP_CUDA_BUILD

* Bumped python requirement from >=3.9 to >=3.10

---------

Co-authored-by: Caitlin Wang <123908556+caitWW@users.noreply.github.com>

Author: berlinchen7

.github/scripts/build.sh

@@ -12,8 +12,9 @@ export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY
 
 # Limit MAX_JOBS otherwise the github runner goes OOM
 export MAX_JOBS=2 
-export MAMBA_FORCE_BUILD="TRUE" 
-export MAMBA_FORCE_CXX11_ABI=$CXX11_ABI 
+export MAMBA_FORCE_BUILD="TRUE"
+export MAMBA_KEEP_CUDA_BUILD="TRUE"
+export MAMBA_FORCE_CXX11_ABI=$CXX11_ABI
 
 # 5h timeout since GH allows max 6h and we want some buffer
 EXIT_CODE=0

.github/workflows/publish.yaml

@@ -138,8 +138,6 @@ jobs:
           pip install torch --index-url https://download.pytorch.org/whl/cpu
 
       - name: Build core package
-        env:
-          MAMBA_SKIP_CUDA_BUILD: "TRUE"
         run: |
           python setup.py sdist --dist-dir=dist
 

README.md

@@ -25,19 +25,42 @@ with an efficient hardware-aware design and implementation in the spirit of [Fla
 
 ## Installation
 
-Install PyTorch first, then:
-- [Option] `pip install causal-conv1d>=1.4.0 --no-build-isolation`: an efficient implementation of a simple causal Conv1d layer used inside the Mamba block.
-- `pip install mamba-ssm --no-build-isolation`: the core Mamba package.
-- `pip install mamba-ssm[causal-conv1d] --no-build-isolation`: To install core Mamba package and causal-conv1d.
-
-`--no-build-isolation` is required so that pip uses your existing CUDA-enabled PyTorch instead of installing torch-cpu in an isolated build environment.
-
-NOTE: To use Mamba-3, please install from source `MAMBA_FORCE_BUILD=TRUE pip install --no-cache-dir --force-reinstall git+https://github.com/state-spaces/mamba.git --no-build-isolation`.
-
-Other requirements:
+Install PyTorch first. By default, `mamba-ssm` installs the core package without compiling the
+`selective_scan_cuda` extension and without downloading cached CUDA-enabled wheels.
+
+| Install mode | Command | Behavior |
+| --- | --- | --- |
+| Core package | `pip install mamba-ssm --no-build-isolation` | Installs `mamba-ssm` without `selective_scan_cuda` and does not compile CUDA extensions. |
+| Core package plus causal-conv1d | `pip install "mamba-ssm[causal-conv1d]" --no-build-isolation` | Installs the core package and the `causal-conv1d` extra, still without `selective_scan_cuda`. |
+| Force local core package build | `MAMBA_FORCE_BUILD=TRUE pip install mamba-ssm --no-build-isolation` | Builds the default pure Python wheel locally, still without `selective_scan_cuda`. |
+| CUDA selective scan opt-in | `MAMBA_KEEP_CUDA_BUILD=TRUE pip install mamba-ssm --no-build-isolation` | Installs `selective_scan_cuda`; pip first tries a matching prebuilt CUDA/HIP wheel, then compiles locally if no wheel is available. |
+| Force local CUDA selective scan build | `MAMBA_FORCE_BUILD=TRUE MAMBA_KEEP_CUDA_BUILD=TRUE pip install mamba-ssm --no-build-isolation` | Skips cached wheels and compiles `selective_scan_cuda` locally. |
+
+`--no-build-isolation` is required for CUDA builds so that pip uses your existing CUDA-enabled
+PyTorch instead of installing torch-cpu in an isolated build environment.
+
+Source installs use the same defaults and opt-in flags:
+
+| Install mode | Command |
+| --- | --- |
+| Source default, no `selective_scan_cuda` | `pip install . --no-build-isolation` |
+| Source default from GitHub, no `selective_scan_cuda` | `pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation` |
+| Source forced local core build, no `selective_scan_cuda` | `MAMBA_FORCE_BUILD=TRUE pip install . --no-build-isolation` |
+| Source CUDA selective scan opt-in | `MAMBA_KEEP_CUDA_BUILD=TRUE pip install . --no-build-isolation` |
+| Source forced local CUDA selective scan build | `MAMBA_FORCE_BUILD=TRUE MAMBA_KEEP_CUDA_BUILD=TRUE pip install . --no-build-isolation` |
+
+NOTE: To use Mamba-3 from the latest source tree, install from source. For the default source
+install without `selective_scan_cuda`, run `pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation`.
+For the CUDA selective scan extension, add `MAMBA_KEEP_CUDA_BUILD=TRUE`; add
+`MAMBA_FORCE_BUILD=TRUE` as well to force local CUDA compilation.
+
+Core requirements:
 - Linux
-- NVIDIA GPU
+- Python 3.10+
 - PyTorch 1.12+
+
+Additional requirements for CUDA `selective_scan_cuda` builds and GPU execution:
+- NVIDIA GPU
 - CUDA 11.6+
 
 For AMD cards, see additional prerequisites below.

mamba_ssm/ops/selective_scan_interface.py

@@ -17,14 +17,26 @@
 
 from mamba_ssm.ops.triton.layer_norm import _layer_norm_fwd
 
-import selective_scan_cuda
+try:
+    import selective_scan_cuda
+except ImportError:
+    selective_scan_cuda = None
+
+
+def _check_selective_scan_cuda():
+    if selective_scan_cuda is None:
+        raise RuntimeError(
+            "selective_scan_cuda is not installed. To use Mamba-1, reinstall with:\n"
+            "  MAMBA_KEEP_CUDA_BUILD=TRUE pip install mamba-ssm --no-build-isolation"
+        )
 
 
 class SelectiveScanFn(torch.autograd.Function):
 
     @staticmethod
     def forward(ctx, u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,
                 return_last_state=False):
+        _check_selective_scan_cuda()
         if u.stride(-1) != 1:
             u = u.contiguous()
         if delta.stride(-1) != 1:
@@ -192,6 +204,7 @@ def forward(ctx, xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weigh
         """
              xz: (batch, dim, seqlen)
         """
+        _check_selective_scan_cuda()
         assert causal_conv1d_fwd_function is not None, "causal_conv1d_cuda is not available. Please install causal-conv1d."
         assert checkpoint_lvl in [0, 1]
         L = xz.shape[-1]