mamba-architecture

Mamba 是一种状态空间模型（State Space Model）架构，用于序列建模，可实现 O(n) 线性复杂度。据资料介绍，相比 Transformer 的 O(n²) 复杂度，Mamba 推理速度约为其 5 倍，支持百万级 token 序列，且无需 KV 缓存。

主要能力

• 选择性 SSM（Selective SSM）：通过选择性机制实现 O(n) 复杂度。
• 线性复杂度：随序列长度线性扩展，而非二次方。
• 推理更快：资料称推理速度约为 Transformer 的 5 倍。
• 无 KV 缓存：内存占用更低，资料称相比 Transformer 减少约 50%。
• 流式处理：每个 token 内存占用恒定。
• 提供 Mamba-1 与 Mamba-2 两种实现，Mamba-2 具备多头结构、RMSNorm、支持张量并行。

安装方式 安装 causal-conv1d（可选，用于提升效率）： pip install causal-conv1d>=1.4.0

安装 Mamba： pip install mamba-ssm 或两者一起安装： pip install mamba-ssm[causal-conv1d]

前置条件：Linux、NVIDIA GPU、PyTorch 1.12+、CUDA 11.6+。

基本用法（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, # Model dimension d_state=16, # SSM state dimension d_conv=4, # Conv1d kernel size expand=2 # Expansion factor ).to("cuda")

y = model(x) # O(n) complexity! assert y.shape == x.shape

常见工作流

1. 使用 Mamba-2 构建语言模型：通过 MambaConfig 配置 d_model、n_layer、vocab_size 及 ssm_cfg（layer="Mamba2"、d_state=128、headdim=64、ngroups=1），再用 MambaLMHeadModel 加载并调用 generate 生成文本。
2. 使用预训练 Mamba 模型：从 HuggingFace 加载，例如 state-spaces/mamba-2.8b，配合兼容分词器（如 EleutherAI/gpt-neox-20b），使用 MambaLMHeadModel.from_pretrained 加载后生成。可用模型包括 state-spaces/mamba-130m、mamba-370m、mamba-790m、mamba-1.4b、mamba-2.8b。
3. Mamba-1 与 Mamba-2 对比：Mamba-1 状态维度较小（d_state=16）；Mamba-2 状态维度更大（d_state=128），具备多头结构、使用 RMSNorm、支持张量并行。
4. 与 Transformer 基准对比：通过 benchmarks/benchmark_generation_mamba_simple.py 脚本，分别对 Mamba 与 Transformer 模型进行生成速度测试。资料称预期结果为 Mamba 推理快 5 倍、无需 KV 缓存、随序列长度线性扩展。

适用场景

• 需要长序列（100K+ token）。
• 希望比 Transformer 更快的推理速度。
• 内存受限（无 KV 缓存）。
• 构建流式应用。
• 线性扩展很重要。

限制与替代

• 若需要最佳性能且有充足算力，可考虑 Transformer。
• 若想要 RNN+Transformer 混合，可考虑 RWKV。
• 若需要基于保留（retention）的架构，可考虑 RetNet。
• 若想要基于卷积的方法，可考虑 Hyena。

硬件要求

• GPU：NVIDIA，CUDA 11.6+。
• 显存：130M 模型约 2GB；370M 约 4GB；790M 约 8GB；1.4B 约 14GB；2.8B 约 28GB（FP16）。

常见问题

• CUDA 显存不足：减小 batch size 或启用梯度检查点（model.gradient_checkpointing_enable()）。
• 安装缓慢：安装二进制 wheel（pip install mamba-ssm --no-build-isolation）。
• 缺少 causal-conv1d：单独安装 pip install causal-conv1d>=1.4.0。
• 模型无法从 HuggingFace 加载：使用 MambaLMHeadModel.from_pretrained（而非 AutoModel）。

进阶主题

• 选择性 SSM：见 references/selective-ssm.md。
• Mamba-2 架构：见 references/mamba2-details.md。
• 性能优化：见 references/performance.md。

资源

• 论文（Mamba-1）：https://arxiv.org/abs/2312.00752
• 论文（Mamba-2）：https://arxiv.org/abs/2405.21060
• GitHub：https://github.com/state-spaces/mamba
• 模型：https://huggingface.co/state-spaces