Technisches Resümee: PTStore (Prefix Tensor Store)

Problemstellung

Workloads der Inferenz großer Sprachmodelle (LLMs) sind zum dominierenden Lastfaktor in High-Performance-Computing (HPC)-Rechenzentren geworden und übertreffen das Training sowohl beim Energieverbrauch als auch bei den Ressourcenanforderungen. Die LLM-Inferenz besteht aus zwei Phasen: Prefill (parallele Verarbeitung des Eingabe-Prompts) und Decode (sequenzielle Generierung von Token). Um redundante Berechnungen von Attention-Mechanismen zu vermeiden, nutzen Systeme einen Key-Value (KV)-Cache, um Zwischenergebnisse zu speichern.

Während modernste Laufzeitumgebungen (z. B. vLLM) das KV-Caching innerhalb eines einzelnen GPU- oder Knotens optimieren, stoßen sie bei Skalierung auf erhebliche Einschränkungen:

1. Mangelnde knotenübergreifende Wiederverwendung: Bestehende Systeme versäumen es oft, den Speicher über verteilte Rechenknoten hinweg zu aggregieren. Wenn eine Anfrage auf einem Knoten ein gemeinsames Präfix mit einer Anfrage auf einem anderen Knoten teilt, berechnet der zweite Knoten das Präfix typischerweise neu, anstatt die gecachten Tensoren wiederzuverwenden.
2. Metadaten- und Latenz-Engpässe: Ansätze, die verteiltes Caching versuchen (z. B. LMCache, EvoStore), leiden häufig unter hohen I/O-Overheads durch Remote-Spechtzugriffe oder komplexe Metadaten-Synchronisation (z. B. die Skalierung von Radix-Attention über einen einzelnen Knoten hinaus).
3. Speicherbeschränkungen: Der individuelle GPU-Speicher reicht für große Kontextfenster nicht aus, und das Auslagern in den Host-Speicher oder auf SSDs führt zu Latenzen, die die Vorteile des Caching zunichtemachen.

Die zentrale Herausforderung besteht darin, eine skalierbare, latenzarme Wiederverwendung von KV-Cache-Präfixen über eine große Anzahl von GPUs zu ermöglichen, die über viele Rechenknoten verteilt sind, ohne prohibitive I/O- oder Metadaten-Overheads zu verursachen.

Methodik: PTStore-Architektur

PTStore (Prefix Tensor Store) ist ein verteilter, replizierter Tensor-Speicher, der darauf ausgelegt ist, diese Einschränkungen durch die Verteilung und Replikation populärer KV-Cache-Präfixe zu adressieren. Das System verwendet ein Client-Server-Modell, bei dem jeder Rechenknoten einen Server ausführt, der den lokalen Host-Speicher und SSDs aggregiert, um sowohl lokale als auch remote befindliche GPU-Clients zu bedienen.

Zentrale Designprinzipien

1. Inkrementelle Tensor-Speicherung (Trie-ähnliche Struktur):

   • Anstatt vollständige KV-Blöcke zu speichern, speichert PTStore inkrementelle Differenzen (Tensoren) zwischen einem neuen Objekt und dem längsten gemeinsamen Präfix (Longest Common Prefix, LCP) zuvor gespeicherter Objekte.
   • Dies ermöglicht es Präfixen, über die Zeit redundantfrei in divergierende Richtungen zu wachsen, ähnlich einem Trie, jedoch implementiert auf der Granularität von Tensoren.
   • Konsolidierte Metadaten: Um eine teure verteilte Trie-Traversierung zu vermeiden, nutzt PTStore eine flache Metadatenstruktur. Die Metadaten jedes Objekts enthalten eine Liste eindeutiger Tensor-IDs. Ein Ladeoperation iteriert durch diese IDs, um die lokale Existenz im Replikations-Cache zu prüfen; falls diese fehlt, werden sie remote vom „Besitzer“-Server abgerufen.

2. Verteiltes hierarchisches Caching mit Replikation:

   • Owned Cache (Besitz-Cache): Speichert die inkrementellen Tensoren, für die ein spezifischer Server verantwortlich ist.
   • Replication Cache (Replikations-Cache): Speichert Kopien von „heißen“ (populären) Präfixen lokal auf dem Server, um die Zugriffslokalität zu verbessern.
   • Trade-off-Management: Das System verwaltet einen konfigurierbaren Schwellenwert zwischen Owned- und Replication-Caches. Es priorisiert das Verwerfen von replizierten Tensoren (die erneut abgerufen werden können) gegenüber dem Verwerfen von Owned-Tensoren (die in langsameren Speicher ausgelagert werden müssen), um die Abrufgeschwindigkeit gegen die Speicherkapazität abzuwägen.

3. Zugriffsmuster-bewusste Eviktierung:

   • PTStore nutzt eine frequenzbasierte Eviktierungsstrategie (adaptiert von GDSF) anstelle von Least Recently Used (LRU), da die Struktur der Präfixe bedeutet, dass frühere Tensoren häufiger aufgerufen werden.
   • Es berücksichtigt den Trade-off zwischen Größe und Frequenz, um sicherzustellen, dass kleine, häufig aufgerufene Tensoren keine größeren, schwerer abrufbaren Tensoren verdrängen.

4. RDMA-bewusste Konsolidierung:

   • Um Streuung zu minimieren, werden die an ein LCP angehängten Inkremente in einer einzigen zusammenhängenden Region auf dem Besitzer-Server konsolidiert.
   • Ladeoperationen nutzen Bulk-RDMA (Remote Direct Memory Access), um verstreute Segmente parallel über einen einzigen RPC abzurufen, wodurch der Overhead entfällt, Daten erst in eine zusammenhängende Region kopieren zu müssen, bevor sie übertragen werden.

Kernbeiträge

1. Designprinzipien: Ein Satz hochrangiger Prinzipien für ein verteiltes Repository, das inkrementelle Tensor-Speicherung, konsolidierte Metadaten und Präfix-Replikation integriert.
2. PTStore-Prototyp: Ein Forschungsprototyp, der diese Prinzipien implementiert und über eine C++ Low-Level-API sowie eine Python-Schnittstelle für die nahtlose Integration in LLM-Laufzeitumgebungen wie vLLM verfügt.
3. Leistungsvalidierung: Umfangreiche Experimente, die eine signifikante Reduktion des I/O-Overheads und der End-to-End-Laufzeit im Vergleich zu State-of-the-Art-Baselines demonstrieren.

Experimentelle Ergebnisse

Die Autoren evaluierten PTStore auf dem ALCF Polaris HPC-Testbed (560 Knoten, A100 GPUs) unter Verwendung von zwei extraktiven QA-Workloads: WikiQA (langer Kontext) und SQUAD (hohes Volumen an Fragen). Das verwendete LLM war Mistral-7B-instruct-V2.

Baselines

• vLLM Vanilla: Standard-vLLM ohne lokales Präfix-Sharing zwischen Anfragen.
• vLLM Prefix: vLLM mit lokalem Präfix-Sharing (innerhalb eines Knotens).
• EvoStore: Ein verteilter Tensor-Speicher, der inkrementelle Speicherung und RDMA nutzt, aber keine lokale Präfix-Replikation besitzt.
• PTStore: Das vorgeschlagene System mit verteilter Bewusstheit und lokaler Replikation.

Erkenntnisse

• Schwache Skalierbarkeit (8–32 GPUs): PTStore übertraf EvoStore und vLLM Prefix signifikant. Während EvoStore unter hohem RDMA-I/O-Overhead litt, wenn entfernte Präfixe abgerufen wurden, milderte PTStores lokale Replikation dies ab, was zu einem „detached Advantage“ bei der Time to First Token (TTFT) führte.
• Sequenzlängen-Skalierbarkeit (1k–8k Token):
  • Bei kurzen Sequenzen (1k) war das lokale Caching von vLLM wettbewerbsfähig.
  • Mit zunehmender Sequenzlänge wuchs der Vorteil von PTStore. Bei 8k Token war PTStore fast 2-mal schneller als das Präfix-Caching von vLLM und 20 % schneller als EvoStore.
  • Die Leistungsdifferenz vergrößerte sich bei längeren Kontexten, da die Kosten für die Neuberechnung oder den Remote-I/O den Nutzen eines rein lokalen Cachings überstiegen.
• Effizienzgewinne: Bei langen Passage-QA-Datensätzen führte PTStore Inferenzen 5- bis 6-mal effizienter aus als Baselines, die keine Speicherkapazitäten über Knoten hinweg aggregieren und KV-Caches neu generieren müssen.

Bedeutung und Ansprüche

Das Paper behauptet, dass PTStore eine kritische Lücke in der skalierbaren LLM-Inferenz-Bereitstellung schließt: die Unfähigkeit aktueller Systeme, KV-Cache-Präfixe effizient über verteilte Knoten hinweg wiederzuverwenden. Durch die Kombination von inkrementeller Speicherung zur Minimierung von Redundanz, konsolidierten Metadaten für schnelle Abfragen und einer Replikationsstrategie zur Optimierung der Lokalität ermöglicht PTStore:

• Eine um Größenordnungen Erweiterung der effektiven KV-Cache-Größe durch Aggregation des Speichers über den Cluster hinweg.
• Eine signifikante Reduktion der TTFT, insbesondere bei Long-Context-Workloads, bei denen die Neuberechnung kostspielig ist.
• Skalierbarkeit, die die Kommunikationsengpässe und Metadaten-Synchronisationsprobleme vermeidet, die bisherige verteilte Ansätze plagen.

Die Autoren positionieren PTStore als einen grundlegenden Schritt hin zur skalierbaren KI-Inferenz und merken an, dass sich zukünftige Arbeiten auf dynamisches Memory-Balancing, ML-basierte Eviktierungsrichtlinien und breiteres Benchmarking gegen Systeme wie LMCache und Mooncake auf realen Konversations- und Code-Completion-Traces konzentrieren werden.