Meerhoofdige Lagedimensionele Aandacht

Samenvatting

Lang-context inferentie in grote taalmodellen wordt beperkt door het laden van de Key-Value (KV) cache tijdens de decoderingfase, waarbij de sequentiële aard van genereren herhaaldelijk overdracht van de KV-cache van off-chip High-Bandwidth Memory (HBM) naar on-chip Static Random-Access Memory (SRAM) vereist bij elke stap. Hoewel Multi-Head Latent Attention (MLA) de totale KV-cachegrootte aanzienlijk reduceert, kampt het met een sharding-beperking tijdens gedistribueerde decodering via Tensor Parallelism (TP). Omdat de enkele latente head niet kan worden gepartitioneerd, wordt elk apparaat gedwongen om de volledige KV-cache redundant te laden voor elke token, wat excessief geheugenverkeer verbruikt en TP-voordelen zoals gewichtssharding vermindert. In dit werk stellen we Multi-Head Low-Rank Attention (MLRA) voor, dat partitieerbare latente toestanden mogelijk maakt voor efficiënte 4-way TP-decodering. Uitgebreide experimenten tonen aan dat MLRA state-of-the-art perplexiteit en prestaties op downstreamtaken bereikt, terwijl het ook een 2,8x versnelling in decoderingssnelheid behaalt ten opzichte van MLA. Code is beschikbaar op https://github.com/SongtaoLiu0823/MLRA. Vooraf getrainde gewichten, samen met de trainings- en evaluatiedata, zijn beschikbaar op 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.