arXiv cs.AI· Longwei Zou, Lin Zhong·· 11 小时前AI 评分38
KVBuffer:面向线性注意力的 IO 感知推理服务机制
KVBuffer: IO-aware Serving for Linear Attention
AI 导读
KVBuffer 是一种面向线性注意力的 IO 感知推理服务机制,通过缓存近期 key 和 value,将状态更新延迟并批量应用,实现分块解码计算。该机制已在 SGLang 中针对 Qwen3.5 和 Nemotron-3 系列模型实现,批量解码最高加速 1.47x;在验证 4 个 draft token 的投机解码场景下,最大服务请求数提升 5x。
正文 · AI 翻译
摘要:线性注意力近来因其解码成本相对于上下文长度保持恒定,在长上下文推理中受到广泛关注。然而,现有服务系统通常在每个解码步骤中循环计算并更新一个大型线性注意力状态来服务线性注意力。由于该状态远大于每个 token 的键和值,循环解码会带来大量内存访问,导致线性注意力的服务效率低下。本文提出 KVBuffer,一种面向线性注意力的 IO 感知服务机制。通过缓存近期的键和值,KVBuffer 使服务系统能够以更灵活、更节省内存的方式计算线性注意力输出。在解码方面,KVBuffer 支持分块计算,通过推迟状态更新并批量应用,降低平均内存访问和解码延迟。在推测解码方面,KVBuffer 并行验证草稿 token,并避免存储临时状态。对于短上下文,KVBuffer 直接从缓存的键和值计算注意力输出,无需创建或更新线性注意力状态。我们在 SGLang 中为 Qwen3.5 和 Nemotron-3 系列模型实现了 KVBuffer。评估表明,KVBuffer 可将批量解码加速最高 1.47 倍,并在验证四个草稿 token 时将推测解码的最大服务请求数提升 5 倍。
| 学科: | 机器学习(cs.LG);人工智能(cs.AI) |
| 引用为: | arXiv:2605.19049 [cs.LG] |
| (或 arXiv:2605.19049v2 [cs.LG] 用于此版本) | |
| https://doi.org/10.48550/arXiv.2605.19049 由 DataCite 颁发的 arXiv DOI |
提交历史
来自:Longwei Zou [查看邮箱]
[v1]
2026年5月18日 星期一 19:14:52 UTC(130 KB)
[v2]
2026年11月5日 星期四 15:09:10 UTC(243 KB)
来源:arXiv cs.AI · arxiv.org