SketchSSM 发布：写全状态、读 sketch，解码加速 2.8 倍

sehoonkim418 · x · 2026-10-08

问题：Mamba、线性注意力等模型维护固定大小循环状态，但每步解码读取完整状态成为主要瓶颈。剪枝/量化会累积误差，长推理任务即使 2–4 倍流量降低也明显掉准。
方法：SketchSSM 保持全量状态更新，只近似读取。低秩状态加权 query 近似可精确还原读输出，且基向量可离线固定；每次状态更新时读一次全状态、预计算基向量输出存入紧凑 sketch，后续每步用 query 相关系数组合 sketch 还原输出。
结果：在 Mamba-2、GDN、KDA 四个模型上，平均 rank 8 下状态访问流量降约 10 倍，四个解码基准平均精度与 FP32 全量基线持平，RULER 检索任务 recall 保持。单卡 NVIDIA B300 上相对标准 vLLM 的线性注意力 kernel 提速达 7.30×、5.02×，整体解码最高 2.8×。
论文：arXiv 2609.33051（Seoul National University 等，Kurt Keutzer 参与）。

「Infra」频道最新