Multi-Head Low-Rank Attention
2026-03-02 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address a slowdown in large language models caused by moving key-value data during text generation. They note that an existing method, Multi-Head Latent Attention (MLA), reduces memory size but can’t split work well across devices, leading to inefficiency. To fix this, they introduce Multi-Head Low-Rank Attention (MLRA), which allows better splitting across four devices and speeds up decoding. Their experiments show MLRA improves performance and is almost three times faster than MLA.
large language modelskey-value cachedecoding stageMulti-Head Latent Attention (MLA)Tensor Parallelism (TP)Multi-Head Low-Rank Attention (MLRA)perplexitydistributed decodinghigh-bandwidth memory (HBM)static random-access memory (SRAM)
Authors
Songtao Liu, Hongwu Peng, Zhiwei Zhang, Zhengyu Chen, Yue Guo
Abstract
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.