Technische Samenvatting: GQLA – Group-Query Latent Attention voor Hardware-Adaptieve Decoding van Grote Taalmodellen

1. Probleemstelling

Moderne autoregressieve decoding van Grote Taalmodellen (LLM) wordt fundamenteel geblokkeerd door het verkeer van de Key-Value (KV) cache, aangezien elke gegenereerde token het lezen van de volledige geschiedenis van de in het geheugen buiten de chip opgeslagen sleutels en waarden vereist. Hoewel recente architecturen zoals Multi-head Latent Attention (MLA), gebruikt in DeepSeek-V2/V3, state-of-the-art reductie van de KV-cache bereiken door sleutels en waarden gezamenlijk te comprimeren tot een low-rank latent, lijden ze aan drie gekoppelde hardware- nadelen bij implementatie op diverse inferentiehardware:

1. Hardware Koppeling: De getrainde gewichten van MLA onthullen slechts één decoding-pad (een "geabsorbeerde" Multi-Query Attention of MQA-vorm). Dit pad is perfect afgestemd op de reken- naar bandbreedte-verhouding van high-end GPU's zoals de NVIDIA H100 (ridge $\approx$ 295 FLOPs/byte). Echter, op export-beperkte of commodity GPU's zoals de H20 (die bandbreedte behoudt maar de rekenkracht met $\sim$7x reduceert, waardoor de ridge daalt tot $\approx$ 37 FLOPs/byte), wordt dit pad ernstig rekenkracht-beperkt, wat resulteert in een slechte doorvoer.
2. Verlies van Tensor Parallelisme (TP): De geabsorbeerde MQA-vorm voert alle query-koppen via een enkele gedeelde latent KV. Deze structuur verhindert tensor parallelisme langs de kop-as, aangezien elk apparaat de volledige latent moet repliceren, wat de schaalbaarheid beperkt.
3. Zero Multi-Token Prediction (MTP) Winst: MTP, dat meerdere tokens per stap genereert om de rekenintensiteit te verhogen, duwt de al rekenkracht-beperkte H20-decoding verder voorbij zijn efficiency-ridge, wat resulteert in nul doorvoerwinst. Omgekeerd overschrijdt MTP op de H100 de ridge, wat de efficiency van de winst vermindert.

2. Methodologie: Group-Query Latent Attention (GQLA)

De auteurs stellen Group-Query Latent Attention (GQLA) voor, een minimale architecturale modificatie van MLA die de gezamenlijke low-rank latent compressie behoudt terwijl het twee algebraïsch equivalente decoding-paden blootlegt over dezelfde getrainde parameters.

Architectuur

GQLA indexeert de up-projectiematrices ($W_{UK}, W_{UV}$) via $g$ groepen in plaats van ze te repliceren over alle $h_q$ query-koppen. Dit ontwerp maakt twee verschillende uitvoeringspaden mogelijk:

1. MQA-Absorb Pad (Geheugen-vriendelijk):

   • Mecanisme: Identiek aan MLA. De up-projecties worden geabsorbeerd in de query- en output-projecties. Alle $h_q$ koppen attenderen direct op de compacte latent cache ($cKV$).
   • Cache: Bewaart alleen de latent ($r_{kv}$) en een gedeelde RoPE-sleutel ($d^R_h$).
   • Doelhardware: Geoptimaliseerd voor H100-klasse hardware waar de rekenintensiteit ($\approx$ 242 FLOPs/byte voor $s_q=1$) net onder de geheugen-rekenridge ligt.

2. GQA Pad (Rekenkracht-vriendelijk):

   • Mecanisme: De latent wordt uitgebreid naar $g$ sleutel/waarde-groepen. Decoding voert vanilla Grouped-Query Attention (GQA) uit zonder per-stap latent uitbreiding.
   • Cache: Bewaart per-groep uitgebreide sleutels en waarden ($2gd_h$) plus de gedeelde RoPE-sleutel.
   • Doelhardware: Geoptimaliseerd voor H20-klasse hardware. Door het aanpassen van het groepenaantal ($g$) en het aantal query-tokens per stap ($s_q$), kan de rekenintensiteit worden vastgezet op de H20-ridge ($\approx$ 37 FLOPs/byte).
   • Schaalbaarheid: Ondersteunt tot 8-voudig zero-redundancy tensor parallelisme langs de groeps-as.

Pad Selectie

De runtime selecteert het pad op basis van de doelhardware zonder hertraining of aangepaste kernels:

• H100: Gebruikt het MQA-absorb pad ($s_q=1$).
• H20: Gebruikt het GQA-pad (bijv. $g=8, s_q=2$ of $g=4, s_q=1$).
• MTP: Op de H20 herstelt het GQA-pad met $s_q=2$ bijna-ideale doorvoerwinsten, terwijl MLA niets wint.

TransGQLA: Conversie van Voortrainde Modellen

Om te voorkomen dat er vanaf nul wordt voortraind, introduceren de auteurs TransGQLA, een extensie van de TransMLA-pijplijn. Het converteert een voortraind GQA-checkpoint naar een GQLA-model via een enkele gerichte modificatie:

• Kop Samenvoegen: In plaats van gegroepeerde koppen te samenvoegen tot een enkele latent en up-projecties te repliceren over alle query-koppen (zoals in TransMLA), houdt TransGQLA de up-projecties geïndexeerd per groep.
• Resultaat: Het model behoudt het GQA-uitvoeringspad structureel terwijl het het MQA-absorb-pad activeert via de standaard absorptie-operatie. Dit behoudt tensor parallelisme langs de groeps-as.

Sparse GQLA

Het ontwerp breidt zich uit naar fijnkorrelige sparse attention. Omdat de canonieke GQLA-configuratie ($h_q/g = 16$) overeenkomt met de Tensor-Core MMA-taakdimensie ($m=16$), kan sparse GQLA het GQA-pad behouden op H20-klasse hardware. Daarentegen is sparse MLA structureel vergrendeld op het geabsorbeerde MQA-pad vanwege taakbeperkingen, en erft dezelfde rekenkracht-overheadproblemen.

3. Belangrijkste Bijdragen

1. Identificatie van MLA Nadelen: Het artikel identificeert dat het single-path ontwerp van MLA het koppelt aan H100-klasse verhoudingen, tensor parallelisme langs de kop-as opgeeft, en nul MTP-winsten oplevert op commodity GPU's.
2. GQLA Architectuur: Een minimale variant van MLA die twee algebraïsch equivalente paden (MQA-absorb en GQA) blootlegt over dezelfde gewichten. Een enkele configuratie ($h_q=128, g=8$) met één MTP-kop spijst de daklijnen van zowel H100 als H20 simultaan.
3. TransGQLA Pijplijn: Een conversiemethode die voortrainde GQA-checkpoints transformeert naar GQLA-modellen met minimale architecturale wijzigingen, tensor parallelisme behoudt en een soepele overgang naar de latent-vorm mogelijk maakt.
4. Roofline Analyse: Een theoretische verificatie die aantoont dat GQLA-gewichten bijna piekdoorvoer kunnen bereiken op zowel H100 als H20 door paden te wisselen, terwijl MLA inefficiënt is op de H20.

4. Resultaten

De auteurs evalueerden TransGQLA op het LLaMA-3-8B model:

• Behoud van Capabiliteit: Het converteren van de GQA-gewichten naar de GQLA latent-vorm (zonder verdere training) resulteerde in een verlies van slechts $\sim$9,7 gemiddelde punten op benchmarks voor gezond verstand redeneren (MMLU, ARC, enz.) vergeleken met het originele 15T-token voortrainde model.
• Cache Compressie: Op het MQA-absorb-pad werd de per-token KV-cache gecomprimeerd tot 28,125% van de originele GQA-baseline (een reductie van 71,875%).
• Doorvoer Projecties:
  • H100: Het MQA-absorb-pad bereikt 354K tokens/sec.
  • H20: Het GQA-pad (met $g=8, s_q=2$) bereikt 221K tokens/sec, een 3,4x verbetering ten opzichte van MLA op hetzelfde apparaat (dat vastzit op 65K tokens/sec vanwege rekenkracht-beperking).
• Voortgezette Training: De auteurs projecteren dat voortgezette training (volgend op de TransMLA-trajectorie) het verloren vermogen zal herstellen tot binnen 0,5 punten van het originele model na 30B tokens, terwijl de KV-cache compressie behouden blijft.

5. Betekenis en Claims

Het artikel claimt dat GQLA een praktisch, hardware-onafhankelijk ontwerpprincipe biedt voor LLM-inferentie. Door meerdere algebraïsch equivalente decoding-paden bloot te leggen over één set getrainde gewichten, stelt GQLA een model in staat om zich aan te passen aan verschillende hardware-beperkingen (rekenkracht-rijk versus geheugen-rijk) op het moment van implementatie zonder hertraining of aangepaste kernels.

De auteurs positioneren GQLA als een oplossing voor het "hardware koppeling" probleem, waardoor efficiënte implementatie mogelijk is op zowel vlaggeschip-accelerators (H100) als commodity/export-beperkte hardware (H20). Zij benadrukken dat deze aanpak Multi-Token Prediction en Tensor Parallelisme ondersteunt op hardware waar MLA dit niet efficiënt kan doen.

Erkende Beperkingen:

• Hardware claims zijn gebaseerd op Roofline-modellering; volledige kernel-niveau benchmarks op daadwerkelijke H20/H100 hardware staan nog uit.
• Resultaten van voortgezette training voor TransGQLA zijn momenteel in uitvoering en nog niet afgerond in het artikel.
• Evaluatie is momenteel beperkt tot de LLaMA-3-8B backbone en benchmarks voor gezond verstand redeneren.
• Sparse GQLA is structureel beschreven maar nog niet gebenchmarkt tegen bestaande sparse attention-methoden.