AI 原生开放教材

外观

AI 工程从零到一（中文）源码研读·中文原文 ↗原件 ↓仓库 ↗

Jamba —— 混合 SSM-Transformer ​

状态空间模型（SSM）和 transformer 想要不同的东西。Transformer 以二次成本的注意力换质量。SSM 以一个递归换线性时间推理和恒定内存，但质量落后。AI21 的 Jamba（2024 年 3 月）和 Jamba 1.5（2024 年 8 月）把它们放进同一个模型：每 7 个 Mamba 层配 1 个 Transformer 层，每隔一个块用 MoE，以及一个能塞进单张 80GB GPU 的 256k context 窗口。Mamba-3（ICLR 2026）用复值状态空间和 MIMO 投影收紧了 SSM 那一侧。本节课端到端读两个架构，并解释为什么混合配方在纯 SSM 和纯 Transformer 的长 context 尝试都没撑住时，挺过了三年的规模化。

类型： Learn 语言： Python（stdlib，层混合计算器） 前置要求： 阶段 10 · 14（开放模型架构），阶段 10 · 17（原生稀疏注意力） 预计时间： ~60 分钟

学习目标 ​

  • 解释一个 Jamba 块里的三种原语——Transformer 层、Mamba 层、MoE——以及 1:7:隔一交错的配方。
  • 高层陈述一个 SSM 的递归长什么样，以及为什么它能实现恒定内存推理。
  • 计算一个 Jamba 模型在 256k context 下的 KV cache 占用，和一个纯 Transformer 模型会需要的相比。
  • 说出三个 Mamba-3 创新（指数-梯形离散化、复值状态更新、MIMO）以及每个针对的问题。

问题背景 ​

注意力对序列长度是二次的。状态空间模型是线性的。这个差异复合：在 256k token 下，一张 Transformer 注意力图每头 650 亿个条目；一个 SSM 的递归状态不管序列多长都是定长的。

纯 SSM 模型（Mamba、Mamba-2）在小规模匹配 Transformer 困惑度，但在状态追踪任务上落后，并在某些类别的 in-context 检索上失败。直觉是：SSM 把历史压进一个定长状态，历史长时信息泄漏。注意力精确地记住一切，但付二次成本。

显而易见的修法：两个都用。把 Transformer 层放在精确召回要紧的地方。其他地方用 SSM 层。调比例。Jamba 是第一个大规模发布这套混合配方的生产级模型（总 52B，激活 12B，256k context，单张 80GB GPU）。Jamba 1.5 把家族扩到总 398B / 激活 94B。Mamba-3（ICLR 2026）是当前最好的纯 SSM 基线，混合模型可以围绕它重建。

本节课读这三篇论文，产出 "挑对比例" 的心智模型。

核心概念 ​

一页讲完 SSM ​

一个状态空间模型经一个定长状态 h 处理序列 x_1, ..., x_N：

h_t = A h_{t-1} + B x_t
y_t = C h_t

每一步状态经一个线性动力学 A 演化，吃输入 B x_t，发出输出 C h_t。A, B, C 可以是学来的。注意那个关键性质：计算 y_t 只需要 h_\{t-1\} 和 x_t，不需要任何更早的 x。内存恒定。推理每 token O(1)。

建模质量的诀窍在 A 的结构。S4（Gu 2021）用了一个高度结构化的矩阵，训练时能作为长卷积高效求值。Mamba（Gu, Dao 2023）把固定的 A, B, C 换成数据相关的（那个 "选择性" 部分）。Mamba-2（2024）进一步简化了结构。Mamba-3（2026）在特定地方重新加回复杂性。

关键性质：对一个解码器 LLM，一个 SSM 层是注意力层的直接替代品，用定长的每层状态而不是增长的 KV cache。

Jamba 块 ​

一个 Jamba 块按两个数字交错层：

  • l：注意力对 Mamba 的比例。Jamba 用 l = 8，意思是每 7 个 Mamba 层配 1 个 Transformer 层（每组 7 个 Mamba + 1 个 Attention = 8 层）。
  • e：MoE 频率。Jamba 用 e = 2，意思是每隔一层应用 MoE。

一个块内的层序列：

M  M  M  M  M  M  M  A    (7 Mamba + 1 Attention)
|  M  |  M  |  M  |  M    (| 标记应用了 MoE 的地方)

每个 Jamba 块是 8 层。4 块深（共 32 层）时，你得到 28 个 Mamba 和 4 个 Attention 层。其中 16 个用 MoE。

为什么是 1:7 比例 ​

AI21 跑了消融：什么样的注意力对 Mamba 比例在他们的长 context eval 上给出最佳的每参数困惑度 和 in-context 召回？

  • 注意力太多（1:1）：质量上升但内存和速度退化。
  • 注意力太少（1:15）：内存极好但 in-context 检索失败。
  • 甜点区：1:7 或 1:8。

直觉是：Transformer 层处理精确召回和状态追踪。Mamba 层处理廉价的大宗处理。

位置编码 ​

Mamba 层本身就感知位置（经递归）。最初基于 Mamba 的混合里的注意力层不用 RoPE——SSM 层提供位置信息。Jamba 1.5 给注意力层加了 RoPE 以做更长 context 的泛化，这是基于经验长 context 评估的事后改进。

内存预算 ​

对一个 Jamba-1 形状（32 层：28 Mamba + 4 Attention，hidden 4096，32 个注意力头）：

  • KV cache（仅注意力层）：256k BF16 下 2 * 4 * 32 * 128 * 256k * 2 = 8.4 GB。只有那 4 个注意力层有贡献。
  • SSM 状态：每 token 前缀 28 * hidden * state_size，但这是每层定长的，不随序列长度增长。典型 Mamba 状态每特征 16、hidden 4096：总共 28 * 4096 * 16 * 2 = 3.7 MB。

和一个 32 层、相同 hidden、32 头全 MHA 的纯 Transformer 对比：256k BF16 下 2 * 32 * 32 * 128 * 256k * 2 = 128 GB。KV cache 减少 8 倍。即便对比大多数 2024 年模型用的 GQA(8) 基线（2 * 32 * 8 * 128 * 256k * 2 = 32 GB），Jamba 的 1:7 混合在 16 GB 仍小 2 倍。

这就是 AI21 说的 "256k context 在单张 80GB GPU 上"。一个全 MHA 纯 Transformer 的 KV cache 装不下；即使 GQA 基线也不给权重和激活留空间；Jamba 的能装下。

Mamba-3：2026 年的纯 SSM 基线 ​

Mamba-3（ICLR 2026，arXiv:2603.15569）在纯 SSM 那一侧引入三个创新：

  1. 指数-梯形离散化。 用一个更有表达力的递归替换 Mamba-2 里的欧拉法离散化。在核心递归内对状态-输入应用类卷积操作，而不是作为对 x_t 的外层卷积。

  2. 复值状态更新。 之前的 Mamba 把状态矩阵从复（S4）减到实对角（Mamba）再到缩放单位（Mamba-2）。Mamba-3 重新加回复值——等价于对状态做一个数据相关的旋转 embedding。这恢复了之前实值简化所损失的状态追踪能力。

  3. 多输入多输出（MIMO）投影。 不用每特征的标量投影，而用矩阵值投影。提升建模能力和推理时硬件利用率，而不增加 decode 延迟。

在 1.5B 参数下，Mamba-3 比 Gated DeltaNet 把平均下游准确率提升 0.6 个点；MIMO 变体再加 1.2，总计 1.8 个点的提升。在相同状态大小下，Mamba-3 用一半的状态匹配 Mamba-2。

Mamba-3 还没在大规模生产混合里上线——但它是下一个 Jamba 级模型 SSM 那一侧的显而易见的候选。

什么时候上混合 ​

混合在以下情况胜出：

  • Context 长到纯 Transformer KV cache 变得痛苦（64k+）。
  • 任务混合短程结构（对 SSM 好）和长程召回（需要 Transformer）。
  • 你想部署在单 GPU 内存预算上，那里光 Transformer KV cache 就装不下。

混合在以下情况败下：

  • Context 短（16k 以下）。SSM 开销被浪费；纯 Transformer 就好。
  • 任务需要处处对处处的注意力（深度推理、多文档交叉引用）。混合里注意力层的稀疏性伤害它。
  • 你在扩到万亿参数前沿模型。纯 Transformer + MLA + MoE（DeepSeek-V3 风格）目前在赢能力竞赛。

竞争格局 ​

模型家族规模独特卖点
Mamba-2纯 SSM3B线性时间，恒定内存
Jamba混合52B/12B80GB 上 256k
Jamba 1.5 Large混合398B/94B企业级长 context
Mamba-3纯 SSM1.5B（论文）状态追踪恢复
DeepSeek-V3纯 Transformer + MoE671B/37B前沿能力

2026 年格局：纯 Transformer MoE 主导前沿，但混合占据 256k 以上 context 的利基。Mamba-3 的状态追踪胜利可能在下一代把混合比例推得更低（更多 SSM，更少注意力）。

figure
swiglu-ffn

实际使用 ​

code/main.py 是一个混合架构的内存计算器。给定一个 SSM-Transformer 比例和一个 hidden-size / 层数配置，它计算：

  • 目标 context 下的 KV cache。
  • SSM 状态内存。
  • 一系列模型形状在 context N 下的总内存。

计算器支持：

  • 纯 Transformer 基线（KV cache 随 N 增长）。
  • Jamba 风格 1:7 混合。
  • 纯 SSM（完全没有 KV cache）。

数字对已发布形状直接取自 Jamba-1 和 Jamba-1.5 论文，对假想变体外推。

真实部署的集成考量：

  • 大多数生产推理服务器（vLLM、SGLang）支持 Jamba 和 Mamba。检查具体版本。
  • 在 256k context 下，Jamba 的内存优势体现在并发请求吞吐上。同样的显存你能塞下比 Transformer 序列更多的 Jamba 序列。
  • 作为独立模型的 Mamba-3 还没在生产里上线——1.5B 的研究预览。

拿去用 ​

本节课产出 outputs/skill-hybrid-picker.md。给定一个工作负载规格（context 长度画像、任务混合、内存预算），它在纯 Transformer、Jamba 风格混合和纯 SSM 之间推荐，并对内存和质量权衡给出明确推理。

练习 ​

  1. 跑 code/main.py 计算一个 32 层纯 Transformer（hidden 4096，32 头）和一个相同形状的 Jamba-1 混合在 256k context 下的 KV cache。验证 AI21 论文宣称的 ~8 倍内存减少。

  2. 改计算器去建模一个 1:3 混合（4 Mamba : 1 Attention）和一个 1:15 混合（14 Mamba : 1 Attention）。画 KV cache vs 比例。在什么比例下 KV cache 等于 SSM 状态内存？

  3. 读 Jamba 论文（arXiv:2403.19887）第 3 节。解释为什么 AI21 用 Mamba-1 而不是 Mamba-2，尽管 Mamba-2 更快。提示：混合消融部分记录了这个。

  4. 计算 Jamba 1.5 Large（总 398B，激活 94B）里 MoE-隔层 的参数开销。把激活比例和 DeepSeek-V3（37B/671B）对比，解释为什么 Jamba 的架构把激活比例推得更高。

  5. 读 Mamba-3 论文（arXiv:2603.15569）第 3 节。用三句话解释为什么复值状态更新等价于一个数据相关的旋转 embedding。把答案和阶段 7 · 第 04 课的 RoPE 推导联系起来。

关键术语 ​

术语人们怎么说它实际是什么
状态空间模型（SSM）"带定长状态的递归"一个带习得递归 h_t = A h_\{t-1\} + B x_t 的层；每 token 恒定内存
选择性 SSM"Mamba 的诀窍"数据相关的 A、B、C 参数，让模型在线性时间下有类门控的选择性
注意力对 Mamba 比例"有多少注意力层"在 Jamba 里，l = 8 意味着每 7 个 Mamba 层 1 个注意力层
Jamba 块"那个 8 层组"一个注意力 + 七个 Mamba + 在交替位置的 MoE
SSM 状态"那个隐藏缓冲"Mamba 层用来替代 KV cache 的定长每层状态
256k context"Jamba 的旗舰数字"Jamba-1 能塞进单张 80GB GPU 的序列长度；纯 Transformer 在那个尺寸做不到
Mamba-3"2026 纯 SSM"当前最好的纯 SSM 架构，带复状态 + MIMO；混合模型重建围绕的基线
MIMO"多输入多输出"Mamba-3 的创新，用矩阵值投影而非每特征标量
指数-梯形离散化"Mamba-3 的递归"一个更有表达力的递归，涵盖了 Mamba-2 的欧拉法离散化
混合架构"混合注意力和 SSM"任何交错 Transformer 和 SSM 层的模型；Jamba 是生产原型

延伸阅读 ​