Pooling Engram Conditional Memory in Large Language Models using CXL

Dit paper introduceert een schaalbare en kostenefficiënte oplossing voor Engram-gebaseerde conditionele geheugen in grote taalmodellen door het gebruik van CXL-geheugenpools, wat naadloze integratie met SGLang mogelijk maakt en prestaties dicht bij die van DRAM behoudt.

De kern

Stel je voor dat een grote taalmodel (zoals een slimme chatbot) een enorme bibliotheek is. Om antwoorden te geven, moet deze bibliotheek twee dingen doen:

1. Rekenen: De eigenlijke "denken" en zinnen construeren.
2. Zoeken: Snel feiten opzoeken uit een gigantische lijst van woorden en zinsdelen (de "Engram").

Het probleem is dat deze zoeklijst (de Engram) zo groot wordt dat hij niet meer in het snelle, dure geheugen van de computer past. Als je hem in het langzamere, goedkopere geheugen stopt, wordt de chatbot traag en hapert hij.

Deze paper lost dat op met een slimme truc: CXL.

Hier is de uitleg in simpele taal, met wat creatieve vergelijkingen:

1. Het Probleem: De "Grote Boekentas"

Stel je voor dat de chatbot een student is die een examen doet.

• De snelle hersenen (DRAM): Dit is het werkgeheugen van de student. Het is supersnel, maar er past maar weinig op.
• De Engram: Dit is een enorme tas met duizenden boeken die de student nodig heeft om feiten te checken. Deze tas is te groot voor op het bureau (de snelle RAM).
• Het oude probleem: Als je de boeken in een kelder (langzamer geheugen) zet, moet de student elke keer hardlopen naar de kelder, een boek pakken en terugkomen. Dat kost te veel tijd. De chatbot wordt traag.

2. De Oplossing: De "CXL Super-Hoofdtelefoon"

De auteurs van dit paper gebruiken een nieuwe technologie genaamd CXL (Compute Express Link).

Stel je voor dat CXL een magische, supersnelle telefoonlijn is die de student direct verbindt met de kelder, zonder dat hij hoeft te lopen.

• Geen gedoe: Bij de oude methode (RDMA) moest de student eerst een formulier invullen, wachten op een bevestiging en dan pas het boek ophalen. Dat is als een postbode die langskomt.
• Met CXL: De student pakt het boek direct uit de kelder alsof het op zijn bureau ligt. Het voelt alsof de boeken daar zijn, terwijl ze eigenlijk ergens anders staan.

3. Waarom werkt dit zo goed voor Engram?

De "Engram" (de zoeklijst) heeft een heel specifiek gedrag:

• Het is statisch: De boeken veranderen nooit tijdens het examen. Je hoeft ze niet te herschrijven, alleen te lezen.
• Het is versnipperd: De student haalt niet één dik boek, maar duizenden kleine kaartjes (woorden) op.
• Het is snel nodig: De student moet de kaartjes halen terwijl hij al aan het rekenen is aan de volgende vraag.

De oude methoden (zoals RDMA) zijn goed voor het verplaatsen van hele vrachtwagens vol data (grote pakketten), maar slecht voor het snel oppikken van duizenden losse kaartjes. CXL is daarentegen perfect voor het snel oppikken van losse kaartjes. Het is alsof je een robotarm hebt die direct naar het juiste plankje grijpt, in plaats van een vrachtwagen die moet parkeren.

4. Het Resultaat: Snel én Goedkoop

In hun experiment hebben ze laten zien dat:

• Snelheid: De chatbot is net zo snel als wanneer de boeken op het bureau lagen (dichtbij de snelheid van het dure geheugen).
• Kosten: Omdat je de boeken nu in één gedeelde "kelder" kunt zetten die door meerdere studenten (servers) wordt gebruikt, hoef je niet voor elke student een eigen dure kast te kopen.

De analogie van de kosten:

• Oude manier: Elke student koopt zijn eigen dure, snelle boekenkast. Als je 10 studenten hebt, heb je 10 dure kasten nodig.
• Nieuwe manier (CXL): Je bouwt één enorme, centrale bibliotheek. Alle studenten gebruiken die ene bibliotheek via de magische telefoonlijn. Voor een grote groep is dit veel goedkoper, zonder dat ze hoeven te wachten.

Conclusie

Deze paper toont aan dat we met CXL de "Engram" (de kennis van de AI) kunnen verplaatsen naar goedkoper, gedeeld geheugen, zonder dat de AI traag wordt. Het is alsof we de AI een super-snel geheugen geven dat oneindig groot kan worden, maar niet duur. Dit maakt de toekomst van slimme AI-systemen veel haalbaarder en goedkoper.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Pooling Engram Conditional Memory in Large Language Models using CXL", geschreven in het Nederlands.

Probleemstelling

Grote Taalmodellen (LLM's) vertrouwen vaak op Mixture-of-Experts (MoE) voor conditionele berekening, maar deze architectuur mist een efficiënt mechanisme voor het opzoeken van statische vocabulairekennis. Dit dwingt modellen om kennisopslag te simuleren via berekeningen, wat inefficiënt is.

Om dit op te lossen introduceert Engram "Conditionele Memory". Dit decoupeert statische kennisopslag (N-Gram embeddings) van dynamische berekening. Hoewel Engram de prestaties verbetert door de N-Gram-tabel exponentieel te laten groeien, introduceert het een aanzienlijk geheugenvolume (potentieel honderden GB's).

• De uitdaging: Het volledig in het dure GPU-DRAM houden van deze enorme Engram-tabellen is kostbaar en schaalt slecht.
• De specificiteit: Engram-toegang is echter spaarsam, lees-only en zeer fijnkorrelig. Per token worden slechts kleine, discrete fragmenten (bijv. 5 KB) opgehaald. Traditionele oplossingen zoals RDMA (Remote Direct Memory Access) zijn hier niet ideaal voor, omdat ze inefficiënt zijn bij het verwerken van kleine, discrete pakketten (hoge overhead, lage doorvoersnelheid voor kleine berichten) en vaak een CPU-gestuurde "bounce buffer" vereisen.

Methodologie

Het paper stelt een systeem voor dat Engram-parameters uitlaadt naar een gedeeld CXL (Compute Express Link) geheugenpool.

1. Architectuur:

   • Het systeem maakt gebruik van een CXL-switch (bijv. XConn XC50256) die meerdere rekennodes (servers) verbindt met een centraal CXL-geheugen.
   • In tegenstelling tot RDMA, biedt CXL hardware-ondersteuning voor native load/store-primitieven op cache-line-niveau. Dit elimineert de overhead van netwerkprotocollen en stelt de CPU/GPU in staat om direct naar het externe geheugen te adresseren alsof het lokaal is.

2. Implementatie in SGLang:

   • De Engram-module is geïntegreerd in de inferentie-engine SGLang.
   • Initialisatie: Engram-parameters worden één keer geladen in het gedeelde CXL-geheugen.
   • Prefetching: Bij het binnenkomen van een batch tokens worden de benodigde Engram-embeddings asynchroon opgehaald uit de CXL-pool.
   • Toegangsroutines:
     • CXL → CPU: Gebruikt multi-threaded memcpy via OpenMP, waarbij het CXL-geheugen is gemap als DAX (Direct Access) apparaat.
     • CXL → GPU: Een aangepaste CUDA-kernel zorgt voor directe P2P (Peer-to-Peer) data-overdracht. De GPU's gebruiken DMA om direct naar de CXL-adressen te schrijven, wat de CPU volledig omzeilt na de initiële mapping. Dit maximaliseert de parallelle verwerking van duizenden kleine verzoeken.

3. Analyse van Toegangsvereisten:

   • De auteurs analyseren dat de bandbreedte-eis voor Engram relatief laag is (ongeveer 0,7 GB/s voor een hoge doorvoer), maar de latentie-eis zeer streng is (minder dan 56 µs voor vroege lagen) om de inferentie niet te blokkeren. CXL voldoet hieraan, terwijl RDMA faalt vanwege de hoge latentie bij kleine pakketten.

Belangrijkste Bijdragen

• Eerste Systeem: Dit is het eerste systeem dat Engram-parameters uitlaadt naar een CXL-gebaseerde geheugenpool.
• Vergelijkende Analyse: Het paper demonstreert dat CXL superieur is aan RDMA voor Engram-toegangspatronen vanwege de fijnkorrelige, lage-latentie eisen.
• Implementatie en Validatie: Een werkende implementatie in SGLang die bewijst dat CXL-pooling haalbaar is zonder in te leveren op inferentieprestaties.

Resultaten

De experimenten zijn uitgevoerd op een testbed met Intel Xeon CPU's en NVIDIA L20 GPU's, verbonden via een CXL-switch naar een 256 GB CXL-geheugenmodule.

• Latentie:
  • CXL-to-CPU-latentie is vergelijkbaar met lokaal DRAM.
  • CXL-to-GPU-latentie is iets hoger dan lokaal DRAM, maar binnen het aanvaardbare bereik voor prefetching.
  • In vergelijking met RDMA (Mooncake Store) is de latentie van CXL aanzienlijk lager, vooral bij kleine batchgroottes, wat essentieel is voor Engram.
• Doorvoer (Throughput):
  • Bij het testen met Qwen3-4B en Qwen3-8B (geëmuldeerd met Engram) toonde het CXL-systeem een doorvoer aan die bijna gelijk was aan die van een lokaal DRAM-systeem.
  • Voorbeeld Qwen3-4B: Baseline (6183 tokens/s) vs. Engram met CXL (5614 tokens/s). Het verlies is marginaal en vergelijkbaar met het gebruik van lokaal DRAM.
• Schaalbaarheid:
  • Het systeem schaalt goed. Het verhogen van Data Parallelism (DP) en het aantal knooppunten leidt tot slechts een verwaarloosbare prestatiedaling, wat aantoont dat de CXL-pool concurrente toegang van meerdere hosts aankan.
• Kostenanalyse:
  • Voor kleine configuraties (bijv. 2 knooppunten) is CXL duurder door de vaste kosten van de switch en adapters.
  • Voor grotere schalen (bijv. 8 of 16 knooppunten) is CXL aanzienlijk goedkoper. Bij een 16-knooppunten cluster met een 400B Engram-tabel bespaart het systeem $166.040 ten opzichte van lokale DRAM-oplossing, omdat de dure geheugenkapaciteit gedeeld wordt in plaats van gedupliceerd.

Betekenis en Conclusie

Dit paper toont aan dat CXL een ideale infrastructuur is voor het opslaan van Engram-conditional memory in toekomstige LLM's. Door de unieke toegangspatronen van Engram (spaarsam, lees-only, fijnkorrelig) te koppelen aan de lage-latentie, fijnkorrelige toegang van CXL, kunnen ontwikkelaars:

1. De hardwarekosten voor grote taalkennis-tabellen drastisch verlagen door geheugen te delen.
2. De inferentieprestaties behouden die nodig zijn voor real-time toepassingen.
3. De "GPU-memory wall" omzeilen zonder de snelheid van het model te vertragen.

De studie legt de basis voor een schaalbare, kostenefficiënte architectuur voor de volgende generatie geheugen-verrijkte Large Language Models.