在 RTX 4090 + Python 3.10 上成功安装 Mamba‑3 的避坑指南

记录一次"官方命令让我重装了 8 次 PyTorch"的血泪史，最终找到一条稳定可靠的安装路径。

背景

Mamba‑3 是近期非常受关注的状态空间模型（SSM）架构，其 MIMO（多输入多输出）模式在长序列建模上展现了极强的性能。官方给出的安装命令看似简单，但实际在特定环境下（如 RTX 4090 + 高版本驱动）会引发一连串的依赖冲突，导致反复卸载、重装 PyTorch，甚至卡死在源码编译阶段。

本文详细记录了我从零搭建 Mamba‑3 环境的完整过程，以及中间遇到的各种错误和最终解决方案，希望能帮助同样踩坑的朋友们一劳永逸。

实验环境

操作系统 ：Ubuntu 22.04（Linux）
GPU ：NVIDIA GeForce RTX 4090（24GB）
驱动：570.211.01，支持 CUDA 12.8
Python ：3.10（conda 环境 mamba3）
目标：安装支持 PyTorch 2.6 + Mamba‑3（MIMO 模式）的开发环境

官方命令为什么不行？

官方 README 中推荐的安装步骤为：

复制代码

pip install torch  # 先装 PyTorch`
`pip install causal-conv1d>=1.4.0 --no-build-isolation`
`pip install mamba-ssm --no-build-isolation`
`# 或者 Mamba‑3 专用命令：`
`MAMBA_FORCE_BUILD=TRUE pip install --no-cache-dir --force-reinstall \`
`    git+https://github.com/state-spaces/mamba.git --no-build-isolation`
`

但在我的环境中，直接执行这些命令会触发以下问题：

--force-reinstall 强制重装所有依赖

即使你已经手动安装好正确版本的 PyTorch（例如 torch 2.6.0+cu124），--force-reinstall 也会让 pip 重新解析依赖 ，它会在默认 PyPI 上寻找最新的 torch（当前为 2.12.0），并自动卸载你的定制版 PyTorch ，换上不兼容的版本，导致 torchvision/torchaudio 报错。

Triton 版本冲突

torch 2.6.0+cu124 写死了依赖 triton==3.2.0，而 mamba-ssm 需要 triton>=3.5.0。若直接一起装，pip 会陷入依赖解析地狱，报错：

复制代码

ERROR: Cannot install mamba-ssm and torch==2.6.0+cu124 because these package versions have conflicting dependencies.`
`

网络不稳定导致 git clone 失败

从 GitHub 克隆 Mamba 源码时，常因国内网络波动出现 RPC failed; curl 16 Error in the HTTP2 framing layer 错误，安装过程卡死或直接中断。

运行时共享内存溢出

成功安装后，若使用不当的参数（例如在 float32 精度下给 chunk_size=16），会触发 Failed to set the allowed dynamic shared memory size 错误，看似又像是安装问题，实际是参数配置与硬件特性不符。

一步一步走向成功

下面是我最终验证可行的安装流程，只需严格按顺序执行 ，无需提前下载任何本地源码，全程在线完成（若网络实在太差，可替换为 Gitee 镜像，文中一并给出）。

1. 安装 CUDA Toolkit（提供 nvcc）

Mamba‑3 源码编译需要 nvcc 编译器。你的系统驱动已经是 12.8，只需装 Toolkit。为免污染系统，安装到用户目录：

复制代码

wget https://developer.download.nvidia.com/compute/cuda/12.8.0/local_installers/cuda_12.8.0_570.86.10_linux.run`
`sh cuda_12.8.0_570.86.10_linux.run --toolkit --silent --override --installpath=$HOME/local/cuda-12.8`
`

设置环境变量（可写入 ~/.bashrc）：

复制代码

export` `CUDA_HOME=$HOME/local/cuda-12.8`
`export` `PATH=$CUDA_HOME/bin:$PATH`
`export` `LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH`
`

执行 nvcc --version 应显示 12.8。

2. 创建 Python 环境并安装指定版本的 PyTorch

复制代码

conda create -n mamba3 python=3.10 -y`
`conda activate mamba3`

`pip install` `torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 \`
`    --index-url https://download.pytorch.org/whl/cu124`
`

此时 torch 为 2.6.0+cu124，自带 triton 3.2.0。

3. 解决 Triton 版本冲突：强制升级到 3.5+

关键步骤 ：用 --no-deps 覆盖安装，避免 pip 因为 torch 的版本限制而拒绝执行。

复制代码

pip install --upgrade --force-reinstall "triton>=3.5.0" --no-deps`
`

这一步之后，triton 升级到 3.7.0，但 torch 依然保持 2.6.0，兼容性在实践上没有任何问题。

4. 安装 Mamba‑3（跳过依赖，避免 triton 降级）

复制代码

MAMBA_FORCE_BUILD=TRUE pip install git+https://github.com/state-spaces/mamba.git \`
`    --no-build-isolation --no-deps`
`

--no-deps：防止 pip 重新解析依赖并把 triton 拉回 3.2.0。
--no-build-isolation：复用当前环境的 PyTorch 进行编译，不会新建一个隔离环境并下载 CPU 版 PyTorch。
MAMBA_FORCE_BUILD=TRUE：确保编译 Mamba‑3 所需的全部 CUDA 算子。

如果 GitHub 克隆失败 ，改用 Gitee 镜像：

复制代码

MAMBA_FORCE_BUILD=TRUE pip install git+https://gitee.com/mirrors/state-spaces-mamba.git \`
`    --no-build-isolation --no-deps`
`

5. 补全其余依赖（不会动 triton 和 torch）

复制代码

pip install einops transformers tilelang apache-tvm-ffi z3-solver`
`

注意：此处没有安装 causal-conv1d，因为 Mamba‑3 的 MIMO 模式当前可以脱离它运行，且其编译容易触发网络和 CUDA 版本问题（本文最后会提及如何处理）。

验证安装

运行一个简单的 MIMO 前向传播测试：

复制代码

python -c "`
`import torch`
`from mamba_ssm import Mamba3`

`batch, length, dim = 2, 64, 64`
`x = torch.randn(batch, length, dim, dtype=torch.float32, device='cuda')`
`model = Mamba3(`
`    d_model=dim,`
`    d_state=128,`
`    headdim=64,`
`    is_mimo=True,`
`    mimo_rank=4,`
`    chunk_size=8,           # float32 时务必用 32/mimo_rank = 8`
`    dtype=torch.float32`
`).to('cuda')`
`y = model(x)`
`print('y.shape:', y.shape)`
`print('Mamba-3 前向传播成功！')`
`"`
`

首次运行会看到 TileLang 编译内核的日志：

复制代码

TileLang begins to compile kernel `mamba_mimo_fwd_kernel` ...`
`TileLang completes to compile kernel `mamba_mimo_fwd_kernel``
`y.shape: torch.Size([2, 64, 64])`
`Mamba-3 前向传播成功！`
`

表示环境 100% 可用。

常见错误及解决方案

1. 共享内存溢出

复制代码

tvm.error.InternalError: Failed to set the allowed dynamic shared memory size to 168128`
`

原因：chunk_size 设置过大，超出了 RTX 4090 默认动态共享内存限制。

解决：

float32 精度：chunk_size = 32 / mimo_rank（例如 mimo_rank=4 时用 8）
bfloat16 精度：chunk_size = 64 / mimo_rank（可用 16）

2. causal-conv1d 编译失败

若后续需要该包（例如使用 SISO 模式或某些训练配置），安装时务必加上 --no-build-isolation，否则会触发 PyTorch 版本和 CUDA 版本不匹配的错误。

复制代码

# 正确的安装方式（使用当前环境的 PyTorch 编译）`
`pip install causal-conv1d>=1.4.0 --no-build-isolation`
`

若网络不佳，可尝试预编译 wheel：

复制代码

pip install causal-conv1d --only-binary causal-conv1d`
`

3. 每次更改模型参数都需要重新编译内核

TileLang 会缓存编译结果于 ~/.cache/tilelang，若修改了 chunk_size、headdim 等参数，记得清除缓存：

复制代码

rm -rf ~/.cache/tilelang`
`

总结：为什么这套方案可靠？

锁定 PyTorch 版本 ：选择了与 RTX 4090 驱动完美匹配的 2.6.0+cu124，避免使用官方的 latest 带来的不确定性。
手动调和 Triton 冲突 ：通过 --no-deps 覆盖安装，绕开 PyTorch 的严格版本限制，实际运行完全兼容。
跳过 Mamba 的依赖传递 ：用 --no-deps 安装 Mamba‑3，再手动补上缺失的包，让 pip 不再"自作聪明"。
参数遵循官方硬件建议 ：chunk_size 的计算公式直接避免运行时共享内存溢出。

整个过程虽然走了不少弯路，但最终整理出的步骤只需 5 条命令，10 分钟内即可完成。希望这篇记录能让你在安装 Mamba‑3 时不再被依赖折磨，一次成功！