tl;dr
Mamba-3 是一种新的状态空间模型（SSM），其首要目标是推理效率——这与 Mamba-2 不同，后者针对训练速度进行了优化。关键升级包括更具表现力的递推公式、复数值状态跟踪，以及一种多输入多输出（MIMO）变体，可在不降低解码速度的情况下提高准确性。
结果：在 1.5B 规模下，Mamba-3 SISO 在所有序列长度上的预填充+解码延迟均优于 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），需要大量生成的展开，而最近，代理工作流（如 Codex、Claude Code 甚至 OpenClaw）将推理需求推向了顶峰。
尽管推理的重要性日益凸显，但许多线性架构（包括 Mamba-2）都是从训练优先的角度开发的。为了加速预训练，底层 SSM 被逐步简化（例如，对角转移被简化为标量乘以单位矩阵）。虽然这带来了训练速度，但使得推理步骤“过于简单”且完全受内存限制——GPU 大部分时间不是在高速计算，而是在移动内存。
在这个推理新时代，我们非常关注突破质量-效率边界：我们希望更好的模型运行更快。
一个自然的问题出现了：
以推理为设计目标的 SSM 会是什么样子？
Mamba-3 模型
缺少什么？ 线性模型的主要吸引力在于其名称：计算量随序列长度线性扩展，因为状态大小固定。不幸的是，没有免费的午餐。同样的固定状态大小虽然实现了高效计算，却迫使模型将所有过去信息压缩到一个表示中，这与 Transformer 完全相反——Transformer 通过不断增长的状态（KV 缓存）存储所有过去信息——这是一个根本性的区别。那么，如果我们不能增长状态，如何让这个固定状态做更多工作呢？
我们看到，早期的设计简化了递推和转移矩阵以加快训练速度。然而，这种变化也降低了动态的丰富性，并使解码受限于内存：每次令牌更新相对于内存移动执行的计算非常少。这为我们提供了三个可以调整的杠杆：(1) 使递推本身更具表现力，(2) 使用更丰富的转移矩阵，以及 (3) 在每次更新中添加更多并行（且几乎免费）的工作。
基于这些见解，我们通过三种核心方式改进了 Mamba-2：
- 通过从我们的指数梯形离散化方案推导出的更通用的递推，提高了 SSM 机制的表现力，
- 通过建模复值 SSM 系统扩展了状态跟踪能力，以及
- 通过使用多输入多输出 (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 一起使用，但经验性地决定不使用。我们发现重新添加标准短卷积：
- 不会提高性能；实际上，它略微降低了性能，并且
- 不会降低在更现实任务（例如 NIAH）上的检索能力。也就是说，如果没有短卷积，在像 MQAR 这样的小规模合成任务上进行训练会变得有些困难。然而，由于现实世界的检索行为不受影响，我们不认为这是一个主要限制。
至于原因？我们没有研究理论机制，但在论文中，我们假设BC偏置和指数-梯形递归都执行了类似的卷积机制，经验上它们与外部短卷积功能相同。
关于短卷积的快速历史课。
短卷积如今是大多数高性能线性模型的核心组件。短卷积的变体最早由H3（以“移位SSM”的形式，受Anthropic的“模糊”归纳头工作启发）和RWKV-4（通过其“token移位”机制）在循环架构中使用，随后由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规模上进一步将下游任务准确率提升了超过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的预填充保持了与Mamba-2几乎相同的性能，表明新的离散化和数据相关的RoPE嵌入没有引入额外开销，而Mamba-3 MIMO由于其高效的TileLang实现，仅在预填充时出现适度减速。两种Mamba-3变体的强解码性能部分归功于CuTe DSL实现，而Mamba-3组件的简洁性使这一实现变得显著容易。
设计选择
我们花了大量时间思考如何在不影响易用性的前提下尽可能快地实现内核。最终我们采用了以下技术栈：Triton、TileLang 和 CuTe DSL。
使用 Triton 是一个相当容易的选择。它几乎是架构开发的标准（优秀的 flash linear attention 仓库完全基于 PyTorch 和 Triton），这是有充分理由的，因为它通过支持可控的分块和内核融合，实现了比标准 PyTorch 更好的性能，同时是一种平台无关的语言。Triton 还有一些非常巧妙的功能，比如 PTX（一种面向 GPU 的汇编语言）注入，以及其张量内存加速器支持（在 Hopper GPU 上），用于从全局内存到共享内存的批量异步传输。
我们的 MIMO 预填充内核是使用 TileLang 开发的。与变体对应的额外投影提供了一个机会，我们可以通过跨 GPU 内存层次结构的策略性操作来减少内存 IO。不幸的是，Triton 没有提供我们所需的内存控制粒度，因此我们选择了 TileLang，它允许我们显式声明和控制共享内存块，创建寄存器片段，更高效地重用内存，同时仍然足够高层，使我们能够快速开发内核。
由于我们一直在强调推理和解码的重要性，我们决定使用 CuTe DSL 来开发解码内核。通过其 Python 接口，我们能够使用 CUTLASS 的高层抽象生成底层内核。在这里，我们几乎拥有 CUDA 级别的控制，从而能够开发针对硬件规格（本例中为 Hopper GPU）定制的高性能内核。通过对张量布局和 warp 特化的细粒度控制，我们构建了一个充分利用 GPU 所有功能的内核。
重要的是，这些跨不同 GPU 抽象级别的实现得益于 Mamba-3 简单轻量的附加组件及其巧妙实例化的底层算法设计。我们在完整发布版中更深入地讨论了诸如精确融合结构和内核 DSL 等细节。
下一步
很高兴你读完了第一部分！关于我们的内核、实验结果和消融实验，还有很多细节我们没来得及在这篇文章中涵盖，但别担心！所有内容都可以在我们的论文中找到，内核已在 mamba-ssm 开源！
接下来，本系列第二篇（也是最后一篇）将深入探讨 Mamba-3 的三个核心改进及其 SSM 基础，并给出我们特别感兴趣的一些方向。
参考文献
Mamba: Linear-Time Sequence Modeling with Selective State Spaces[PDF]
Gu, A. and Dao, T., 2024.Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality[PDF]
Dao, T. and Gu, A., 2024.Gated Delta Networks: Improving Mamba2 with Delta Rule[PDF]
Yang, S., Kautz, J. and Hatamizadeh, A., 2025.学习（在测试时学习）：具有表达性隐藏状态的RNN[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.饥饿的河马：迈向基于状态空间模型的语言建模[PDF]
Fu, D.Y., Dao, T., Saab, K.K., Thomas, A.W., Rudra, A. and Ré, C., 2023.上下文学习与归纳头
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.RWKV：为Transformer时代重塑RNN[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.测试时回归：设计具有关联记忆的序列模型的统一框架[PDF]
Wang, K.A., Shi, J. and Fox, E.B., 2025.基于Mamba的语言模型实证研究[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.