优化注意力机制,让大模型推理更快更省资源。
Hardware-Efficient Attention for Fast Decoding
- 用分组共享键值状态减少内存读取,提升硬件效率。
- 新方法在长文本下比FlashMLA快2倍,吞吐量提升2倍。
- 适合追求高速推理的在线服务和资源受限场景。
大语言模型解码在大规模批量处理和长上下文时,受高带宽内存中键值(KV)缓存加载的限制,导致单令牌延迟增加,且解码的顺序性限制了并行度。本文分析了计算强度、并行化与模型质量之间的相互作用,质疑现有架构是否充分挖掘现代硬件潜力。为此,重新设计注意力机制,以每字节加载数据执行更多计算,最大化硬件效率,同时保持良好的可并行扩展性。首先提出分组绑定注意力(GTA),通过合并和复用键值状态,在不损失模型质量的前提下减少内存传输。随后引入分组潜在注意力(GLA),一种支持并行的潜在注意力结构,配合底层优化实现快速解码,并维持高质量。实验表明,GTA 在匹配分组查询注意力(GQA)性能的同时,仅需约一半的KV缓存;而GLA性能媲美多头潜在注意力(MLA),且更易分片。优化后的GLA内核在查询长度超过1时,比FlashMLA快达2倍;通过降低每设备所需获取的KV缓存大小,GLA在在线服务基准测试中将端到端延迟降低,吞吐量最高提升2倍。
原文摘要 · Abstract (English)
LLM decoding is bottlenecked for large batches and long contexts by loading the key-value (KV) cache from high-bandwidth memory, which inflates per-token latency, while the sequential nature of decoding limits parallelism. We analyze the interplay among arithmetic intensity, parallelization, and model quality and question whether current architectures fully exploit modern hardware. This work redesigns attention to perform more computation per byte loaded from memory to maximize hardware efficiency without trading off parallel scalability. We first propose Grouped-Tied Attention (GTA), a simple variant that combines and reuses key and value states, reducing memory transfers without compromising model quality. We then introduce Grouped Latent Attention (GLA), a parallel-friendly latent attention paired with low-level optimizations for fast decoding while maintaining high model quality. Experiments show that GTA matches Grouped-Query Attention (GQA) quality while using roughly half the KV cache and that GLA matches Multi-head Latent Attention (MLA) and is easier to shard. Our optimized GLA kernel is up to 2$\times$ faster than FlashMLA, for example, in a speculative decoding setting when the query length exceeds one. Furthermore, by fetching a smaller KV cache per device, GLA reduces end-to-end latency and increases throughput in online serving benchmarks by up to 2$\times$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。