用途 Mamba 是面向序列建模的状态空间模型架构，可作为 Transformer 的替代方案，用于语言模型、长序列处理和流式生成等任务。资料称其复杂度为 O(n)，随序列长度线性增长，而 Transformer 为 O(n²)。
主要能力
- 线性复杂度：O(n) 序列建模，支持百万级 token 序列。
- 推理效率：资料称推理速度约为 Transformer 的 5 倍，且不需要 KV cache，内存占用更低。
- 选择性 SSM：通过选择性机制实现线性复杂度（资料指引见 references/selective-ssm.md）。
- 两个版本：Mamba-1（d_state=16）与 Mamba-2（d_state=128、多头结构、RMSNorm、支持张量并行）。
- 预训练权重：提供 mamba-130m、370m、790m、1.4b、2.8b 等规模，可通过 HuggingFace 加载。
使用方式 安装（Linux + NVIDIA GPU + PyTorch 1.12+ + CUDA 11.6+）： pip install causal-conv1d>=1.4.0 pip install mamba-ssm 或合并安装：pip install mamba-ssm[causal-conv1d]
基础用法为实例化 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)
语言模型可使用 MambaLMHeadModel 配合 MambaConfig 构建并调用 generate 生成文本；也可用 from_pretrained 加载 HuggingFace 上的 state-spaces 预训练模型（需使用兼容 tokenizer，例如 EleutherAI/gpt-neox-20b）。资料还提供与 Transformer（EleutherAI/pythia-2.8b）对比的基准脚本 benchmarks/benchmark_generation_mamba_simple.py。
适用场景 需要 100K+ 长序列、希望比 Transformer 更快推理、内存受限（无 KV cache）、构建流式应用、以及重视线性扩展性的场景。资料建议：若追求最佳性能且有充足算力，可继续使用 Transformer；也可考虑 RWKV、RetNet、Hyena 等替代架构。
限制与注意事项
- 环境要求：Linux、NVIDIA GPU、CUDA 11.6+，非 NVIDIA 环境未在资料中说明支持。
- 显存参考：130M 约 2GB、370M 约 4GB、790M 约 8GB、1.4B 约 14GB、2.8B 约 28GB（FP16）。
- 显存不足时可减小 batch 或启用 gradient_checkpointing_enable()。
- 安装慢可用 pip install mamba-ssm --no-build-isolation 安装二进制包。
- 加载 HuggingFace 模型应使用 MambaLMHeadModel.from_pretrained，而非 AutoModel。
- 资料中的速度与内存对比数据（5 倍推理、内存减少 50% 等）为预期/宣称值，实际表现依环境而定。
参考：论文 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。