PIM-SHERPA: Software Method for On-device LLM Inference by Resolving PIM Memory Attribute and Layout Inconsistencies

Dit paper introduceert PIM-SHERPA, een softwaremethode die de inconsistenties in geheugenattributen en lay-out oplost om efficiënte on-device inferentie van grote taalmodellen op PIM-systemen mogelijk te maken.

De kern

Stel je voor dat je een zeer slimme, digitale assistent (een "Grote Taalmodel" of LLM) op je telefoon wilt laten draaien. Deze assistent kan praten, vertalen en vragen beantwoorden. Maar er is een probleem: deze assistent is enorm en vraagt veel van het geheugen van je telefoon.

Om dit snel te laten werken, hebben onderzoekers een slimme techniek bedacht genaamd PIM (Processing-In-Memory). In plaats van dat de processor (de hersenen) alle gegevens moet ophalen bij het geheugen (de bibliotheek), doet de assistent zijn rekenwerk direct in de bibliotheek zelf. Dit is veel sneller.

Maar hier zit een addertje onder het gras, en dat is precies wat dit papier, PIM-SHERPA, oplost.

Het Probleem: Twee verschillende werelden

Het probleem is dat de assistent twee verschillende manieren van werken heeft, en deze botsen met elkaar:

1. De "Voorbereiding" (Prefill): Als je een lange tekst invoert, moet de assistent eerst alles lezen en begrijpen. Hierbij is het slim om de gegevens in een cache te stoppen. Denk aan een cache als een werktafel vlak bij de assistent. Als de assistent iets vaak nodig heeft, pakt hij het snel van de tafel. Dit is supersnel.
2. De "Antwoord" (Decode): Als de assistent één woord per keer gaat schrijven, moet hij direct naar de bibliotheek (het geheugen) gaan om de regels te raadplegen. Hier werkt de PIM-techniek het beste. Maar voor PIM te werken, mag de assistent geen gebruik maken van de werktafel (cache). Hij moet direct naar de bibliotheek rennen. Als de gegevens op de werktafel liggen, denkt de assistent: "Ah, daar heb ik het al!" en hij gaat niet naar de bibliotheek. Dan werkt de snelle PIM-techniek niet.

De botsing:

• Voor de voorbereiding wil je de gegevens op de werktafel (cache).
• Voor het antwoord wil je dat ze niet op de werktafel liggen, maar direct in de bibliotheek (non-cacheable).

Als je dit oplost door twee volledige kopieën van de assistent te maken (één voor de werktafel, één voor de bibliotheek), heb je dubbel zoveel ruimte nodig. Je telefoon is dan vol, en je kunt geen grote modellen meer draaien.

De Oplossing: PIM-SHERPA

De onderzoekers hebben PIM-SHERPA bedacht. Dit is een slimme software-methode die geen extra hardware nodig heeft. Ze gebruiken twee slimme trucs:

Truc 1: De "Dubbele Buffer" (DDB) – Het Vervoerbandje

Stel je voor dat je een fabriek hebt.

• De ene machine (de CPU) werkt aan de huidige taak.
• Tegelijkertijd draait er een vervoerbandje (de copy-thread) die de volgende set instructies alvast van de bibliotheek haalt en op een klein, tijdelijk tafeltje (de cache-buffer) legt.
• Zodra de eerste machine klaar is, pakt hij de nieuwe instructies van dat kleine tafeltje, terwijl het vervoerbandje alweer aan het werk is voor de daarna volgende stap.

Dit betekent dat het "verplaatsen" van de gegevens (het aanpassen van de lay-out) gebeurt terwijl de assistent al aan het rekenen is. Je merkt er niets van, want het wachten valt samen met het werken.

Truc 2: "Online Herverrang" (OWR) – De Slimme Koerier

Soms is de tekst die je invoert zo lang, dat het rekenen zelf de langste tijd kost. In dat geval maakt het niet uit als je even stopt om de gegevens te verplaatsen.

• De software wacht tot het moment dat de assistent klaar is met een stap.
• Dan stuurt hij direct een koerier (de copy-thread) om de benodigde gegevens uit de bibliotheek te halen, ze om te vormen tot een formaat dat de assistent begrijpt, en ze op de werktafel te leggen.
• Omdat het rekenen zo lang duurt, is de koerier al klaar voordat de assistent weer iets nodig heeft.

Waarom is dit geweldig?

1. Ruimtebesparing: In plaats van twee volledige kopieën van de assistent te hebben (wat je telefoon vol zou maken), gebruiken ze maar één kopie en een heel klein extra tafeltje. Dit bespaart bijna 50% aan geheugenruimte. Je kunt dus veel slimmere modellen op je telefoon draaien.
2. Snelheid: Het is net zo snel als de theoretische maximumsnelheid. De assistent hoeft niet te wachten op de gegevens.
3. Geen nieuwe hardware: Je hoeft je telefoon niet te vervangen. Het werkt gewoon met de software die er nu al is.

Samenvattend

PIM-SHERPA is als een slimme manager die ervoor zorgt dat zijn assistent altijd de juiste gereedschappen op de juiste plek heeft, zonder dat hij twee keer zoveel gereedschapskisten nodig heeft. Hij haalt de volgende set gereedschappen alvast op terwijl de assistent aan het werk is, zodat er nooit stilstand ontstaat. Hierdoor kunnen we straks veel krachtige AI-assistenten hebben die direct op onze telefoons werken, zonder dat we een dure nieuwe telefoon hoeven te kopen.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "PIM-SHERPA: SOFTWARE METHOD FOR ON-DEVICE LLM INFERENCE BY RESOLVING PIM MEMORY ATTRIBUTE AND LAYOUT INCONSISTENCIES" in het Nederlands.

1. Het Probleem

De implementatie van Large Language Models (LLM's) op mobiele en edge-apparaten groeit snel. LLM-inferentie bestaat uit twee fasen met fundamenteel verschillende eisen:

• Prefill-fase: Berekeningsintensief (GEMM), waarbij gewichten vaak worden hergebruikt. Dit profiteert van een cacheerbaar geheugenregio en een "host-vriendelijke" lay-out (bijv. kolom-gewijze volgorde voor efficiënte cache-hits).
• Decode-fase: Bandbreedte-intensief (GEMV), waarbij één token per stap wordt verwerkt. Dit is ideaal voor Processing-in-Memory (PIM), zoals LPDDR-PIM, omdat PIM de DRAM-bandbreedte tot 8x kan verhogen. Echter, om PIM te activeren, moeten DRAM-aanvragen de geheugencontroller bereiken. Als gewichten in een cacheerbaar gebied staan, kunnen cache-hits deze aanvragen blokkeren, waardoor PIM niet wordt geactiveerd. Daarom vereist de decode-fase een niet-cacheerbaar gebied en een "PIM-bewuste" lay-out (waarbij volledige rijen van een matrix binnen één DRAM-bank liggen voor SIMD-efficiëntie).

De Kernuitdaging: Er bestaat een inconsistentie tussen deze fasen:

1. Geheugenattribuut-inconsistentie: Prefill vereist cacheerbaar geheugen, decode vereist niet-cacheerbaar geheugen.
2. Gewichtslay-out-inconsistentie: Prefill prefereert een host-vriendelijke lay-out, decode vereist een PIM-bewuste lay-out.

Bestaande oplossingen zoals weight duplication (het dupliceren van gewichten voor beide fasen) verdubbelen het geheugengebruik, wat onhaalbaar is voor mobiele apparaten met beperkt DRAM (bijv. 8-12 GB). Hardware-oplossingen zoals FACIL vereisen wijzigingen aan de geheugencontroller, wat de toepasbaarheid beperkt.

2. Methodologie: PIM-SHERPA

PIM-SHERPA is een software-only methode die deze inconsistenties oplost zonder hardware-aanpassingen. Het introduceert twee technieken die kleine cacheerbare buffers gebruiken om gewichten dynamisch te verplaatsen en te herschikken:

A. DRAM Double Buffering (DDB)

• Concept: Deze methode overlapt data-overdracht met berekening. Er worden twee kleine cacheerbare buffers gereserveerd in het DRAM.
• Werking: Terwijl de huidige decoderlaag GEMM-berekeningen uitvoert op gewichten in Buffer 0, worden de gewichten voor de volgende laag (in de PIM-bewuste lay-out) uit het niet-cacheerbare gebied gehaald, herschikt (swizzled) en opgeslagen in Buffer 1.
• Voordeel: De latency van het herschikken van gewichten wordt volledig verdoezeld door de berekeningstijd van de huidige laag. Dit vereist geen offline duplicatie van het volledige model.

B. Online Weight Rearrangement met Swizzled Memory Copy (OWR)

• Concept: Deze methode is geoptimaliseerd voor scenario's met lange invoersequenties (zoals interactieve LLM's of RAG), waarbij de berekeningstijd (GEMM) domineert.
• Werking: Onmiddellijk voor de berekening van een laag worden de benodigde gewichten uit het niet-cacheerbare gebied gekopieerd naar een enkele cacheerbare buffer. Tijdens deze kopie worden de data direct omgezet van de PIM-bewuste lay-out naar een host-vriendelijke lay-out ("Swizzled Memory Copy").
• Voordeel: Vereist geen complexe synchronisatie of load-balancing tussen kopieer- en berekeningsthreads, wat de implementatiecomplexiteit verlaagt. Omdat de kopieerlatentie statisch is en de GEMM-tijd groeit met de sequentielengte, wordt de relatieve overhead verwaarloosbaar bij lange sequenties.

Swizzled Memory Copy (SMC):
Een kerntechniek die de data transformeert tijdens het kopiëren. De host leest gewichten uit het niet-cacheerbare gebied (PIM-bewust) en schrijft ze direct in de juiste volgorde naar het cacheerbare gebied (host-vriendelijk), zodat de standaard GEMM-kernels zonder wijzigingen kunnen worden gebruikt.

3. Belangrijkste Bijdragen

• Identificatie van het probleem: Dit is het eerste werk dat de specifieke inconsistentie in geheugenattributen en gewichtslay-outs tussen prefill en decode fasen in PIM-systemen expliciet identificeert en kwantificeert.
• Software-oplossingen: Voorstel van twee software-only methoden (DDB en OWR) die kleine buffers gebruiken in plaats van volledige duplicatie, waardoor het geheugengebruik drastisch wordt verminderd.
• Geen hardware-aanpassingen: De oplossing werkt op bestaande producten zonder wijzigingen aan de geheugencontroller of het DRAM-chip.
• Kwantificering: Analyse van hoe naive duplicatie het DRAM-gebruik onhaalbaar maakt voor mobiele apparaten en hoe PIM-SHERPA dit oplost.

4. Resultaten

De evaluatie werd uitgevoerd op een Samsung Galaxy S24+ met Llama 3.2-modellen (1B en 3B parameters) en een emulatie van LPDDR5X-PIM.

• Geheugencapaciteit: PIM-SHERPA bespaart ongeveer 47,8% tot 49,7% DRAM-capaciteit in vergelijking met de methode van gewichtsduplicatie (WD). Dit maakt het mogelijk om grotere modellen (zoals Llama 3.2 3B) op smartphones te draaien die PIM ondersteunen, wat anders onmogelijk zou zijn door geheugenbeperkingen.
• Prestaties (TTFT):
  • De Time-To-First-Token (TTFT) van PIM-SHERPA (zowel DDB als OWR) is vergelijkbaar met of zelfs iets lager dan die van FACIL-O (een hypothetische versie van FACIL die het geheugenprobleem zou oplossen).
  • Bij invoersequenties van 128 tokens of langer, verbergt DDB de kopieerlatentie volledig.
  • OWR presteert iets slechter bij korte sequenties door de seriële kopieeroverhead, maar benadert de prestaties van DDB en FACIL-O bij lange sequenties (bijv. 192 tokens).
• Snelheidswinst: PIM-SHERPA behoudt de snelheidswinst van PIM (tot 3,3x sneller dan een niet-PIM baseline) en levert eind-tot-eind versnellingen die vergelijkbaar zijn met de theoretische maximums.

5. Betekenis en Impact

• Commercialisatie van PIM: Dit werk is cruciaal voor de praktische toepassing van PIM in consumentenelektronica. Het lost de fundamentele blokkade op die eerder leidde tot onhaalbaar geheugengebruik of hardware-afhankelijke oplossingen.
• Toekomstbestendigheid: De methode is ontworpen voor de trend naar langere invoersequenties (door interactieve LLM's en RAG), waarbij de overhead van online herschikking minder impact heeft.
• Universele toepasbaarheid: Hoewel getest op CPU's, is de aanpak ook toepasbaar op GPU's en NPUs die een laatste niveau cache of SRAM hebben, wat de relevantie voor diverse hardware-platforms vergroot.

Kortom, PIM-SHERPA bewijst dat efficiënte LLM-inferentie op PIM-apparaten mogelijk is door slimme software-strategieën die de conflicterende eisen van prefill en decode fasen oplossen zonder de beperkte geheugenbronnen van mobiele apparaten te overbelasten.