Atención Multi-Cabezal de Bajo Rango

Resumen

La inferencia de contexto largo en modelos de lenguaje grandes se ve limitada por la carga de la caché Clave-Valor (KV) durante la etapa de decodificación, donde la naturaleza secuencial de la generación requiere transferir repetidamente la caché KV de la memoria de alto ancho de banda (HBM) fuera del chip a la memoria estática de acceso aleatorio (SRAM) dentro del chip en cada paso. Si bien la atención latente multicabezal (MLA) reduce significativamente el tamaño total de la caché KV, sufre un cuello de botella de fragmentación durante la decodificación distribuida mediante paralelismo de tensores (TP). Dado que su única cabeza latente no puede particionarse, cada dispositivo se ve obligado a cargar redundantemente la caché KV completa para cada token, consumiendo un tráfico de memoria excesivo y disminuyendo los beneficios del TP, como la fragmentación de pesos. En este trabajo, proponemos la Atención Multicabezal de Bajo Rango (MLRA), que permite estados latentes particionables para una decodificación TP eficiente de 4 vías. Experimentos exhaustivos muestran que MLRA logra una perplexidad y un rendimiento en tareas específicas de vanguardia, al mismo tiempo que ofrece una aceleración de la decodificación 2.8 veces mayor que MLA. El código está disponible en https://github.com/SongtaoLiu0823/MLRA. Los pesos preentrenados, junto con los datos de entrenamiento y evaluación, están disponibles en 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.