通过混合注意力机制减少缓存占用,提升长文本推理速度。
Inference-Friendly Models With MixAttention
- 结合滑动窗口与层间缓存共享,仅存储最近部分标记。
- 内存占用降低60%,长序列推理速度提升2.1倍。
- 适合需要高并发、长上下文的部署场景。
键值(KV)缓存大小在现代语言模型推理中至关重要,直接影响最大上下文长度和并发请求数量。KV缓存大小随注意力头数量和处理标记数线性增长,导致内存消耗增加和推理变慢。本文探索了MixAttention架构,该结构源自Character.AI的博客文章,结合滑动窗口注意力(仅存储近期部分标记)与层间缓存共享。实验表明,MixAttention在不牺牲短/长上下文任务性能的前提下,显著降低内存使用并提升推理速度。我们还评估了多种配置,识别出在保持评测指标质量的同时优化资源效率的方案。
原文摘要 · Abstract (English)
The size of the key-value (KV) cache plays a critical role in determining both the maximum context length and the number of concurrent requests supported during inference in modern language models. The KV cache size grows proportionally with the number of attention heads and the tokens processed, leading to increased memory consumption and slower inference for long inputs. In this work, we explore the use of MixAttention, a model architecture modification closely related to a blog published by Character.AI. MixAttention combines sliding window attention, where only a small subset of recent tokens is stored in the KV cache, with KV cache sharing across layers. Our experiments demonstrate that MixAttention significantly reduces memory usage and improves inference speed without sacrificing model performance in both short and long-context tasks. We also explore various configurations of this architecture, identifying those that maintain quality across evaluation metrics while optimizing resource efficiency.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。