Pooling Engram Conditional Memory in Large Language Models using CXL

Die vorgestellte Arbeit schlägt vor, Engramm-basierte konditionelle Gedächtnisstrukturen für Large Language Models über einen CXL-Memory-Pool auszulagern, um durch feingranularen Zugriff und Prefetching eine skalierbare, kosteneffiziente Lösung zu bieten, die die Inferenzleistung im Vergleich zu DRAM kaum beeinträchtigt.

Das Wesentliche

Stellen Sie sich vor, ein großes Sprachmodell (wie ein sehr kluger KI-Assistent) ist wie ein riesiger Bibliothekar, der in einem riesigen Gebäude arbeitet.

Bisher hatte dieser Bibliothekar ein Problem: Um eine Frage zu beantworten, musste er erst sein eigenes riesiges Gedächtnis (die neuronalen Netze) durchsuchen, um einfache Fakten zu finden. Das war wie wenn er für jede kleine Frage erst durch alle Regale laufen müsste, nur um ein einzelnes Buch zu finden. Das war langsam und ineffizient.

Die Lösung, die in diesem Papier vorgestellt wird, heißt Engram. Man kann sich das wie einen speziellen Schnellzugriff-Index vorstellen. Statt alles im Kopf zu behalten, schaut der Bibliothekar in ein separates, riesiges Nachschlagewerk, das nur Fakten enthält. Das ist viel schneller.

Das neue Problem:
Dieses Nachschlagewerk ist so groß, dass es nicht mehr auf den Schreibtisch des Bibliothekars (den schnellen Arbeitsspeicher des Computers) passt. Es würde den ganzen Platz wegnehmen und den Bibliothekar daran hindern, schnell zu arbeiten.

Die alte Lösung (RDMA):
Früher dachte man: "Wir stellen das Buch in einen anderen Raum und schicken einen Boten, der es holt." Aber dieser Boten (das Netzwerk) braucht Zeit, um die Tür zu öffnen, das Buch zu holen und zurückzulaufen. Für kleine, schnelle Fragen ist dieser Boten zu langsam. Er braucht zu lange für so winzige Buchstaben-Fetzen.

Die neue Lösung (CXL-Pooling):
Hier kommt die Idee des Papiers ins Spiel. Sie nutzen eine neue Technologie namens CXL (Compute Express Link).

Stellen Sie sich CXL nicht wie einen Boten vor, sondern wie einen magischen Teleporter oder einen direkten Rohrpost-Kanal, der den Bibliothekar direkt mit dem riesigen Nachschlagewerk verbindet.

• Der Vergleich:
  • RDMA (Alt): Wie ein Kurierdienst. Sie rufen an, der Kurier kommt, holt das Paket und bringt es. Das dauert lange, besonders wenn Sie nur ein einzelnes Briefmarken-Paket brauchen.
  • CXL (Neu): Wie ein unsichtbarer Schlauch, der direkt vom Regal in die Hand des Bibliothekars führt. Er greift zu, und zack ist das Wissen da. Es fühlt sich fast so an, als wäre das Buch direkt auf seinem Schreibtisch, obwohl es eigentlich in einem anderen Raum liegt.

Was haben die Forscher gemacht?
Sie haben dieses "magische Rohr" (CXL) mit dem "Schnellzugriff-Index" (Engram) verbunden. Sie haben das riesige Nachschlagewerk in einen gemeinsamen Speicherpool gelegt, auf den mehrere Computer gleichzeitig zugreifen können, ohne sich zu stören.

Das Ergebnis:

1. Geschwindigkeit: Der Bibliothekar ist fast genauso schnell wie vorher, obwohl das Buch eigentlich woanders liegt. Der "Teleporter" ist so schnell, dass der Bibliothekar kaum merkt, dass er nicht direkt greift.
2. Kosten: Da viele Bibliothekare (Computer) sich jetzt ein riesiges Nachschlagewerk teilen können, müssen sie nicht mehr jeder ihr eigenes kaufen. Das spart enorm viel Geld.
3. Skalierbarkeit: Wenn die Bibliothek noch größer wird (mehr KI-Modelle), müssen sie nur mehr "Rohre" (CXL-Kabel) verlegen, statt neue ganze Bibliotheken zu bauen.

Zusammengefasst in einem Satz:
Die Forscher haben einen Weg gefunden, wie KI-Modelle riesige Mengen an Wissen in einem günstigen, geteilten Speicher ablegen können, ohne dabei langsamer zu werden – dank einer neuen Technologie, die den Speicher so schnell macht, als wäre er direkt am Prozessor angeschlossen.

Technische Zusammenfassung

Hier ist eine detaillierte technische Zusammenfassung des Papers „Pooling Engram Conditional Memory in Large Language Models using CXL" auf Deutsch:

Problemstellung

Große Sprachmodelle (LLMs) nutzen zunehmend Architekturen wie Mixture-of-Experts (MoE) für effiziente Berechnungen. Ein neuer Ansatz, Engram, führt jedoch eine „bedingte Erinnerung" (Conditional Memory) ein, um statisches Wissen (N-Gram-Einbettungen) von dynamischen Berechnungen zu entkoppeln. Dies ermöglicht eine skalierbare Wissensspeicherung mit $O(1)$-Lookup.

Das Hauptproblem bei Engram ist der enorme Speicherbedarf: Zukünftige Modelle könnten Hunderte von Gigabyte an Engram-Parametern benötigen. Da diese Parameter jedoch nur selten und in kleinen, diskreten Paketen abgerufen werden (sparse access patterns), ist es ineffizient, sie vollständig im teuren lokalen DRAM der GPUs zu halten.

• Herausforderung: Wie kann man diesen Speicherbedarf kosteneffizient in einen geteilten Memory-Pool auslagern, ohne die Inferenz-Latenz zu erhöhen?
• Limitierung bestehender Lösungen: Herkömmliche Ansätze basieren oft auf RDMA (Remote Direct Memory Access). RDMA leidet jedoch unter hohen Latenzen bei kleinen, diskreten Datenpaketen (Overhead durch Protokollstapel und Nachrichtenverarbeitung), was für die feinkörnigen Abrufmuster von Engram ungeeignet ist.

Methodik

Die Autoren schlagen vor, Compute Express Link (CXL) als Technologie für das Memory Pooling zu nutzen.

1. Architektur:

   • Ein zentraler CXL-Switch verbindet mehrere Rechenknoten (Server mit CPUs/GPUs) mit einem gemeinsamen, geteilten CXL-Speicherpool.
   • Im Gegensatz zu RDMA bietet CXL hardwareseitige Unterstützung für native Load/Store-Operationen auf Cache-Line-Ebene. Dies ermöglicht einen direkten Zugriff ohne den Overhead traditioneller Netzwerkstapel.

2. Implementierung in SGLang:

   • Die Engram-Parameter werden in einen CXL-Speicherpool ausgelagert.
   • CXL-zu-CPU: Nutzung von mmap und DAX (Direct Access) für parallele Lesevorgänge mittels OpenMP, die sich wie lokaler DRAM verhalten.
   • CXL-zu-GPU: Entwicklung eines benutzerdefinierten CUDA-Kernels, der Peer-to-Peer (P2P) Datenübertragungen ermöglicht. Dies umgeht die CPU nach dem initialen Mapping und nutzt den PCIe-Bus direkt, um Tausende kleiner Abrufanfragen parallel zu bündeln und die PCIe-Bandbreite voll auszulasten.
   • Prefetching: Die Abrufe werden asynchron gestartet, um die Latenz mit der Berechnung vorheriger Transformer-Blöcke zu überlappen.

3. Analyse der Zugriffsmuster:

   • Engram erfordert nur ca. 5 KB pro Token und Schicht.
   • Die Bandbreitenanforderungen sind gering (ca. 0,7 GB/s im Testfall), aber die Latenzanforderungen sind streng (Prefetch-Fenster von ca. 56 µs für frühe Schichten).

Hauptbeiträge

• Erstes System: Präsentation des ersten Systems, das Engram-Parameter auf einem CXL-basierten Memory Pool auslagert.
• Vergleich RDMA vs. CXL: Eine detaillierte Analyse zeigt, dass CXL aufgrund seiner feinkörnigen Zugriffsmöglichkeiten und niedrigeren Latenz RDMA für Engram-Abrufmuster übertrifft.
• Integration und Performance: Implementierung in das State-of-the-Art Inferenz-Framework SGLang mit Erzielung von Performance, die nahe an lokalem DRAM liegt.

Ergebnisse

Die Experimente wurden auf einem Testbed mit zwei Servern, Intel Xeon-Prozessoren und NVIDIA L20 GPUs durchgeführt, verbunden über einen CXL-Switch und einen 256 GB CXL-Speicher.

• Latenz:
  • CXL-zu-CPU: Die Latenz ist mit lokalem DRAM vergleichbar.
  • CXL-zu-GPU: Obwohl die direkte GPU-Latenz etwas höher ist als bei DRAM, bleibt der Overhead im akzeptablen Bereich und skaliert stabil, selbst bei größeren Engram-Parametern (bis 40B).
  • Im Vergleich dazu liegt die RDMA-Latenz um Größenordnungen höher und ist für Engram ungeeignet.
• End-to-End Durchsatz:
  • Bei Integration in SGLang (getestet mit Qwen3-4B und Qwen3-8B) zeigt sich nur ein marginaler Durchsatzverlust gegenüber dem Baseline-Modell.
  • CXL vs. DRAM: Der Durchsatz mit CXL-Pooling (z. B. 5614,4 tokens/s für Qwen3-4B) ist nahezu identisch mit dem lokalen DRAM-Setup (5683,7 tokens/s).
• Skalierbarkeit:
  • Das System skaliert robust auf mehrere Knoten (nnode=2) und Data-Parallelism (DP=2) mit nur vernachlässigbaren Performance-Einbußen.
• Kostenanalyse:
  • Während CXL für sehr kleine Cluster (z. B. 2 Knoten) aufgrund der Infrastrukturkosten (Switch, Adapter) zunächst teurer sein kann, führt es bei Skalierung (8+ Knoten, große Modelle) zu massiven Kosteneinsparungen.
  • Beispiel: Für ein 400B-Engram-Modell auf 16 Knoten spart die CXL-Lösung ca. 166.000 USD im Vergleich zur lokalen DRAM-Ausstattung jedes Knotens.

Bedeutung und Ausblick

Dieses Paper demonstriert, dass CXL die ideale Infrastruktur für die nächste Generation von speichererweiterten LLMs ist. Es löst das „GPU-Memory Wall"-Problem für Engram-Modelle, indem es:

1. Die Kosten für Speicher drastisch senkt (Shared Memory Pool).
2. Die Performance erhält (nahe DRAM-Latenz durch CXL).
3. Eine skalierbare Architektur bietet, die mit dem Wachstum der Modelle mithält.

Die Arbeit legt den Grundstein für kosteneffiziente, hochperformante Inferenzsysteme, die statisches Wissen effizient von der Rechenleistung entkoppeln können. Zukünftige Arbeiten könnten die Integration von Engram und KV-Cache im selben Pool sowie die Optimierung von RDMA für ähnliche Muster weiter erforschen.