ReplaySSM:重新思考混合模型解码中的状态存储
NVIDIA Generativ14 天前
背景
混合模型会在网络中交错使用状态空间模型(SSM)层和少量注意力层。文中提到,NVIDIA Nemotron 3、Qwen3.5 到 Qwen3.8、Kimi-K3、GLM-5.3-Flash 都采用了类似设计,其中 SSM 层可能是 Mamba-2、Gated DeltaNet(GDN)或 Kimi Delta Attention(KDA)。
这类设计的核心目标之一,是让解码阶段更快。与传统 Transformer 在生成过程中维护会随上下文长度增长的 KV cache 不同,SSM 层使用固定大小的状态来表示历史信息,因此理论上能降低长上下文解码时的存储压力。
关键问题
文章围绕 ReplaySSM 展开,重点是重新思考混合模型解码中的状态存储方式。
在混合模型中,注意力层仍然需要 KV cache,而 SSM 层则维护自己的固定大小状态。随着模型结构越来越复杂,如何存储、复用和调度这些状态,会直接影响解码效率。
ReplaySSM 试图从状态存储角度优化这一过程,而不是只关注单个算子或单层计算性能。
为什么值得关注
对于推理系统来说,解码性能通常受以下因素影响:
- 历史上下文带来的缓存开销;
- 多层状态在显存中的存放与访问方式;
- SSM 层与注意力层混合时的调度成本;
- 批处理、并发请求和长上下文场景下的内存压力。
SSM 层用固定大小状态替代随序列长度增长的 KV cache,是混合模型的重要优势。但如果状态管理方式不够高效,这一优势可能无法充分转化为端到端吞吐提升。
适合讨论的方向
- SSM 与注意力混合架构是否会成为大模型推理的新主流?
- 固定大小状态相比 KV cache,在长上下文场景中的实际收益如何衡量?
- 推理框架应如何抽象 SSM 状态,以兼容 Mamba-2、GDN、KDA 等不同实现?
- 对服务端推理而言,状态存储优化和算子优化哪个更关键?
这篇文章适合关注大模型推理优化、混合架构和长上下文解码性能的开发者阅读。
