Together.ai 发布了 Mamba-3,这是一种专为推理工作负载而全新构建的开源状态空间模型 (SSM),标志着线性架构设计中的重大哲学转变。据报道,新模型的性能优于其前身 Mamba-2,并在 16K 序列长度下与 Transformer 解码速度相匹配,随着代理 AI 工作流程将推理需求推向前所未有的水平。
在 16,384 序列长度下,Mamba-3 的 SISO(单输入、单输出)变体表现出显着的性能提升,预填充和解码的时钟速度为 140.61 秒,而 Mamba-2 为 149.02 秒,在 vLLM 上运行的 Llama-3.2-1B 则为惊人的 976.50 秒。这意味着在相同的 H100 GPU 硬件上,性能比 Transformer 基准快近7 倍。
为什么推理优化现在很重要
发布时间与不断变化的行业优先事项保持一致。虽然 Mamba-2 的重点是在 2024 年中期提高训练速度,但当前的形势使推理成为主要瓶颈。为编码和数学应用提供可验证奖励的强化学习需要大规模部署,Codex、Claude Code 和 OpenClaw 等工具强化了推理需求,而不是预训练。
以前的线性架构简化了其底层机制以加速训练,通常使推理步骤“过于简单”且受内存限制,导致 GPU 在数据洗牌而不是计算上花费过多时间。
三大核心架构改进
Mamba-3 通过基于经典控制理论的根本性改变解决了这些限制:
指数梯形离散化创建了更具表现力的递归,消除了困扰早期 Mamba 版本的短因果卷积——自从 H3 和 RWKV-4 普及以来,该组件已成为线性模型的标准组件。
复值 SSM 系统扩展了状态跟踪功能,使模型能够处理 Mamba-2 无法可靠解决的合成任务,例如奇偶校验和算术推理。
多输入多输出 (MIMO) 架构并行运行多个 SSM。与标准 Mamba-3 相比,MIMO 变体在 1B 规模下将下游精度提高了 1 个百分点以上,其关键优势是虽然训练时间更长,但解码延迟不受影响。
最后一点强调了一个关键的工程见解:训练受计算限制,而推理受内存限制。每个时间步添加 FLOP 对推理延迟的影响最小,因为空闲 GPU 核心可以有效吸收额外的工作负载。
性能基准
在下游语言建模评估中,Mamba-3 在预训练模型规模上的表现优于 Mamba-2 和门控 DeltaNet。 SISO 变体与 Mamba-2 的架构相匹配,同时提供卓越的精度,而 MIMO 配置则进一步提升性能。
检索任务呈现出更加细致入微的画面。纯线性模型在这个领域自然表现不佳——它们固定大小的状态无法匹配不断增长的 KV 缓存以实现精确召回。然而,Mamba-3 在次二次替代方案中保持了具有竞争力的性能,MIMO 在不增加状态大小的情况下提高了检索能力。
开发团队预测,将线性层与全局自注意力相结合的混合模型将主导未来的语言建模。他们的实验表明,这种组合在检索任务上优于普通 Transformers,同时保持效率提升。
开源可用性
内核可立即在 mamba-ssm 存储库中使用,该存储库根据具体操作跨 Triton、TileLang 和 CuTe DSL 构建。该堆栈反映了务实的工程考虑:Triton 用于标准架构开发,TileLang 用于 MIMO 预填充期间的细粒度内存控制,CuTe DSL 用于在解码期间最大化 Hopper GPU 性能。
业界对 SSM 架构的兴趣似乎正在加速增长,NVIDIA 最近发布的 Nemotron 3 Super 就证明了这一点,该产品在混合配置中利用了 Mamba-2 层。 Mamba-3 的推理优先方法可以显着加快在生产环境中的采用,在生产环境中,令牌生成速度直接影响运营成本和用户体验。
完整的研究论文可在 arXiv 上找到,第二篇博客文章详细介绍了预计很快将推出的三个核心改进的数学基础。
