提出MLRA，让大模型长文本生成更快更省内存。
Multi-Head Low-Rank Attention
- 用多头低秩注意力替代单头潜变量，支持分布式并行分片。
- 在4路并行下实现2.8倍解码速度提升，性能达当前最优。
- 适合需要高效长文本生成的部署场景，如对话系统、文档生成。
大型语言模型在长文本推理中受制于解码阶段的键值（KV）缓存加载，因生成具有顺序性，每步需从片外高带宽内存（HBM）重复传输KV缓存至片上静态随机存取内存（SRAM）。尽管多头潜在注意力（MLA）显著减少了总KV缓存大小，但在通过张量并行（TP）进行分布式解码时仍存在分片瓶颈——其单个潜在头无法被分割，导致每个设备必须冗余加载完整缓存，造成大量内存流量并削弱了权重分片带来的TP优势。本文提出多头低秩注意力（MLRA），实现了可分片的潜在状态，支持高效的4路TP解码。大量实验表明，MLRA在困惑度和下游任务性能上均达到当前最优水平，同时相比MLA实现2.8倍解码速度提升。代码已开源：https://github.com/SongtaoLiu0823/MLRA；预训练权重及数据集见：https://huggingface.co/Soughing/MLRA。
原文摘要 · Abstract (English)
Long-context inference in large language models is bottlenecked by Key--Value (KV) cache loading during the decoding stage, where the sequential nature of generation requires repeatedly transferring the KV cache from off-chip High-Bandwidth Memory (HBM) to on-chip Static Random-Access Memory (SRAM) at each step. While Multi-Head Latent Attention (MLA) significantly reduces the total KV cache size, it suffers from a sharding bottleneck during distributed decoding via Tensor Parallelism (TP). Since its single latent head cannot be partitioned, each device is forced to redundantly load the complete KV cache for every token, consuming excessive memory traffic and diminishing TP benefits like weight sharding. In this work, we propose Multi-Head Low-Rank Attention (MLRA), which enables partitionable latent states for efficient 4-way TP decoding. Extensive experiments show that MLRA achieves state-of-the-art perplexity and downstream task performance, while also delivering a 2.8$\times$ decoding speedup over MLA. Code is available at https://github.com/SongtaoLiu0823/MLRA. Pretrained weights, along with the training and evaluation data, are available at https://huggingface.co/Soughing/MLRA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。