Attenzione Multi-Testa a Basso Rango

Abstract

L'inferenza a contesto lungo nei grandi modelli linguistici è limitata dal caricamento della cache Chiave-Valore (KV) durante la fase di decodifica, dove la natura sequenziale della generazione richiede il trasferimento ripetuto della cache KV dalla memoria ad alta larghezza di banda (HBM) off-chip alla memoria statica ad accesso casuale (SRAM) on-chip a ogni passo. Sebbene l'attenzione latente multi-testa (MLA) riduca significativamente le dimensioni totali della cache KV, essa soffre di un collo di bottiglia dovuto allo sharding durante la decodifica distribuita tramite parallelismo tensoriale (TP). Poiché la sua singola testa latente non può essere partizionata, ogni dispositivo è costretto a caricare ridondantemente l'intera cache KV per ogni token, consumando un traffico di memoria eccessivo e diminuendo i vantaggi del TP come lo sharding dei pesi. In questo lavoro, proponiamo l'attenzione a basso rango multi-testa (MLRA), che abilita stati latenti partizionabili per una decodifica efficiente con TP a 4 vie. Esperimenti estensivi mostrano che MLRA raggiunge uno stato dell'arte in termini di perplexity e prestazioni su task downstream, offrendo anche un aumento di velocità di decodifica di 2,8 volte rispetto a MLA. Il codice è disponibile all'indirizzo https://github.com/SongtaoLiu0823/MLRA. I pesi pre-addestrati, insieme ai dati di addestramento e valutazione, sono disponibili su https://huggingface.co/Soughing/MLRA.

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.8times 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.