Mamba

GitHub:https://github.com/state-spaces/mamba

描述信息:基于结构化状态空间模型SSM的新型深度学习架构

GitHub

Mamba（包含 Mamba-2）的官方实现代码库，聚焦于基于选择性状态空间模型（Selective State Space Model, SSM）的线性时间序列建模，由 Albert Gu 和 Tri Dao 主导开发，核心目标是解决传统 Transformer 在信息密集型数据（如语言建模）上的效率瓶颈，同时保持甚至超越其性能。

一、核心定位

Mamba 是一种全新的状态空间模型架构，基于结构化状态空间模型（SSM）的研究进展，结合硬件感知的高效设计（借鉴 FlashAttention 思路），实现了线性时间复杂度的序列建模，相比 Transformer 类模型在长序列任务上更高效。

二、核心内容

1. 核心模块

Selective SSM：Mamba 的核心层（对应论文 Algorithm 2），实现选择性状态空间计算，源码位于 mamba_ssm/ops/selective_scan_interface.py；
Mamba Block：封装选择性 SSM 的核心架构块，基础版本在 mamba_ssm/modules/mamba_simple.py；
Mamba-2：Mamba 的升级版本（基于 State Space Dual 模型），实现位于 mamba_ssm/modules/mamba2.py/mamba2_simple.py，还提供了极简版 SSD 模块（ssd_minimal.py）；
Mamba 语言模型：完整的语言模型示例（Mamba 块堆叠 + 语言模型头），支持大语言模型训练 / 推理。

2. 预训练模型

提供多个规模的预训练模型（上传至 Hugging Face state-spaces 仓库），涵盖 Mamba 和 Mamba-2 系列，参数规模包括 130M、370M、790M、1.4B、2.8B 等，训练数据基于 Pile/SlimPajama 数据集（300B/600B tokens）。

3. 硬件与环境支持

基础依赖：Linux 系统、NVIDIA GPU（CUDA 11.6+）、PyTorch 1.12+；

扩展支持：AMD GPU（ROCM 6.0+，需补丁）、AWS Trainium 2 硬件；

编译与构建：支持多版本 CUDA/PyTorch/Python 编译，提供自动化构建脚本（GitHub Actions），可生成适配不同环境的 wheel 包。

三、使用方式

1. 安装

支持 pip 快速安装或源码编译：

# 核心包
pip install mamba-ssm
# 含因果卷积依赖
pip install mamba-ssm[causal-conv1d]
# 源码构建
pip install .

2. 基础使用示例

Mamba 基础块：

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

Mamba-2 块（用法兼容）：

from mamba_ssm import Mamba2
model = Mamba2(d_model=dim, d_state=64, d_conv=4, expand=2).to("cuda")
y = model(x)

四、工程化设计

编译构建：通过 setup.py 封装 CUDA/HIP 扩展编译，支持不同 CUDA 版本（11.6+）、C++11 ABI 适配，自动生成适配多环境的 wheel 包；

CI/CD：通过 GitHub Actions 实现自动化构建、测试、发布，支持多平台（x86/ARM）、多依赖版本组合，构建产物上传至 GitHub Release 并同步发布至 PyPI；

代码结构：模块化设计，核心计算（selective scan）通过 CUDA 实现高性能算子，上层封装模型块和语言模型，配套基准测试（benchmarks/）、单元测试（tests/）。

五、许可证与引用

许可证：基于 Apache License 2.0 开源，允许商用但需遵循许可证条款；

学术引用：若使用该代码库，需引用 Mamba 和 Mamba-2 的论文：


@article{mamba,
  title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
  author={Gu, Albert and Dao, Tri},
  journal={arXiv preprint arXiv:2312.00752},
  year={2023}
}

@inproceedings{mamba2,
  title={Transformers are {SSM}s: Generalized Models and Efficient Algorithms Through Structured State Space Duality},
  author={Dao, Tri and Gu, Albert},
  booktitle={International Conference on Machine Learning (ICML)},
  year={2024}
}