flash-mamba

flash-mamba 是官方 mamba-ssm 的轻量包装库。 统一创建 Mamba / Mamba2，提供懒加载和环境诊断，返回官方模型实例。 训练、推理缓存、参数默认值、state_dict 和内核行为均由官方实现决定。

这是独立项目，与官方团队没有隶属关系。项目名称不代表新增加速内核或性能提升。 v0.1.0 的封装目标是 PyPI mamba-ssm==2.3.2.post1；Mamba3 不在本次范围内。 当前本机只验证了包装层和发行包，真实 GPU 验证需在 Linux 环境完成。

安装

包装层需要 Python 3.10+，自身没有运行时依赖，可在 Windows 导入和运行诊断。 发行包发布前，从本地安装：

python -m pip install .

发布到 PyPI 后才可使用：

python -m pip install flash-mamba

实际模型运行需要先在 Linux（Windows 用户可使用 WSL2） 安装适配设备的 PyTorch，再安装官方 GPU 后端。包装层不提供 CPU 替代实现或原生 Windows 内核。 CUDA/ROCm、编译工具和模型的具体要求以 官方安装文档 为准。

# 先按 https://pytorch.org/get-started/locally/ 安装对应的 GPU 版 PyTorch
python -m pip install packaging ninja wheel setuptools
python -m pip install "mamba-ssm[causal-conv1d]==2.3.2.post1" --no-build-isolation
python -m pip install .
flash-mamba doctor

ssm extra 声明 mamba-ssm>=2.3.2.post1,<3。安装前仍需先配置 GPU 版 PyTorch 和官方构建依赖：

python -m pip install ".[ssm]" --no-build-isolation

显式两步安装更方便排查官方编译或二进制兼容问题。包装层不会自动安装、下载或 替换 PyTorch。官方 main 与 PyPI 发行版的安装默认值可能不同；当前 main 默认不构建 selective_scan_cuda，Mamba-1 使用该扩展时需要显式开启：

# 仅当选择官方最新源码时使用；并非 v0.1.0 已验证的发行版组合
MAMBA_KEEP_CUDA_BUILD=TRUE python -m pip install \
  "mamba-ssm[causal-conv1d] @ git+https://github.com/state-spaces/mamba.git" \
  --no-build-isolation

使用

import torch
from flash_mamba import create_mamba

x = torch.randn(2, 64, 128, device="cuda")
model = create_mamba(
    d_model=128,
    variant="mamba2",  # 默认；Mamba-1 使用 "mamba"
    d_state=64,
    d_conv=4,
    expand=2,
    headdim=64,
    device="cuda",
)
y = model(x)
assert y.shape == x.shape
y.square().mean().backward()

create_mamba 只校验正整数 d_model 和 variant；其余关键字直接传给官方构造函数。 不额外添加残差、归一化、转置或模型包装。Mamba2 的 d_ssm（默认为 expand * d_model，分布式时还受 world_size 影响）必须满足官方 headdim 约束；例如 d_model=16, expand=2, headdim=64 不合法。

也可以直接导入官方类，第一次访问时才加载 mamba_ssm：

from flash_mamba import Mamba, Mamba2

model = Mamba(d_model=128, device="cuda")

返回的实例可直接调用 .to()、.parameters()、.state_dict()、 .allocate_inference_cache() 和官方支持的 forward 关键字。 import flash_mamba 本身不导入 PyTorch 或 mamba_ssm。 后端导入失败会抛出 BackendUnavailableError，保留原异常链；构造函数参数错误 和模型执行错误直接沿用官方异常。

环境诊断

flash-mamba doctor
flash-mamba doctor --json
python -m flash_mamba doctor --json

检查已安装版本、PyTorch/官方后端能否导入、CUDA/ROCm 是否可访问。 退出码 0 表示可以尝试 GPU smoke test；退出码 1 表示检查未通过。 成功导入不等于内核可运行，真实前向/反向需执行下面的测试。 flash_mamba.diagnose() 返回相同的字典结果，不修改环境。

开发与验证

python -m pip install -e ".[dev]"
python -m ruff check .
python -m ruff format --check .
python -m pytest
python -m build
python -m twine check --strict dist/*
python scripts/verify_dist.py

仓库约定：没有其他环境要求时，以上 Python 命令在 conda torch2 中运行。 默认测试不依赖官方后端，检查懒加载、参数传递、异常处理、诊断和 CLI。 发行包验证在临时虚拟环境中安装 wheel 以及从 sdist 构建的 wheel。

在已配置官方库的 Linux GPU 环境中，必须执行真实测试再确认模型可运行：

FLASH_MAMBA_RUN_GPU_TESTS=1 python -m pytest tests/test_gpu.py -v

显式开启后，环境缺失或模型运行失败会导致测试失败，不会静默跳过。 测试比较工厂实例与官方实例的前向结果、参数/输入梯度，并检查 checkpoint 可直接互载。未开启时，GPU 测试明确标记为 skipped。

发布

发布准备、TestPyPI/PyPI 命令和 GitHub Trusted Publishing 配置见 发布说明。CI 检查包装层并构建发行包，不能替代真实 GPU 测试。

许可与来源

本项目的新包装代码使用 MIT 许可证；不复制、不重新分发官方源码或 CUDA 内核。 mamba-ssm 由 Albert Gu、Tri Dao 等作者开发，使用 Apache-2.0 许可证。 使用该模型进行研究时，请按官方说明 引用相关论文。