MLKV: Multi-Lagen Sleutel-Waarde Koppen voor Geheugenefficiënte Transformer-Decodering

MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding

June 13, 2024

Auteurs: Zayd Muhammad Kawakibi Zuhri, Muhammad Farid Adilazuarda, Ayu Purwarianti, Alham Fikri Aji

cs.AI

Samenvatting

Auto-regressieve inferentie van transformers profiteert sterk van Key-Value (KV)-caching, maar kan leiden tot grote geheugenproblemen naarmate de modelgrootte, batchgrootte en sequentielengte op schaal toenemen. We introduceren Multi-Layer Key-Value (MLKV)-deling, een nieuwe aanpak die KV-deling uitbreidt over transformer-lagen om het geheugengebruik verder te verminderen dan mogelijk was met Multi-Query Attention (MQA) en Grouped-Query Attention (GQA). Evaluaties op verschillende NLP-benchmarks en inferentiemetrieken met behulp van bijgetrainde Pythia-160M-varianten tonen aan dat MLKV het geheugengebruik aanzienlijk vermindert met minimaal prestatieverlies, waarbij de KV-cachegrootte wordt teruggebracht tot een factor 6x vergeleken met MQA. Deze resultaten onderstrepen het potentieel van MLKV voor efficiënte implementatie van transformer-modellen op schaal. We bieden code aan op https://github.com/zaydzuhri/pythia-mlkv.

English

Auto-regressive inference of transformers benefit greatly from Key-Value (KV) caching, but can lead to major memory bottlenecks as model size, batch size, and sequence length grow at scale. We introduce Multi-Layer Key-Value (MLKV) sharing, a novel approach extending KV sharing across transformer layers to reduce memory usage beyond what was possible with Multi-Query Attention (MQA) and Grouped-Query Attention (GQA). Evaluations on various NLP benchmarks and inference metrics using uptrained Pythia-160M variants demonstrate that MLKV significantly reduces memory usage with minimal performance loss, reducing KV cache size down to a factor of 6x compared to MQA. These results highlight MLKV's potential for efficient deployment of transformer models at scale. We provide code at https://github.com/zaydzuhri/pythia-mlkv