SketchSSM：写入完整状态，从紧凑草图中读取

arXiv:2609.33051v2 公告类型：替换 摘要：混合注意力模型用线性注意力替换大多数 softmax 注意力层，从而减少 KV 缓存的增长，并支持更大的解码批次，此时循环状态访问成为主要瓶颈。ReplaySSM 通过缓冲键和值来摊销状态更新开销，但每个新查询仍需读取完整状态，即使状态在两次更新之间保持不变。我们观察到，采用状态加权的低秩查询近似能够准确保留状态读取的输出。尽管未来的查询未知，用于近似这些查询的基向量却可以离线确定。基于这一观察，我们提出了 SketchSSM，它在保留完整状态更新的同时，对读取进行近似。在每次状态更新时，SketchSSM 读取一次完整状态，预先计算这些基向量对应的输出，并将其存储在紧凑草图中。随后的每个解码步骤将草图向量与依赖于查询的系数组合起来，无需读取完整状态即可重建输出。在四个基于 Mamba-2、GDN 和 KDA 的模型上，当草图的平均秩为 8 时，SketchSSM 将状态访问流量降低至原来的约十分之一，同时在四项解码基准测试中达到与 FP32 完整状态基线相当的平均准确率，并在四项 RULER 检索任务中保持召回率。在此秩下，使用单张 NVIDIA B300 时，相对于 Standard vLLM 基线，Mamba-2、GDN 和 KDA 的线性注意力内核分别实现了 7.30 倍、5.02 倍和 5.24 倍的加速，Nemotron 3 Super 的解码吞吐量最高达到基线的 2.77 倍。