论文

CommunityKV:通过图分区进行高效的长上下文解码

CommunityKV: Efficient Long-Context Decoding via Graph Partitioning

模型推理KV Cache

摘要

将 Transformer 扩展到长上下文受到自注意力的二次成本和键值缓存传输的线性增长的限制。稀疏注意力通过仅检索相关标记来缓解这一问题,但当前的方法要么需要大规模训练,要么在无训练的情况下依赖于语义上粗糙的启发式方法或昂贵的聚类,而这些聚类在解码过程中难以有效更新。我们引入了 CommunityKV,一个将稀疏注意力表述为社区检测问题的框架。 CommunityKV 根据在标准预填充期间计算出的 $QK^T$ 分数构建词元图,并将该图划分为社区,以便能够检索语义一致的词元组。本地更新规则在恒定时间内将新生成的词元分配给社区,从而在整个流解码过程中实现稀疏检索,而无需全局重新分区。我们在三个长上下文基准测试中评估 Qwen3 和 Llama-3.1 模型上的 CommunityKV。每个查询头一个图,CommunityKV 可提供高达 $1.25\times$ 的密集注意力的端到端生成吞吐量,而查询组图聚合则可提供高达 $1.71\times$ 的准确度。