ReplaySSM: 重新思考状态存储以加速混合模型解码

空逐月AI 前沿📡 BestBlogs·AI高分精选⭐ 912026-10-13462 阅读💛 286 收藏

📌 One-Sentence Summary

ReplaySSM 通过将输出计算与状态更新解耦来优化状态空间模型（SSM），从而减少内存流量、简化投机解码并缓解量化误差。

📝 Summary

本文介绍了 ReplaySSM，这是一种提高结合了 SSM 层（如 Mamba）与注意力层的高混合模型效率的技术。传统的 SSM 解码受限于内存带宽，因为它在每一步都会更新、读取和存储状态。ReplaySSM 将状态表示为检查点（checkpoint）加上近期输入的缓冲区，并仅在“刷新（flush）”步骤中更新检查点。这种方法提供了三个主要优势：减少内存流量、通过避免昂贵的状态回滚实现高效的投机解码，以及允许使用不同的精度（递归使用高精度，读取使用低精度）来防止量化误差的传播。

💡 Main Points

将输出计算与状态更新解耦

通过将状态表示为检查点和输入缓冲区，ReplaySSM 减少了贵的 HBM 写入频率，解决了 SSM 解码中的内存带宽瓶颈。

优化投机解码

传统的 SSM 在投机解码期间需要贵的贵的回滚；ReplaySSM 在缓冲区上使用指针方法，使得验证过程显著加快且更具扩展性。

感知精度的状态

ReplaySSM 允许在递归更新时保持高精度检查点，同时生成输出时使用低精度副本，防止舍入误差在序列中累积。

💬 Key Quotes

SSM 层的设计旨在实现快速解码：它们使用固定大小的状态来代替随上下文长度增长的键值缓存（KV cache）。

ReplaySSM 将 SSM 状态表示为一个检查点加上一个近期输入的缓冲区。这使得输出计算与状态更新能够分离。

低精度副本永远不会被读回递归循环中，因此其舍入误差仅影响当前窗口的输出，而不会跨越递归进行传播。

📊 Article Meta