Multi-Head Low-Rank Attention

2026-03-02

AI Summary

This work addresses the memory bandwidth bottleneck in long-context large language model inference, where frequent accesses to the key-value (KV) cache in high-bandwidth memory (HBM) during decoding severely limit performance. While existing low-rank methods like MLA reduce cache size, their single-head latent representation is not shardable, leading to redundant data loading and communication overhead under tensor parallelism. To overcome this limitation, we propose Multi-head Low-Rank Attention (MLRA), which introduces a sharded multi-head latent state design—enabling efficient 4-way tensor parallelism while preserving both cache compression and distributed efficiency. Experimental results demonstrate that MLRA achieves a 2.8× speedup in decoding over MLA, while maintaining state-of-the-art perplexity and downstream task performance.

Technology Category

• Machine Learning: Hardware-aware ML
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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.

Problem

• KV cache
• Tensor Parallelism
• Multi-Head Latent Attention
• distributed decoding
• memory bottleneck

Innovation

• Multi-Head Low-Rank Attention
• KV cache optimization
• Tensor Parallelism
• distributed decoding
• low-rank approximation