Multi-query och grouped-query attention

I vanlig multi-head attention (MHA) har varje huvud egna Q, K och V. KV-cachen skalar med antalet huvuden — och blir minnesboven vid lång kontext.

MQA (multi-query): alla huvuden delar ett K/V-par. Cachen krymper med antalet huvuden (32×), men kvaliteten sjunker märkbart.

GQA (grouped-query): huvudena delas i grupper som delar K/V. Med 32 Q-huvuden och 8 KV-grupper blir cachen 4× mindre med nästan oförändrad kvalitet.

GQA är standard i Llama 3, Mistral och Qwen — en av de få ändringar som ger stor vinst till nästan inget pris.

KV-cachens storlek:

$\text{bytes} = 2 \times L \times H_{kv} \times d_{head} \times T \times B \times \text{bytes/tal}$

där $L$ = lager, $H_{kv}$ = antal KV-huvuden (inte Q-huvuden), $T$ = sekvenslängd, $B$ = batch.

Exempel, 7B-modell (L = 32, d_head = 128, fp16, T = 8 192, B = 1):

Med GQA ryms fyra gånger fler samtidiga användare på samma kort — vilket direkt översätts till kostnad per användare.

Vad som offras: uttrycksförmåga. Varje grupp måste dela på samma nycklar och värden, så huvuden inom en grupp kan inte specialisera sig lika fritt. Empiriskt är GQA med 4–8 grupper nästan omöjlig att skilja från MHA, medan MQA syns i kvaliteten.

Uptraining: en MHA-modell kan konverteras till GQA genom att medelvärdesbilda K/V-projektionerna inom varje grupp och sedan fortsätta träna på ~5 % av det ursprungliga beräkningsbudgeten (Ainslie m.fl. 2023).

def kv_cache_gb(lager=32, kv_huvuden=32, d_head=128, seq=8192, batch=1, bytes_per_tal=2):
    return 2 * lager * kv_huvuden * d_head * seq * batch * bytes_per_tal / 1e9

for namn, h in (("MHA", 32), ("GQA-8", 8), ("GQA-4", 4), ("MQA", 1)):
    print(f"{namn:6s} {kv_cache_gb(kv_huvuden=h):6.2f} GB")
# MHA     4.29 GB
# GQA-8   1.07 GB
# GQA-4   0.54 GB
# MQA     0.13 GB

# Hur många samtidiga användare ryms på ett 80 GB-kort efter modellvikterna (14 GB)?
for namn, h in (("MHA", 32), ("GQA-8", 8)):
    print(namn, int((80 - 14) / kv_cache_gb(kv_huvuden=h)))
# MHA 15
# GQA-8 61