通过共享键值注意力提升长序列大模型推理效率
MoSKA: Mixture of Shared KV Attention for Efficient Long-Sequence LLM Inference
- 区分请求独有与高频复用的上下文，动态优化缓存结构
- 共享数据处理从内存瓶颈转为计算密集型，吞吐最高提升538.7倍
- 适合高上下文复用场景，如对话系统、文档摘要
大型语言模型（LLMs）上下文长度不断增长，导致键值（KV）缓存成为严重性能瓶颈，其内存受限特性造成显著的GPU利用率下降。本文提出混合共享键值注意力（MoSKA），通过利用上下文数据的异质性，区分每请求独有的内容与大规模复用的共享序列。核心是新型共享键值注意力机制，将共享数据的注意力计算从一系列内存密集型的GEMV操作转变为单次计算密集型的GEMM，通过批量处理并发请求实现。该机制结合受MoE启发的稀疏注意力策略以剪枝搜索空间，并采用专门设计的解耦基础设施，分别优化硬件对独有与共享数据的处理。在具有高上下文复用的工作负载中，该方法相较基线吞吐量最高提升538.7倍，为可扩展的大模型推理提供了清晰的架构路径。
原文摘要 · Abstract (English)
The escalating context length in Large Language Models (LLMs) creates a severe performance bottleneck around the Key-Value (KV) cache, whose memory-bound nature leads to significant GPU under-utilization. This paper introduces Mixture of Shared KV Attention (MoSKA), an architecture that addresses this challenge by exploiting the heterogeneity of context data. It differentiates between per-request unique and massively reused shared sequences. The core of MoSKA is a novel Shared KV Attention mechanism that transforms the attention on shared data from a series of memory-bound GEMV operations into a single, compute-bound GEMM by batching concurrent requests. This is supported by an MoE-inspired sparse attention strategy that prunes the search space and a tailored Disaggregated Infrastructure that specializes hardware for unique and shared data. This comprehensive approach demonstrates a throughput increase of up to 538.7x over baselines in workloads with high context sharing, offering a clear architectural path toward scalable LLM inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。