Together AI 等联合发布 Mamba-3 状态空间模型并开源内核

tl;dr

Mamba-3 是一种新的状态空间模型（SSM），以推理效率为首要设计目标——这与 Mamba-2 不同，后者针对训练速度进行了优化。关键升级包括更具表达力的递归公式、复数值状态跟踪，以及一种 MIMO（多输入、多输出）变体，可在不拖慢解码速度的情况下提升准确率。

结果是：在 1.5B 规模下，Mamba-3 SISO 在所有序列长度的 prefill+decode 延迟上均优于 Mamba-2、Gated DeltaNet，甚至 Llama-3.2-1B（Transformer）。

团队还开源了这些内核，使用 Triton、TileLang 和 CuTe DSL 混合构建，以实现最大硬件性能。

本博客同步发布于 Goomba Lab 博客，涵盖卡内基梅隆大学、普林斯顿大学、Cartesia AI 和 Together AI 的研究人员合作完成的工作。

自 2024 年年中 Mamba-2 发布以来，大多数架构已从 Mamba-1 切换过来。为什么？Mamba-2 押注训练效率是状态空间模型（SSM）的最大瓶颈，因此简化了底层 SSM 机制，相比前代实现了 2−8 倍的训练加速，从而获得更广泛的采用。

自那以后，LLM 格局开始转变。虽然预训练仍然极其重要，但更多注意力已转向后训练和部署，这两者都极其依赖推理。后训练方法的规模化，尤其是针对编程或数学的可验证奖励强化学习（RLVR），需要大量生成的 rollout，而最近，诸如 Codex、Claude Code 甚至 OpenClaw 等智能体工作流，已将推理需求推至顶峰。

尽管推理的重要性日益明显，许多线性架构（包括 Mamba-2）仍是从训练优先的角度开发的。为了加速预训练，底层 SSM 被逐步简化（例如，对角转移被简化为标量乘以单位矩阵）。这虽然带来了训练速度，却让推理步骤变得“过于简单”，并完全受限于内存——GPU 大部分时间不是在飞速计算，而是在搬运内存。

在这个推理的新时代，我们非常关注推动质量-效率前沿的边界：我们希望更好的模型运行得更快。

一个自然的问题出现了：

如果设计 SSM 时以推理为核心，它会是什么样子？

Mamba-3 模型

缺失了什么？线性模型的主要吸引力就在其名字中：由于状态大小固定，计算量随序列长度线性增长。不幸的是，天下没有免费的午餐。正是这个固定状态大小实现了高效计算，却也迫使模型将所有过去信息压缩到一个表示中，这与 Transformer 完全相反——后者通过不断增长的状态（KV 缓存）存储所有过去信息——这是一个根本性差异。那么，如果我们无法增大状态，如何让这个固定状态做更多工作？

我们看到，早期的设计简化了递归和转移矩阵以加快训练速度。然而，这一改变也降低了动态的丰富性，并导致解码受限于内存：每次令牌更新相对于内存移动所执行的计算量非常少。这为我们提供了三个可以调整的杠杆：(1) 使递归本身更具表现力，(2) 使用更丰富的转移矩阵，以及 (3) 在每次更新中添加更多并行（且几乎免费）的工作。

基于这些见解，我们通过三个核心方式改进了 Mamba-2：

1. 通过源自我们的指数梯形离散化方案的更通用递归，提高 SSM 机制的表现力，
2. 通过建模复值 SSM 系统来扩展状态跟踪能力，以及
3. 通过使用多输入多输出 (MIMO) SSM（并行建模多个 SSM）而非当前的单输入单输出 (SISO) SSM，在几乎不影响解码延迟的情况下提升模型的整体性能。

通过这三项改变，Mamba-3 在保持相似推理延迟的同时，将性能推向了前沿。

值得注意的是，这三项改变都受到了更“经典”的控制理论和状态空间模型文献的启发。

我们的工作与许多现代线性架构背道而驰，这些架构使用递归的替代解释（如线性注意力或测试时训练），这些解释不易捕捉这些概念。

架构

Mamba-2 层有什么变化？除了上述对核心 SSM 的三项方法升级外，我们还对架构进行了一些调整，使其更符合传统现代语言模型。

根据图示，你会注意到我们改变了几处。在高层次上，

归一化。 我们添加了 QKNorm 或 SSM 术语中的“BCNorm”，这在经验上稳定了 Mamba-3 模型的训练。添加此归一化使 Mamba-3 与当代 Transformer 和 Gated DeltaNet (GDN) 模型保持一致。使用 QKNorm 后，Mamba-2 中的 RMSNorm 变为可选。然而，我们经验发现，在混合模型中可能仍值得保留，因为它有助于长度外推能力。稍后会详细介绍。

告别短卷积。 我们通过结合 (1) BCNorm 后 B 和 C 上的简单偏置与 (2) 我们新的基于离散化的递归，成功摆脱了 Mamba-1/2 中烦人的短因果卷积。新的递归隐式地对隐藏状态的输入应用了卷积，我们在博客的第 2 部分展示了这一点。

短卷积真的可以移除吗？

Mamba-3 中的变化在SSM 递归内部添加了类似卷积的组件，但并不能完全与放置在SSM 递归外部的标准短卷积互换。

后者仍然可以与 Mamba-3 一起使用，但不使用的决定是基于经验的。我们发现添加回标准短卷积：

1. 不会提高性能；事实上，它略微降低了性能，并且
2. 在更多真实世界任务（例如 NIAH）上不会降低检索能力。话虽如此，如果没有短卷积，在 MQAR 这类小规模合成任务上的训练会变得有些困难。不过，由于真实世界的检索行为仍然不受影响，我们不认为这是一个重大局限。

至于原因？我们没有研究理论机制，但在论文中，我们假设 BC 偏置和指数梯形递归都执行类似的类卷积机制，这在经验上起到了与外部短卷积相同的功能。

关于短卷积的快速历史课。

短卷积如今是大多数高性能线性模型的核心组件。短卷积的各个版本最早被 H3（以“shift SSM”的形式，其灵感来自 Anthropic 关于“smeared”归纳头的工作）和 RWKV-4（通过其“token shift”机制）用于递归架构，之后由 Mamba-1 以其当前形式推广开来。

它之所以如此普遍，是因为先前的工作反复表明，短卷积既能提升经验性能，也能在理论上支持归纳式检索能力。

最后，你会注意到几个新组件，即 RoPE 和 MIMO 投影。RoPE 模块通过将复数转移解释为旋转来表达复值 SSM，从而避免了代价高昂的核函数重新实现。MIMO 投影则将 B 和 C 矩阵扩展为 MIMO SSM 所需的适当表示。

我们将在博客第二部分更详细地探讨这两者的动机和确切实现（那里有很多好东西 🎁），所以现在只需将它们视为独立的、根本性的改进，它们各自都有助于提升模型的性能和/或能力。

最后，我们的整体架构现在采用交错 MLP 层，遵循 Transformer 和其他线性模型的标准惯例。

实证结果

我们将最终的 Mamba-3 模型与其他流行的线性替代方案以及 Transformer 基线进行评估对比。

语言建模

我们发现，新的 Mamba-3 模型在多种预训练模型规模的语言建模上优于先前的 Mamba-2 模型以及 GDN 等强线性注意力替代方案。Mamba-3-SISO 可直接与先前的线性模型比较；例如，它在架构形状（模型维度、状态大小等）上与 Mamba-2 完全一致，并且训练时间相当。我们的 MIMO 版 Mamba-3 在 1B 规模下，将下游任务准确率比常规 Mamba-3 进一步提升超过 1 个百分点，但需注意 MIMO 需要更长的训练时间，而解码延迟并不会更长！

训练成本怎么会上升而推理不会？

虽然我们会在博客第二部分详细讨论这一点，但这里先给读者一个预告：

这种二分法可以追溯到训练和推理各自受计算限制与受内存限制的特性。当前的线性模型被设计为大量使用 GPU 张量核心（这是 Mamba-2 的主要贡献之一）以实现快速训练，但在解码期间，每个时间步所需的计算量如此之少，以至于硬件大部分时间都处于闲置状态。

因此，如果我们围绕仅增加每个时间步所需 FLOPs 来设计架构，推理延迟大致保持不变，因为我们可以直接使用一些空闲核心——训练则不然！

检索任务

线性模型凭借其固定大小的状态，在基于检索的任务上天然表现不佳，不如其 Transformer 对应模型。正如预期，在纯模型中，Transformer 在检索任务上更优，但 Mamba-3 在亚二次替代方案这一类中表现良好。有趣的是，加入 MIMO 进一步提升了检索性能，且未增加状态大小。

鉴于这种固有缺陷但整体建模性能强劲，

我们预测，未来线性层将主要与全局自注意力层结合使用。*

$^*$至少对于语言建模而言

将线性层通用的类记忆特性与自注意力 KV 缓存的精确类数据库存储相结合的混合模型，已被实证证明优于纯模型，同时能显著节省内存和计算，我们在此也确实发现，线性层与自注意力的结合相比普通 Transformer 能实现更好的检索。

然而，我们要强调，这些线性模型与自注意力交互的确切方式尚未被完全理解。例如，我们发现，为 Mamba-3 使用可选的输出前投影，能提升在合成 NIAH 任务上的长度泛化性能，但会略微牺牲上下文中的真实世界检索任务。此外，即便是返回范数的细节，如位置（门控前 vs 门控后）和类型（分组 vs 常规），也对由半结构化和非结构化数据组成的任务（如 FDA 和 SWDE）的准确率有不可忽视的影响。

内核无处不在

我们很期待看到人们用 Mamba-3 构建什么。为帮助促进这一点，我们开源了我们的内核，其速度与原始 Mamba-2 Triton 内核相当。

基准测试延迟

预填充延迟

预填充+解码延迟

在单块 H100-SXM 80GB GPU 上，1.5B 模型在不同序列长度下的预填充和预填充+解码（预填充与解码使用相同 token 数）延迟。所有序列长度均使用批量大小 128，报告的是三次重复的挂钟时间（以秒为单位）。在 1.5B 规模上比较模型时，Mamba-3（SISO 变体）实现了最快的预填充+解码延迟，在所有序列长度上均优于 Mamba-2、Gated DeltaNet，甚至优于拥有高度优化 vLLM 生态的 Transformer。此外，Mamba-3 MIMO 在速度上与 Mamba-2 相当，但性能强得多。

Mamba-3 SISO 基于 Triton 的 prefill 性能与 Mamba-2 几乎相同，表明新的离散化和数据依赖的 RoPE 嵌入没有引入额外开销，而 Mamba-3 MIMO 由于高效的 TileLang 实现，prefill 仅出现适度减速。两个 Mamba-3 变体的强劲 decode 性能部分归功于 CuTe DSL 实现，而 Mamba-3 组件的简洁性使这一实现变得显著更容易。

设计选择

我们花了大量时间思考如何在不牺牲易用性的前提下让内核尽可能快。最终我们使用了以下技术栈：Triton、TileLang 和 CuTe DSL。

使用 Triton 是一个相当容易的选择。它几乎是架构开发的标准（出色的 flash linear attention 仓库完全用 PyTorch 和 Triton 编写），这是有充分理由的，因为它通过可控的 tiling 和内核融合实现了比标准 PyTorch 更好的性能，同时是一种平台无关的语言。Triton 还有一些相当巧妙的特性，比如 PTX（一种面向 GPU 的汇编语言）注入，以及在 Hopper GPU 上的 Tensor Memory Accelerator 支持，用于从全局内存到共享内存的批量异步传输。

我们的 MIMO prefill 内核则使用 TileLang 开发。该变体对应的额外投影提供了一个机会，让我们可以通过在 GPU 内存层次结构中进行策略性操作来减少内存 IO。遗憾的是，Triton 没有提供我们想要的细粒度内存控制，因此我们选择了 TileLang，它允许我们显式声明和控制共享内存 tile 并创建寄存器片段，更高效地复用内存，同时仍然足够高层，使我们能够快速开发内核。

由于我们一直在强调推理和 decode 的重要性，我们决定为 decode 内核使用 CuTe DSL。通过其 Python 接口，我们能够使用 CUTLASS 的高层抽象生成底层内核。在这里，我们实际上拥有 CUDA 级别的控制，使我们能够针对硬件规格（此处为 Hopper GPU）开发高性能内核。通过对张量布局和 warp 特化的细粒度控制，我们构建了一个充分利用 GPU 所有功能的 kernel。

重要的是，这些跨越不同 GPU 抽象层次的实现之所以成为可能，要归功于 Mamba-3 简单、轻量级添加的底层算法设计及其巧妙的实例化。我们在完整发布中更深入地讨论了确切的融合结构和内核 DSL 等细节。

接下来

很高兴你读到了第 1 部分的结尾！关于我们的内核、实验结果和消融研究，有很多细节我们没时间在这篇文章中涵盖，但别担心！一切都可以在我们的论文中找到，内核已在 mamba-ssm 开源！

接下来，本系列的第二部分（也是最后一部分）将深入探讨 Mamba-3 的三项核心改进及其 SSM 基础，并给出一些我们特别感兴趣的方向。

‍

参考文献

1. Mamba: Linear-Time Sequence Modeling with Selective State Spaces  [PDF]
   Gu, A. and Dao, T., 2024.
2. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality  [PDF]
   Dao, T. and Gu, A., 2024.
3. Gated Delta Networks: Improving Mamba2 with Delta Rule  [PDF]
   Yang, S., Kautz, J. and Hatamizadeh, A., 2025.
4. Learning to (Learn at Test Time): RNNs with Expressive Hidden States  [PDF]
   Sun, Y., Li, X., Dalal, K., Xu, J., Vikram, A., Zhang, G., Dubois, Y., Chen, X., Wang, X., Koyejo, S., Hashimoto, T. and Guestrin, C., 2025.
5. Hungry Hungry Hippos: Towards Language Modeling with State Space Models  [PDF]
   Fu, D.Y., Dao, T., Saab, K.K., Thomas, A.W., Rudra, A. and Ré, C., 2023.
6. In-context Learning and Induction Heads
   Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S. and Olah, C., 2022. Transformer Circuits Thread.
7. RWKV: Reinventing RNNs for the Transformer Era  [PDF]
   Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Chung, M., Grella, M., GV, K.K., He, X., Hou, H., Lin, J., Kazienko, P., Kocon, J., Kong, J., Koptyra, B., Lau, H., Mantri, K.S.I., Mom, F., Saito, A., Song, G., Tang, X., Wang, B., Wind, J.S., Wozniak, S., Zhang, R., Zhang, Z., Zhao, Q., Zhou, P., Zhou, Q., Zhu, J. and Zhu, R., 2023.
8. Test-time regression: a unifying framework for designing sequence models with associative memory  [PDF]
   Wang, K.A., Shi, J. and Fox, E.B., 2025.
9. An Empirical Study of Mamba-based Language Models  [PDF]
   Waleffe, R., Byeon, W., Riach, D., Norick, B., Korthikanti, V., Dao, T., Gu, A., Hatamizadeh, A., Singh, S., Narayanan, D., Kulshreshtha, G., Singh, V., Casper, J., Kautz, J., Shoeybi, M. and Catanzaro, B., 2024.