PagedAttention: hoe AI-modellen efficiënter met hun geheugen leren omgaan
Wie weleens met ChatGPT of een vergelijkbare chatbot heeft gepraat, weet dat de antwoorden woord voor woord, of eigenlijk stukje voor stukje, verschijnen. Wat je niet ziet, is dat de computer achter de schermen voor elk gesprek een enorme hoeveelheid tussentijdse rekenresultaten moet onthouden zolang het gesprek duurt. Die tussentijdse resultaten heten de KV-cache, en het efficiënt beheren daarvan is een van de grootste technische knelpunten bij het op grote schaal aanbieden van taalmodellen. PagedAttention is een techniek die dat geheugenbeheer drastisch efficiënter maakt.
De kern van het idee is verrassend ouderwets: het is geleend van besturingssystemen zoals Windows of Linux, die al decennia lang het werkgeheugen van een computer opdelen in kleine, vaste "pagina's" in plaats van alles aaneengesloten op te slaan. Vergelijk het met een parkeergarage: in plaats van te eisen dat de auto's van één gezelschap altijd naast elkaar staan (waardoor je bij een vol parkeervak soms geen plek meer vindt, ook al staan er verderop losse lege plekken), laat je auto's gewoon verspreid parkeren en houd je een lijstje bij van wie waar staat. PagedAttention past dat principe toe op het geheugen van AI-taalmodellen, met als resultaat dat servers veel meer gebruikers tegelijk kunnen bedienen op dezelfde hardware.
Wat is het precies?
Moderne taalmodellen zoals GPT, Llama of Mistral zijn gebouwd op een architectuur die "transformer" heet. Het kernmechanisme daarin heet attention (aandacht): voor elk nieuw woord dat het model genereert, vergelijkt het dat woord met alle voorgaande woorden in het gesprek om te bepalen welke daarvan relevant zijn. Om dat niet steeds opnieuw te hoeven berekenen, slaat het model voor elk eerder woord twee soorten getallenreeksen op: een "key" en een "value". Samen vormen die de KV-cache.
Het probleem is dat die cache groot en onvoorspelbaar van omvang is. Hoe langer een gesprek of document, hoe meer geheugen de cache inneemt, en vooraf weet een server niet hoe lang een gebruiker gaat doorpraten. Traditionele systemen reserveerden daarom voor elk gesprek een aaneengesloten geheugenblok, groot genoeg voor de langst mogelijke tekst. In de praktijk bleef daardoor veel gereserveerde ruimte ongebruikt, en onderzoekers van vLLM (zie hieronder) rapporteerden dat op deze manier vaak meer dan de helft van het GPU-geheugen verspild werd aan lege, gereserveerde ruimte. Dat probleem heet geheugenfragmentatie: het geheugen raakt versnipperd in kleine, onbruikbare stukjes en gereserveerde-maar-lege blokken.
PagedAttention lost dit op door de KV-cache, net als het werkgeheugen van een besturingssysteem, op te knippen in kleine, vaste blokken (in de praktijk vaak zestien tokens per blok; een token is ongeveer een woorddeel). Die blokken hoeven niet aaneengesloten in het geheugen te staan. Een apart "blokkentabel" per gesprek houdt bij welke fysieke blokken bij welk gesprek horen en in welke volgorde ze gelezen moeten worden, precies zoals een besturingssysteem een paginatabel bijhoudt om virtueel geheugen naar fysiek geheugen te vertalen. Zo kan een server geheugen vrijwel volledig benutten in plaats van ruimte te reserveren die misschien nooit gebruikt wordt.
Een bijkomend voordeel is geheugendeling. Wanneer een server bijvoorbeeld meerdere antwoordvarianten tegelijk genereert op basis van dezelfde vraag (parallel sampling), of wanneer meerdere gebruikers met hetzelfde stuk voorafgaande tekst beginnen, hoeven de bijbehorende blokken niet gedupliceerd te worden. Ze kunnen simpelweg door meerdere gesprekken tegelijk "gelezen" worden, en pas als een van die gesprekken die tekst wil aanpassen, wordt er een kopie gemaakt (een techniek die "copy-on-write" heet, eveneens geleend uit de wereld van besturingssystemen).
Wat wil men ermee bereiken?
Het uiteindelijke doel is simpel: meer gebruikers tegelijk bedienen met dezelfde hoeveelheid dure GPU-hardware, tegen lagere kosten per gegenereerd woord. GPU-geheugen is schaars en duur, en elke procent die verloren gaat aan verspilling is direct geld en capaciteit die niet aan echte gebruikers ten goede komt.
Concreet streeft men naar een hogere doorvoer (throughput): het aantal tokens dat een server per seconde kan genereren over alle gebruikers heen. Door minder geheugen te verspillen, past een server meer gelijktijdige gesprekken in het beschikbare geheugen, wat de batchgrootte vergroot en de hardware efficiënter benut. Dat heeft weer gevolgen voor de gebruiker: snellere antwoorden bij drukte, lagere prijzen per opgevraagd antwoord, en de mogelijkheid om langere gesprekken of documenten te verwerken zonder dat de server meteen vastloopt op geheugentekort.
Voorbeelden uit de praktijk
PagedAttention werd geïntroduceerd door onderzoekers van de University of California, Berkeley, in het paper "Efficient Memory Management for Large Language Model Serving with PagedAttention" (2023), en gepresenteerd op SOSP 2023, een gerenommeerde conferentie over besturingssystemen. De techniek vormt de basis van vLLM, een open-source softwarepakket voor het serveren van taalmodellen dat sindsdien is uitgegroeid tot een van de meest gebruikte inference-servers voor open taalmodellen zoals Llama, Mistral en Qwen.
Een bekend voorbeeld van toepassing in de praktijk is LMSYS Chatbot Arena, het platform waarop internetgebruikers anoniem verschillende taalmodellen met elkaar kunnen vergelijken; dit initiatief, ontstaan aan Berkeley, gebruikte vLLM om de vele gelijktijdige testgesprekken efficiënt te kunnen bedienen.
Ook buiten het vLLM-project heeft het idee van blokgewijs, niet-aaneengesloten geheugenbeheer voor de KV-cache navolging gekregen. NVIDIA heeft in zijn eigen inference-framework TensorRT-LLM vergelijkbare paged KV-cache-technieken doorgevoerd, en Hugging Face heeft in zijn Text Generation Inference (TGI) ook op blokken gebaseerd geheugenbeheer geïntegreerd. Dit soort kruisbestuiving is typisch voor de sector: een idee dat in één open-sourceproject zijn nut bewijst, wordt binnen afzienbare tijd door concurrerende projecten overgenomen of er wordt een eigen variant op gebouwd.
Hoe ver is de techniek?
PagedAttention is geen experimenteel prototype meer, maar een breed toegepaste, volwassen standaardtechniek in de wereld van LLM-serving. Het vLLM-project wordt actief onderhouden, met regelmatige releases en bijdragen van tientallen ontwikkelaars en bedrijven wereldwijd. In een paar jaar tijd is de techniek van academisch onderzoek naar praktisch gangbare infrastructuur gegaan, wat in het snelle AI-veld niet vanzelfsprekend is.
Rond PagedAttention zijn inmiddels aanvullende optimalisaties ontstaan die met deze onderliggende geheugenstructuur samenwerken. Denk aan continuous batching (het dynamisch samenvoegen van nieuwe en lopende verzoeken in één rekenbatch, in plaats van te wachten tot een hele batch klaar is), chunked prefill (het opsplitsen van de verwerking van lange invoerprompts in kleinere stukken om andere verzoeken niet te laten wachten) en prefix caching (het hergebruiken van reeds berekende blokken wanneer meerdere verzoeken met dezelfde tekst beginnen).
Toch zijn er nog open uitdagingen. Het plannen (scheduling) van welke blokken en verzoeken wanneer verwerkt worden, wordt complexer naarmate workloads diverser worden, bijvoorbeeld bij een mix van hele korte en hele lange gesprekken, of bij multimodale modellen die naast tekst ook beeld of geluid verwerken en dus andere soorten cachegegevens beheren. Ook de keuze van de blokgrootte blijft een afweging: kleinere blokken verminderen verspilling verder maar verhogen het beheeroverhead, terwijl grotere blokken eenvoudiger te beheren zijn maar weer meer ruimte kunnen verspillen. Er is dus geen definitief "opgelost"; het blijft een actief ontwikkelgebied.
Wie werken eraan?
De oorsprong van PagedAttention ligt bij het Sky Computing Lab van UC Berkeley (de opvolger van het bekende RISELab), waar het oorspronkelijke onderzoeksteam rond Woosuk Kwon en collega's de techniek ontwikkelde en als open source vrijgaf. Sindsdien is vLLM uitgegroeid tot een gemeenschapsproject met bijdragen van een brede mix aan partijen: individuele onderzoekers en ontwikkelaars, universiteiten, en commerciële bedrijven die de software in hun eigen producten gebruiken of er financieel en met mankracht aan bijdragen.
Daarnaast werkt de bredere industrie aan vergelijkbare of concurrerende technieken. NVIDIA (met TensorRT-LLM) en Hugging Face (met TGI) zijn de bekendste voorbeelden van bedrijven die eigen implementaties van blokgewijs geheugenbeheer voor taalmodellen hebben gebouwd. Ook grote clouddienstverleners die AI-modellen als dienst aanbieden, hebben belang bij dit soort efficiëntieverbeteringen en dragen bij aan het onderzoek, al is niet altijd publiek gedetailleerd wie precies welke interne aanpassingen gebruikt.