Mamba 与状态空间模型 (SSM)：超越 Transformer 的下一代序列建模架构

序列长度     注意力计算量    相对代价
1K tokens    1M 次运算      1x（基准）
4K tokens    16M 次运算     16x
32K tokens   1,024M 次运算  1,024x
128K tokens  16,384M 次运算 16,384x
1M tokens    1,000,000M     1,000,000x

状态更新方程：  h'(t) = A · h(t) + B · x(t)
输出方程：      y(t) = C · h(t) + D · x(t)

其中：
- x(t)：输入信号（当前时刻的输入）
- h(t)：隐藏状态（对历史信息的压缩表示）
- y(t)：输出信号
- A：状态转移矩阵（控制状态如何演化）
- B：输入矩阵（控制输入如何影响状态）
- C：输出矩阵（控制状态如何映射到输出）
- D：直通矩阵（输入到输出的直接连接，通常省略）

离散化后：
  h[k] = Ā · h[k-1] + B̄ · x[k]
  y[k] = C · h[k]

其中：
  Ā = exp(A · Δ)    （离散化后的状态矩阵）
  B̄ = (Ā - I) · A⁻¹ · B · Δ  （近似简化后的输入矩阵）

递推模式（推理时使用）：
  逐步计算 h[k] = Ā · h[k-1] + B̄ · x[k]
  复杂度：O(n)，每步只需常数时间
  ✅ 适合自回归生成

卷积模式（训练时使用）：
  将 SSM 展开为全局卷积核 K = (C·B̄, C·Ā·B̄, C·Ā²·B̄, ...)
  用 FFT 加速卷积计算
  ✅ 适合 GPU 并行训练

# 传统 SSM（参数固定）
class TraditionalSSM:
    def __init__(self, d_model, d_state):
        self.A = nn.Parameter(...)  # 固定
        self.B = nn.Parameter(...)  # 固定
        self.C = nn.Parameter(...)  # 固定
        self.delta = nn.Parameter(...)  # 固定

    def forward(self, x):
        # 所有输入使用相同的 A, B, C, delta
        ...

# Mamba 的选择性 SSM（参数随输入变化）
class SelectiveSSM:
    def __init__(self, d_model, d_state):
        self.A = nn.Parameter(...)  # A 仍然固定（结构化）
        self.s_B = nn.Linear(d_model, d_state)   # B 由输入动态生成
        self.s_C = nn.Linear(d_model, d_state)   # C 由输入动态生成
        self.s_delta = nn.Linear(d_model, 1)     # Δ 由输入动态生成

    def forward(self, x):
        B = self.s_B(x)          # B(x)：输入决定如何写入状态
        C = self.s_C(x)          # C(x)：输入决定如何读取状态
        delta = softplus(self.s_delta(x))  # Δ(x)：输入决定时间步长
        # 每个位置的参数都不同！
        ...

class MambaBlock(nn.Module):
    """
    Mamba 的基本构建单元
    """
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        d_inner = d_model * expand

        # 输入投影：扩展维度
        self.in_proj = nn.Linear(d_model, d_inner * 2, bias=False)

        # 深度可分离卷积：局部上下文
        self.conv1d = nn.Conv1d(
            d_inner, d_inner,
            kernel_size=d_conv,
            padding=d_conv - 1,
            groups=d_inner
        )

        # SSM 参数投影
        self.x_proj = nn.Linear(d_inner, d_state * 2 + 1, bias=False)
        self.A_log = nn.Parameter(torch.log(
            torch.arange(1, d_state + 1).float()
        ))
        self.D = nn.Parameter(torch.ones(d_inner))

        # 输出投影
        self.out_proj = nn.Linear(d_inner, d_model, bias=False)

    def forward(self, x):
        # 1. 输入投影 + 分支
        xz = self.in_proj(x)
        x_branch, z = xz.chunk(2, dim=-1)

        # 2. 卷积（捕捉局部模式）
        x_branch = self.conv1d(x_branch.transpose(1, 2))
        x_branch = x_branch.transpose(1, 2)[:, :x.size(1), :]
        x_branch = F.silu(x_branch)

        # 3. 选择性 SSM
        y = self.selective_ssm(x_branch)

        # 4. 门控 + 输出
        y = y * F.silu(z)
        return self.out_proj(y)

SSM 视角（递推）：
  h[k] = a · h[k-1] + B[k] · x[k]
  y[k] = C[k] · h[k]

等价的注意力视角：
  y = M ⊙ (Q · K^T) · V

其中 M 是由 a 的幂次构成的半可分矩阵（下三角掩码 + 指数衰减）
Q 和 K 分别对应 C 和 B 的投影

Mamba-1 架构：                    Mamba-2 架构：
┌─────────────┐                  ┌─────────────┐
│  Input Proj  │                  │  Input Proj  │
│ (expand x2)  │                  │ (expand x2)  │
├─────────────┤                  ├─────────────┤
│   Conv1D    │                  │   Conv1D    │
├─────────────┤                  ├─────────────┤
│ Selective SSM│                  │   SSD Layer  │  ← 标量 A + 多头
│ (对角 A)     │                  │ (标量 A·I)   │    矩阵乘法加速
├─────────────┤                  ├─────────────┤
│   Gate + Proj│                  │  Norm + Proj │  ← 组归一化
└─────────────┘                  └─────────────┘

训练速度：基准                    训练速度：2-8x 提升

# 安装 Mamba 核心库
pip install mamba-ssm>=2.2.0

# 安装依赖（需要 CUDA 支持）
pip install causal-conv1d>=1.4.0
pip install torch>=2.1.0

# 或者通过 transformers 使用预训练 Mamba 模型
pip install transformers>=4.39.0

from transformers import MambaForCausalLM, AutoTokenizer
import torch

# 加载预训练 Mamba 模型
model_name = "state-spaces/mamba-2.8b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = MambaForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 文本生成
prompt = "The key advantage of state space models over transformers is"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
import torch

# 加载原生 Mamba 模型（推理更快）
model = MambaLMHeadModel.from_pretrained(
    "state-spaces/mamba-2.8b",
    dtype=torch.float16,
    device="cuda"
)

# 自回归生成（展示 SSM 的递推推理）
input_ids = tokenizer.encode("Explain state space models:", return_tensors="pt")
input_ids = input_ids.to("cuda")

output_ids = model.generate(
    input_ids=input_ids,
    max_length=300,
    temperature=0.7,
    top_k=50,
    cg=True  # 启用 CUDA Graph 优化
)

print(tokenizer.decode(output_ids[0]))

import torch
import time

def benchmark_inference(model, tokenizer, seq_lengths, device="cuda"):
    """对比不同序列长度下的推理效率"""
    results = []

    for seq_len in seq_lengths:
        input_ids = torch.randint(0, 32000, (1, seq_len)).to(device)

        # 预热
        with torch.no_grad():
            _ = model(input_ids)

        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()

        # 正式计时
        start = time.perf_counter()
        with torch.no_grad():
            for _ in range(10):
                _ = model(input_ids)
        torch.cuda.synchronize()
        elapsed = (time.perf_counter() - start) / 10

        peak_mem = torch.cuda.max_memory_allocated() / (1024 ** 3)

        results.append({
            "seq_len": seq_len,
            "latency_ms": elapsed * 1000,
            "peak_memory_gb": peak_mem
        })
        print(f"  长度 {seq_len:>6d}: {elapsed*1000:.1f}ms, {peak_mem:.2f}GB")

    return results

# 运行基准测试
seq_lengths = [512, 1024, 2048, 4096, 8192, 16384]
print("Mamba 推理基准:")
mamba_results = benchmark_inference(mamba_model, tokenizer, seq_lengths)

Granite 4.0-H 架构（以 H-Small 为例）：

Layer 0:  [Mamba-2]  ─┐
Layer 1:  [Mamba-2]   │
Layer 2:  [Mamba-2]   │ 9 个 Mamba-2 层
Layer 3:  [Mamba-2]   │ 高效处理全局上下文
...                    │
Layer 8:  [Mamba-2]  ─┘
Layer 9:  [Transformer] ← 1 个 Transformer 层
Layer 10: [Mamba-2]  ─┐   精细局部解析
Layer 11: [Mamba-2]   │
...                    │ 重复 9:1 模式
Layer 18: [Mamba-2]  ─┘
Layer 19: [Transformer] ← ...
...

推理吞吐量（tokens/sec，batch_size=1，A100 GPU）：

模型               1K tokens   4K tokens   16K tokens  64K tokens
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Transformer-3B     2,100       1,800       950         OOM
Mamba-3B           2,300       2,250       2,200       2,100
Hybrid (9:1)       2,200       2,100       1,950       1,850
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

注意：Transformer 在 64K 时因 KV Cache 内存溢出 (OOM)
      Mamba 在所有长度下吞吐量几乎不变 ✅

推理时 GPU 显存占用（GB，fp16，batch_size=1）：

序列长度           Transformer-3B    Mamba-3B    节省比例
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
1K tokens          6.2               5.8         ~6%
4K tokens          7.1               5.8         ~18%
16K tokens         11.4              5.9         ~48%
64K tokens         28.6              5.9         ~79%
128K tokens        OOM               6.0         ∞
512K tokens        OOM               6.1         ∞
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

关键发现：Mamba 的内存几乎不随序列长度增长 ✅