Windows 系统下安装 Mamba 全流程指南

# import selective_scan_cuda

修改函数实现，将 C++ 调用替换为 Python 引用实现：

def selective_scan_fn(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False, return_last_state=False):
    # ... docstring ...
    return selective_scan_ref(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)

def mamba_inner_fn(
    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
    out_proj_weight, out_proj_bias, A, B=None, C=None, D=None,
    delta_bias=None, B_proj_bias=None, C_proj_bias=None, delta_softplus=True
):
    return mamba_inner_ref(xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight, out_proj_weight, out_proj_bias, A, B, C, D, delta_bias, B_proj_bias, C_proj_bias, delta_softplus)

最后执行安装：

pip install .

7. 常见问题处理

版本不匹配报错

若运行时出现 TypeError: causal_conv1d_fwd(): incompatible function arguments，说明 mamba 与 causal-conv1d 版本不对应。请务必先调整版本再安装。

输入参数缺失

若遇到类似 seq_idx 相关的报错，可能是源码逻辑问题。可在传入激活函数前添加 seq_idx=None 参数进行兼容。

8. 测试验证

运行以下代码验证安装是否成功：

import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    d_model=dim,
    d_state=16,
    d_conv=4,
    expand=2
).to("cuda")
y = model(x)
assert y.shape == x.shape
print('success')

若输出 success，则表明安装完成。