Многоголовое внимание с низким рангом

Аннотация

Вывод больших языковых моделей с длинным контекстом ограничивается загрузкой кэша ключей-значений (KV) на этапе декодирования, где последовательная природа генерации требует многократной передачи KV-кэша из внешней высокоскоростной памяти (HBM) во внутреннюю статическую память с произвольным доступом (SRAM) на каждом шаге. Хотя многоуровневая латентная внимательность (MLA) значительно сокращает общий размер KV-кэша, она сталкивается с проблемой шардинга при распределенном декодировании с помощью тензорного параллелизма (TP). Поскольку ее единственный латентный заголовок не подлежит разделению, каждое устройство вынуждено избыточно загружать полный KV-кэш для каждого токена, потребляя чрезмерный объем памяти и сводя на нет преимущества TP, такие как шардинг весов. В данной работе мы предлагаем многоуровневую внимательность низкого ранга (MLRA), которая обеспечивает разделяемые латентные состояния для эффективного 4-стороннего TP-декодирования. Многочисленные эксперименты показывают, что MLRA достигает наилучших показателей перплексии и производительности на практических задачах, а также обеспечивает ускорение декодирования в 2,8 раза по сравнению с MLA. Код доступен по адресу https://github.com/SongtaoLiu0823/MLRA. Предобученные веса, а также данные для обучения и оценки доступны по адресу 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.