摘要

这项受控研究用同一份 3000 亿 token 数据,分别预训练 attention、Mamba 与 hybrid 三个 3B block-diffusion 模型,并通过统一的 cached decoding 接口对比。只有 Mamba 的状态 cache 不随上下文长度增长。在 256K tokens 下,它约占 7.5 GB、每步 6.8 毫秒;attention 则为 82 GB、29 毫秒,单流吞吐高 2.6 倍,总吞吐高 14 倍。Mamba 与 hybrid 在训练长度的 8–16 倍仍能完成检索,而 attention 在 2 倍处失效,且未测得质量损失。

为什么重要
对探索 diffusion language model 的团队,cache 架构有机会消除大部分长上下文 serving 负担。这项实验成本高,但三条模型线采用同规模预训练,控制较严。由于没有公开 checkpoint 或 serving 代码,目前更适合 WATCH,而不是直接部署。
技术细节
训练 三个 3B block-diffusion 模型,各 3000 亿 token,共用数据、tokenizer、schedule 与目标函数
架构 full attention、bidirectional Mamba-2、每五层 Mamba 插一层 attention 的 hybrid
Cache 复杂度 Mamba 状态对序列长度为 O(1);attention KV cache 为 O(L)
256K 结果 Mamba 约 7.5 GB、6.8 毫秒/步;attention 82 GB、29 毫秒/步
吞吐 256K 下 Mamba 单流 2.6 倍、总吞吐 14 倍
可用性 未链接公开 checkpoint 或实现仓库
标签
diffusion-language-modelmambalong-contextcacheinference