Mamba架构教程：选择性状态空间模型安装与使用指南，实现O(n)序列建模

mamba-architecture by davila7/claude-code-templates

162 周安装量

23,400 GitHub Stars

GitHub

安装命令

npx skills add https://github.com/davila7/claude-code-templates --skill mamba-architecture

AI/机器学习 PyTorch 自然语言处理

🇨🇳中文介绍

Mamba - 选择性状态空间模型

快速开始

Mamba 是一种状态空间模型架构，在序列建模中实现了 O(n) 线性复杂度。

安装：

# 安装 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,      # 模型维度
    d_state=16,       # SSM 状态维度
    d_conv=4,         # Conv1d 卷积核大小
    expand=2          # 扩展因子
).to("cuda")

y = model(x)  # O(n) 复杂度！
assert y.shape == x.shape

常用工作流

工作流 1：使用 Mamba-2 构建语言模型

包含生成的完整语言模型：

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from mamba_ssm.models.config_mamba import MambaConfig
import torch

# 配置 Mamba-2 语言模型
config = MambaConfig(
    d_model=1024,           # 隐藏层维度
    n_layer=24,             # 层数
    vocab_size=50277,       # 词汇表大小
    ssm_cfg=dict(
        layer="Mamba2",     # 使用 Mamba-2
        d_state=128,        # Mamba-2 使用更大的状态维度
        headdim=64,         # 头维度
        ngroups=1           # 分组数量
    )
)

model = MambaLMHeadModel(config, device="cuda", dtype=torch.float16)

# 生成文本
input_ids = torch.randint(0, 1000, (1, 20), device="cuda", dtype=torch.long)
output = model.generate(
    input_ids=input_ids,
    max_length=100,
    temperature=0.7,
    top_p=0.9
)

广告位招租

在这里展示您的产品或服务

触达数万 AI 开发者，精准高效

联系我们

工作流 2：使用预训练的 Mamba 模型

从 HuggingFace 加载：

from transformers import AutoTokenizer
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

# 加载预训练模型
model_name = "state-spaces/mamba-2.8b"
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")  # 使用兼容的分词器
model = MambaLMHeadModel.from_pretrained(model_name, device="cuda", dtype=torch.float16)

# 生成
prompt = "The future of AI is"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
output_ids = model.generate(
    input_ids=input_ids,
    max_length=200,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.2
)
generated_text = tokenizer.decode(output_ids[0])
print(generated_text)

state-spaces/mamba-130m
state-spaces/mamba-370m
state-spaces/mamba-790m
state-spaces/mamba-1.4b
state-spaces/mamba-2.8b

工作流 3：Mamba-1 与 Mamba-2 对比

Mamba-1（状态维度较小）：

from mamba_ssm import Mamba

model = Mamba(
    d_model=256,
    d_state=16,      # 较小的状态维度
    d_conv=4,
    expand=2
).to("cuda")

Mamba-2（多头，状态维度较大）：

from mamba_ssm import Mamba2

model = Mamba2(
    d_model=256,
    d_state=128,     # 较大的状态维度
    d_conv=4,
    expand=2,
    headdim=64,      # 多头结构的头维度
    ngroups=1        # 并行组数
).to("cuda")

状态大小：Mamba-1 (d_state=16) 对比 Mamba-2 (d_state=128)
架构：Mamba-2 具有多头结构
归一化：Mamba-2 使用 RMSNorm
分布式：Mamba-2 支持张量并行

工作流 4：与 Transformers 进行基准测试对比

生成速度对比：

# 对 Mamba 进行基准测试
python benchmarks/benchmark_generation_mamba_simple.py \
  --model-name "state-spaces/mamba-2.8b" \
  --prompt "The future of machine learning is" \
  --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

# 对 Transformer 进行基准测试
python benchmarks/benchmark_generation_mamba_simple.py \
  --model-name "EleutherAI/pythia-2.8b" \
  --prompt "The future of machine learning is" \
  --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

Mamba：推理速度快 5 倍
内存：无需 KV 缓存
扩展性：与序列长度呈线性关系

何时使用及替代方案

在以下情况使用 Mamba：

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

O(n) 复杂度：线性对比二次方
推理速度快 5 倍：无注意力机制开销
无需 KV 缓存：内存使用率更低
百万 token 序列：硬件效率高
流式处理：每个 token 占用恒定内存

在以下情况使用替代方案：

Transformers：需要最佳性能，且计算资源充足
RWKV：想要 RNN+Transformer 混合架构
RetNet：需要基于保留机制的架构
Hyena：想要基于卷积的方法

问题：CUDA 内存不足

减少批处理大小或使用梯度检查点：

model = MambaLMHeadModel(config, device="cuda", dtype=torch.float16)
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）：

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b")

选择性 SSM：请参阅 references/selective-ssm.md 了解数学公式、状态空间方程以及选择性如何实现 O(n) 复杂度。

Mamba-2 架构：请参阅 references/mamba2-details.md 了解多头结构、张量并行和分布式训练设置。

性能优化：请参阅 references/performance.md 了解硬件感知设计、CUDA 内核和内存效率技术。

GPU：支持 CUDA 11.6+ 的 NVIDIA GPU
显存：
- 130M 模型：2GB
- 370M 模型：4GB
- 790M 模型：8GB
- 1.4B 模型：14GB
- 2.8B 模型：28GB (FP16)
推理：比 Transformers 快 5 倍
内存：无需 KV 缓存（低于 Transformers）

性能（与 Transformers 对比）：

速度：推理速度快 5 倍
内存：减少 50%（无需 KV 缓存）
扩展性：线性对比二次方

论文 (Mamba-1): https://arxiv.org/abs/2312.00752 (2023年12月)
论文 (Mamba-2): https://arxiv.org/abs/2405.21060 (2024年5月)
GitHub: https://github.com/state-spaces/mamba ⭐ 13,000+
模型: https://huggingface.co/state-spaces
文档: 仓库 README 和 wiki

🇺🇸English

Mamba - Selective State Space Models

Quick start

Mamba is a state-space model architecture achieving O(n) linear complexity for sequence modeling.

Installation :

# Install causal-conv1d (optional, for efficiency)
pip install causal-conv1d>=1.4.0

# Install Mamba
pip install mamba-ssm
# Or both together
pip install mamba-ssm[causal-conv1d]

Prerequisites : Linux, NVIDIA GPU, PyTorch 1.12+, CUDA 11.6+

Basic usage (Mamba block):

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

Common workflows

Workflow 1: Language model with Mamba-2

Complete LM with generation :

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from mamba_ssm.models.config_mamba import MambaConfig
import torch

# Configure Mamba-2 LM
config = MambaConfig(
    d_model=1024,           # Hidden dimension
    n_layer=24,             # Number of layers
    vocab_size=50277,       # Vocabulary size
    ssm_cfg=dict(
        layer="Mamba2",     # Use Mamba-2
        d_state=128,        # Larger state for Mamba-2
        headdim=64,         # Head dimension
        ngroups=1           # Number of groups
    )
)

model = MambaLMHeadModel(config, device="cuda", dtype=torch.float16)

# Generate text
input_ids = torch.randint(0, 1000, (1, 20), device="cuda", dtype=torch.long)
output = model.generate(
    input_ids=input_ids,
    max_length=100,
    temperature=0.7,
    top_p=0.9
)

Workflow 2: Use pretrained Mamba models

Load from HuggingFace :

from transformers import AutoTokenizer
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

# Load pretrained model
model_name = "state-spaces/mamba-2.8b"
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")  # Use compatible tokenizer
model = MambaLMHeadModel.from_pretrained(model_name, device="cuda", dtype=torch.float16)

# Generate
prompt = "The future of AI is"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
output_ids = model.generate(
    input_ids=input_ids,
    max_length=200,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.2
)
generated_text = tokenizer.decode(output_ids[0])
print(generated_text)

Available models :

state-spaces/mamba-130m
state-spaces/mamba-370m
state-spaces/mamba-790m
state-spaces/mamba-1.4b
state-spaces/mamba-2.8b

Workflow 3: Mamba-1 vs Mamba-2

Mamba-1 (smaller state):

from mamba_ssm import Mamba

model = Mamba(
    d_model=256,
    d_state=16,      # Smaller state dimension
    d_conv=4,
    expand=2
).to("cuda")

Mamba-2 (multi-head, larger state):

from mamba_ssm import Mamba2

model = Mamba2(
    d_model=256,
    d_state=128,     # Larger state dimension
    d_conv=4,
    expand=2,
    headdim=64,      # Head dimension for multi-head
    ngroups=1        # Parallel groups
).to("cuda")

Key differences :

State size : Mamba-1 (d_state=16) vs Mamba-2 (d_state=128)
Architecture : Mamba-2 has multi-head structure
Normalization : Mamba-2 uses RMSNorm
Distributed : Mamba-2 supports tensor parallelism

Workflow 4: Benchmark vs Transformers

Generation speed comparison :

# Benchmark Mamba
python benchmarks/benchmark_generation_mamba_simple.py \
  --model-name "state-spaces/mamba-2.8b" \
  --prompt "The future of machine learning is" \
  --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

# Benchmark Transformer
python benchmarks/benchmark_generation_mamba_simple.py \
  --model-name "EleutherAI/pythia-2.8b" \
  --prompt "The future of machine learning is" \
  --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

Expected results :

Mamba : 5× faster inference
Memory : No KV cache needed
Scaling : Linear with sequence length

When to use vs alternatives

Use Mamba when :

Need long sequences (100K+ tokens)
Want faster inference than Transformers
Memory-constrained (no KV cache)
Building streaming applications
Linear scaling important

Advantages :

O(n) complexity : Linear vs quadratic
5× faster inference : No attention overhead
No KV cache : Lower memory usage
Million-token sequences : Hardware-efficient
Streaming : Constant memory per token

Use alternatives instead :

Transformers : Need best-in-class performance, have compute
RWKV : Want RNN+Transformer hybrid
RetNet : Need retention-based architecture
Hyena : Want convolution-based approach

Common issues

Issue: CUDA out of memory

Reduce batch size or use gradient checkpointing:

model = MambaLMHeadModel(config, device="cuda", dtype=torch.float16)
model.gradient_checkpointing_enable()  # Enable checkpointing

Issue: Slow installation

Install binary wheels (not source):

pip install mamba-ssm --no-build-isolation

Issue: Missing causal-conv1d

Install separately:

pip install causal-conv1d>=1.4.0

Issue: Model not loading from HuggingFace

Use MambaLMHeadModel.from_pretrained (not AutoModel):

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b")

Advanced topics

Selective SSM : See references/selective-ssm.md for mathematical formulation, state-space equations, and how selectivity enables O(n) complexity.

Mamba-2 architecture : See references/mamba2-details.md for multi-head structure, tensor parallelism, and distributed training setup.

Performance optimization : See references/performance.md for hardware-aware design, CUDA kernels, and memory efficiency techniques.

Hardware requirements

GPU : NVIDIA with CUDA 11.6+
VRAM :
- 130M model: 2GB
- 370M model: 4GB
- 790M model: 8GB
- 1.4B model: 14GB
- 2.8B model: 28GB (FP16)
Inference : 5× faster than Transformers
Memory : No KV cache (lower than Transformers)

Performance (vs Transformers):

Speed : 5× faster inference
Memory : 50% less (no KV cache)
Scaling : Linear vs quadratic

Resources

Paper (Mamba-1): https://arxiv.org/abs/2312.00752 (Dec 2023)
Paper (Mamba-2): https://arxiv.org/abs/2405.21060 (May 2024)
GitHub: https://github.com/state-spaces/mamba ⭐ 13,000+
Models: https://huggingface.co/state-spaces
Docs: Repository README and wiki

Weekly Installs

142

Repository

davila7/claude-…emplates

GitHub Stars

22.6K

First Seen

Jan 21, 2026

Security Audits

Gen Agent Trust HubPass SocketPass SnykWarn

Installed on

claude-code117

opencode110

gemini-cli106

cursor103

antigravity95

codex92

超能力技能使用指南：AI助手技能调用优先级与工作流程详解

46,500 周安装