通过分层专属头机制,高效处理长文本上下文。
ZigzagAttention: Efficient Long-Context Inference with Exclusive Retrieval and Streaming Heads
- 将注意力头分为专属检索或流式头,集中于单层以减少延迟。
- 相比基线方法降低推理延迟,性能几乎无损。
- 适合需要长上下文且对延迟敏感的部署场景。
随着大语言模型(LLMs)的快速发展,处理长上下文已成为关键能力。然而,这一能力伴随部署难题,尤其是键值缓存(KV cache)内存开销增大。已有工作尝试优化KV cache内存占用,基于观察:注意力头可分为重要检索头和次要流式头。通常,识别流式头并省略其KV cache可大幅降低开销而性能影响较小。但若在同一层中同时使用两类头,会将一次大规模注意力计算拆分为两次小计算,意外引入额外的张量访问与索引延迟。为此,本文提出改进:设计一种准则,强制使每个层仅包含单一类型的头(检索或流式),从而彻底消除额外延迟,仅带来可忽略的性能下降。所提方法名为\textsc{ZigzagAttention},在降低延迟的同时保持良好性能,优于对比基线。
原文摘要 · Abstract (English)
With the rapid development of large language models (LLMs), handling long context has become one of the vital abilities in LLMs. Such long-context ability is accompanied by difficulties in deployment, especially due to the increased consumption of KV cache. There is certain work aiming to optimize the memory footprint of KV cache, inspired by the observation that attention heads can be categorized into retrieval heads that are of great significance and streaming heads that are of less significance. Typically, identifying the streaming heads and and waiving the KV cache in the streaming heads would largely reduce the overhead without hurting the performance that much. However, since employing both retrieval and streaming heads in one layer decomposes one large round of attention computation into two small ones, it may unexpectedly bring extra latency on accessing and indexing tensors. Based on this intuition, we impose an important improvement to the identification process of retrieval and streaming heads, in which we design a criterion that enforces exclusively retrieval or streaming heads gathered in one unique layer. In this way, we further eliminate the extra latency and only incur negligible performance degradation. Our method named \textsc{ZigzagAttention} is competitive among considered baselines owing to reduced latency and comparable performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。