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

Die Arbeit stellt PIM-SHERPA vor, eine rein softwarebasierte Methode, die durch die Einführung von DRAM-Doppelbuffering und online Gewichtsumordnung die Speicherattribut- und Layout-Inkonsistenzen bei der Inferenz von Large Language Models auf PIM-fähigen Endgeräten löst und dabei erhebliche Speichereinsparungen bei vergleichbarer Leistung ermöglicht.

Das Wesentliche

Stell dir vor, du hast einen riesigen, genialen Koch (das ist dein KI-Modell, wie Llama 3.2), der auf deinem Smartphone kochen soll. Dieser Koch hat zwei sehr unterschiedliche Arbeitsweisen, je nachdem, was er gerade tut:

1. Das "Vorbereiten" (Prefill): Er muss eine ganze Liste von Zutaten (einen langen Text) auf einmal durchgehen. Hier ist er super schnell, wenn er alle Zutaten gleichzeitig auf dem Tisch hat und sie geschickt mischt. Er braucht einen großen, gut organisierten Arbeitsbereich.
2. Das "Servieren" (Decode): Jetzt kommt er Wort für Wort. Er holt sich ein Wort, denkt kurz nach, gibt es aus und holt sich das nächste. Hier ist er nicht mehr durch die Geschwindigkeit seiner Hände begrenzt, sondern durch die Zeit, die es dauert, die Zutaten aus dem Keller (dem Arbeitsspeicher) zu holen.

Das Problem: Der Konflikt zwischen "Keller" und "Tisch"

Das Smartphone hat einen speziellen Keller, den man PIM (Processing-in-Memory) nennt. Das ist wie ein Keller, in dem die Zutaten direkt verarbeitet werden können, ohne dass man sie erst auf den Tisch bringen muss. Das ist extrem schnell für das "Servieren" (Decode).

Aber hier liegt das Problem, das die Forscher PIM-SHERPA genannt haben:

• Der Konflikt der Eigenschaften: Damit der Koch im "Vorbereiten"-Modus schnell ist, müssen die Zutaten auf dem Tisch (im Cache) liegen, damit er sie sofort greifen kann. Damit der Keller (PIM) beim "Servieren" funktioniert, müssen die Zutaten aber im Keller (nicht im Cache) liegen. Wenn sie auf dem Tisch liegen, ruft der Keller sie nicht ab, weil der Tisch sie schon "weggeschnappt" hat.
• Der Konflikt der Anordnung: Im "Vorbereiten"-Modus mag der Koch die Zutaten in einer bestimmten Reihenfolge (wie in einem normalen Kochbuch). Im Keller (PIM) müssen die Zutaten aber ganz anders gestapelt sein, damit der Keller-Verarbeiter sie effizient greifen kann.

Die alte Lösung: Doppelte Vorräte

Bisherige Lösungen sagten: "Okay, wir machen einfach zwei komplette Vorräte."

• Vorrat A: Für den Tisch (gut organisiert für das Vorbereiten).
• Vorrat B: Für den Keller (gut organisiert für das Servieren).

Das Problem? Ein Smartphone hat wenig Platz im Schrank (RAM). Wenn du zwei komplette Vorräte für einen großen Koch hast, passt er nicht mehr in dein Handy. Du müsstest den Koch verkleinern und weniger gut machen.

Die neue Lösung: PIM-SHERPA (Der clevere Butler)

Die Forscher haben eine Software-Lösung namens PIM-SHERPA entwickelt. Sie brauchen keinen neuen Schrank und keine doppelten Vorräte. Stattdessen nutzen sie zwei clevere Tricks, wie ein geschickter Butler:

Trick 1: Der "Zwei-Eimer-Trick" (DRAM Double Buffering)

Stell dir vor, der Koch arbeitet an einem Rezept.

• Eimer 1 steht auf dem Tisch und enthält die Zutaten für das aktuelle Rezept.
• Während der Koch gerade Eimer 1 benutzt, holt der Butler im Hintergrund schon Eimer 2 aus dem Keller und füllt ihn mit den Zutaten für das nächste Rezept.
• Sobald der Koch mit Eimer 1 fertig ist, ist Eimer 2 schon bereit.
• Der Clou: Der Butler sortiert die Zutaten im Eimer 2 genau so um, wie der Keller sie mag, während der Koch arbeitet. So verliert der Koch keine Zeit, und der Butler nutzt die Zeit, die der Koch ohnehin braucht, um die Umordnung zu erledigen.

Trick 2: Der "Sofort-Service" (Online Rearrangement)

Wenn der Koch sehr lange Texte verarbeitet (was immer häufiger wird), dauert das "Vorbereiten" so lange, dass der Butler Zeit hat, die Zutaten für das "Servieren" direkt vor dem Start umzulegen.

• Der Butler holt die Zutaten aus dem Keller, sortiert sie sofort um (wie ein Kartenhaus, das man neu aufbaut) und legt sie auf den Tisch.
• Da der Koch ohnehin lange braucht, um zu kochen, merkt er gar nicht, dass der Butler gerade die Umordnung macht. Es ist wie Warten auf den Kaffee, während man sein Handy benutzt – die Zeit vergeht sowieso.

Warum ist das genial?

1. Platzsparend: Du brauchst nur einen Vorrat (die Zutaten im Keller), plus ein winziges bisschen Platz für den Eimer des Butlers. Das spart etwa 48% Platz im Vergleich zur alten "Doppel-Vorrat"-Methode. Das bedeutet, du kannst viel größere und intelligentere KIs auf deinem Handy laufen lassen.
2. Schnell: Es ist fast so schnell wie die theoretisch beste Methode, aber ohne den riesigen Speicherbedarf.
3. Kein neuer Hardware: Du musst dein Handy nicht umbauen. Es ist nur eine neue Software-Regel, wie der Butler die Zutaten organisiert.

Zusammenfassung in einem Satz:
PIM-SHERPA ist wie ein super-effizienter Butler, der die Zutaten für einen KI-Koch so geschickt hin und her sortiert, dass er sowohl beim schnellen Vorbereiten als auch beim langsamen Servieren immer genau die richtigen Werkzeuge zur Hand hat – und das alles, ohne dass der Schrank (dein Handy-Speicher) überquillt.

Technische Zusammenfassung

Hier ist eine detaillierte technische Zusammenfassung des Papers „PIM-SHERPA" auf Deutsch:

Titel: PIM-SHERPA: Software-Methode für die Inference von LLMs auf Endgeräten durch Behebung von Inkonsistenzen bei PIM-Merkmalen und -Layouts

1. Problemstellung

Die Implementierung von Large Language Models (LLMs) auf mobilen und Edge-Geräten gewinnt zunehmend an Bedeutung. Die Inferenz von LLMs besteht aus zwei Phasen mit unterschiedlichen Anforderungen:

• Prefill-Phase: Verarbeitet viele Tokens gleichzeitig (Matrix-Matrix-Multiplikation, GEMM). Sie ist rechenintensiv und profitiert von einer cachebaren Speicherzuweisung, um Datenwiederverwendung zu maximieren.
• Decode-Phase: Verarbeitet einen Token pro Schritt (Matrix-Vektor-Multiplikation, GEMV). Sie ist speicherbandbreitenintensiv und eignet sich ideal für Processing-in-Memory (PIM), insbesondere LPDDR-PIM, um die DRAM-Bandbreite zu nutzen.

Die Kernprobleme (Inkonsistenzen):

1. Inkonsistenz der Speicherattribute:
   • Für die Prefill-Phase müssen Gewichte im cachebaren Bereich liegen, um Cache-Hits zu nutzen.
   • Für die Decode-Phase (PIM) müssen Gewichte im nicht-cachebaren Bereich liegen. Wenn Gewichte im Cache sind, werden DRAM-Anfragen vom Speichercontroller abgefangen, wodurch die PIM-Ausführung blockiert wird.
2. Inkonsistenz des Gewichts-Layouts:
   • Host-freundliches Layout: Für die CPU/GPU optimiert (z. B. spaltenweise/row-major), um Interleaving über Kanäle und Banken zu nutzen.
   • PIM-bewusstes Layout: Erfordert, dass ganze Matrixzeilen innerhalb einer einzelnen DRAM-Bank zusammenhängend liegen, um die in-Bank-SIMD-Auslastung zu maximieren.

Bisherige Lösungen wie HBM-PIM umgehen dies durch Gewichts-Duplizierung (eine Kopie für den Host, eine für PIM), was den DRAM-Speicherbedarf verdoppelt und auf ressourcenbeschränkten Endgeräten (Smartphones) oft unpraktikabel macht. Andere Ansätze wie FACIL erfordern Hardware-Änderungen am Speichercontroller.

2. Methodik: PIM-SHERPA

PIM-SHERPA ist eine rein softwarebasierte Lösung, die diese Inkonsistenzen zur Laufzeit auflöst, ohne Hardware-Änderungen oder Gewichts-Duplizierung. Sie nutzt zwei Ansätze, die auf kleinen, cachebaren DRAM-Puffern basieren:

A. DRAM Double Buffering (DDB)

• Prinzip: Überlappung von Datenbewegung und Berechnung.
• Mechanismus: Es werden zwei kleine Puffer im cachebaren Bereich reserviert. Während die aktuelle Schicht im ersten Puffer berechnet wird (GEMM), werden die Gewichte der nächsten Schicht aus dem nicht-cachebaren (PIM-)Bereich in den zweiten Puffer vorgezogen (prefetched).
• Swizzled Memory Copy (SMC): Während des Kopierens werden die Gewichte vom PIM-Layout in das Host-freundliche Layout umgewandelt.
• Vorteil: Die Latenz der Layout-Umwandlung wird durch die Rechenzeit der aktuellen Schicht „versteckt" (overlapped). Dies erfordert keine komplexe Synchronisation zwischen Copy- und Compute-Threads.

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

• Prinzip: Serielle Ausführung, optimiert für lange Eingabesequenzen.
• Mechanismus: Unmittelbar vor der Berechnung einer Schicht werden die benötigten Gewichte aus dem nicht-cachebaren Bereich in einen einzelnen cachebaren Puffer kopiert und dabei umsortiert (swizzled).
• Vorteil: Keine komplexe Synchronisation oder Lastverteilung nötig. Die relative Latenz der Umordnung sinkt, da die Rechenzeit (GEMM) bei langen Sequenzen dominiert.
• Anwendung: Besonders effektiv bei interaktiven LLMs und Retrieval-Augmented Generation (RAG), wo die Eingabelänge zunimmt.

3. Schlüsselbeiträge

1. Identifikation der Probleme: Erstmals wurden die spezifischen Inkonsistenzen bei Speicherattributen und Gewichts-Layouts zwischen Prefill- und Decode-Phasen in PIM-fähigen Systemen identifiziert und analysiert.
2. Software-only Lösungen: Entwicklung von DDB und OWR, die ohne Hardware-Modifikationen auskommen und nur einen minimalen Speicher-Overhead (Puffer für 1-2 Schichten) benötigen.
3. Speichereffizienz: Deutliche Reduktion des DRAM-Bedarfs im Vergleich zu Duplizierungsansätzen.
4. Validierung: Implementierung und Evaluation auf einem echten Endgeräte-System (Samsung Galaxy S24+).

4. Ergebnisse

Die Evaluation erfolgte auf einem Samsung Galaxy S24+ mit Exynos 2400 und Llama 3.2 Modellen (1B und 3B Parameter).

• DRAM-Speichereinsparung:
  • PIM-SHERPA (DDB und OWR) spart im Vergleich zur Gewichts-Duplizierung (Weight Duplication) etwa 47,8 % bis 49,7 % DRAM-Kapazität.
  • Dies ermöglicht den Betrieb größerer Modelle auf Smartphones, die sonst durch den Speicherbedarf von Duplizierung ausgeschlossen wären.
• Performance (Time-to-First-Token - TTFT):
  • Bei längeren Eingabesequenzen (Input Sequence Length > 128) erreicht PIM-SHERPA eine TTFT, die mit dem theoretischen Optimum (FACIL-O, das Hardware-Änderungen voraussetzt) vergleichbar ist.
  • DDB: Kann die Umordnungslatenz effektiv hinter der Rechenzeit verstecken.
  • OWR: Zeigt bei sehr langen Sequenzen (z. B. 192 Tokens) eine Performance, die fast mit DDB und FACIL-O mithalten kann.
• Vergleich mit Baseline:
  • Im Vergleich zu einem reinen Host-System (ohne PIM) oder einem System mit nicht-cachebaren Gewichten (NC-GEMM) bietet PIM-SHERPA signifikante Beschleunigungen (bis zu 3,3-fach gegenüber Nicht-PIM-Baselines).
  • Die Emulation von PIM auf dem Gerät zeigte eine hohe Genauigkeit (Fehler < 4 %) im Vergleich zu echter PIM-Hardware.

5. Bedeutung und Fazit

PIM-SHERPA ist ein wegweisender Ansatz, der die praktische Anwendbarkeit von PIM auf kommerziellen Endgeräten erst ermöglicht.

• Überwindung von Hindernissen: Es löst das fundamentale Dilemma, dass Prefill und Decode unterschiedliche Speicheranforderungen haben, ohne den Speicherverbrauch zu verdoppeln.
• Zukunftssicherheit: Mit dem Trend zu längeren Kontextfenstern (Interactive LLMs, RAG) wird der Overhead der Online-Umordnung (OWR) relativ geringer, was die Methode für zukünftige Anwendungen immer attraktiver macht.
• Universalität: Der Ansatz ist nicht auf CPUs beschränkt, sondern kann auch auf GPUs und NPUs mit ähnlichen Speicherhierarchien angewendet werden.

Zusammenfassend demonstriert PIM-SHERPA, dass durch intelligente Software-Methoden die Lücke zwischen theoretischen PIM-Vorteilen und praktischen Implementierungen auf ressourcenbeschränkten Geräten geschlossen werden kann.