Technische Samenvatting: PTStore (Prefix Tensor Store)

Probleemstelling

Workloads voor Large Language Model (LLM) inferentie zijn de dominante belasting geworden in high-performance computing (HPC) datacentra, waarbij ze de energieconsumptie en de vraag naar resources hebben overtroffen ten opzien van training. LLM-inferentie bestaat uit twee fasen: prefill (het parallel verwerken van de input-prompt) en decode (het sequentieel genereren van tokens). Om redundante berekeningen van aandachtmechanismen te voorkomen, maken systemen gebruik van een Key-Value (KV) cache om tussenresultaten op te slaan.

Hoewel state-of-the-art runtimes (bijv. vLLM) KV-caching optimaliseren binnen een enkele GPU of node, kampen zij bij schaalvergroting met significante beperkingen:

1. Gebrek aan Cross-Node Hergebruik: Bestaande systemen falen vaak in het aggregeren van geheugen over gedistribueerde compute-nodes. Als een verzoek op één node een gedeelde prefix heeft met een verzoek op een andere node, zal de tweede node de prefix doorgaans opnieuw berekenen in plaats van de gecachte tensoren te hergebruiken.
2. Metadata en Latentie-bottlenecks: Benaderingen die proberen gedistribueerde caching toe te passen (bijv. LMCache, EvoStore) lijden vaak onder hoge I/O-overhead door remote geheugentoegang of complexe metadata-synchronisatie (bijv. het schalen van Radix-Attention buiten een enkele node).
3. Geheugenbeperkingen: Individueel GPU-geheugen is onvoldoende voor grote contextvensters, en het offloaden naar hostgeheugen of SSD's introduceert latentie die de voordelen van caching tenietdoet.

De kernuitdaging is het mogelijk maken van schaalbare, lage-latentie hergebruik van KV-cache prefixes over een groot aantal GPU's verdeeld over vele compute-nodes, zonder dat dit gepaard gaat met prohibitieve I/O- of metadata-overhead.

Methodologie: PTStore Architectuur

PTStore (Prefix Tensor Store) is een gedistribueerde, gerepliceerde tensor store die ontworpen is om deze beperkingen aan te pakken door populaire KV-cache prefixes te distribueren en te repliceren. Het systeem hanteert een client-servermodel waarbij elke compute-node een server draait die lokaal hostgeheugen en SSD's aggregeert om zowel lokale als remote GPU-clients te bedienen.

Belangrijke Ontwerpprincipes

1. Incrementele Tensoropslag (Trie-achtige Structuur):

   • In plaats van volledige KV-blokken op te slaan, slaat PTStore incrementele verschillen (tensoren) op tussen een nieuw object en de langste gemeenschappelijke prefix (LCP) van eerder opgeslagen objecten.
   • Dit stelt prefixes in staat om zonder redundantie te groeien in divergente richtingen, vergelijkbaar met een trie, maar geïmplementeerd op de granulariteit van tensoren.
   • Geconsolideerde Metadata: Om dure gedistribueerde trie-traversal te vermijden, gebruikt PTStore een vlakke metadata-structuur. De metadata van elk object bevat een lijst met unieke tensor-ID's. Een load-operatie iterereert door deze ID's om te controleren op lokale aanwezigheid in de replicatie-cache; indien ontbrekend, worden ze remote opgehaald bij de "owner" server.

2. Gedistribueerde Hiërarchische Caching met Replicatie:

   • Owned Cache: Slaat de incrementele tensoren op waarvoor een specifieke server verantwoordelijk is.
   • Replication Cache: Slaat kopieën van "hotte" (populaire) prefixes lokaal op de server op om de toegangslocatie te verbeteren.
   • Trade-off Management: Het systeem beheert een configureerbare drempelwaarde tussen owned en gerepliceerde caches. Het geeft prioriteit aan het verwijderen van gerepliceerde tensoren (die opnieuw opgehaald kunnen worden) boven het verwijderen van owned tensoren (die naar tragere opslag geflushed moeten worden) om de balans te bewaren tussen ophaalsnelheid en opslagcapaciteit.

3. Access Pattern-bewuste Evictie:

   • PTStore maakt gebruik van een frequentiegebaseerd evictiebeleid (geadapteerd van GDSF) in plaats van Least Recently Used (LRU), aangezien de prefix-structuren betekenen dat eerdere tensoren vaker worden geraadpleegd.
   • Het houdt rekening met de trade-off tussen grootte versus frequentie, waardoor wordt gewaarborgd dat kleine, frequente tensoren geen grotere, moeilijk op te halen tensoren verdringen.

4. RDMA-bewuste Consolidatie:

   • Om versnippering te minimaliseren, worden increments die aan een LCP worden toegevoegd geconsolideerd in één enkel aaneengesloten gebied op de owner-server.
   • Load-operaties gebruiken bulk RDMA (Remote Direct Memory Access) om verspreide segmenten parallel te fetchen via een enkele RPC, waardoor de overhead van het kopiëren van data naar een aaneengesloten regio vóór de transfer wordt vermeden.

Belangrijkste Bijdragen

1. Ontwerpprincipes: Een reeks hoogwaardige principes voor een gedistribueerde repository die incrementele tensoropslag, geconsolideerde metadata en prefix-replicatie integreert.
2. PTStore Prototype: Een onderzoeksprototype dat deze principes implementeert, voorzien van een C++ low-level API en een Python-interface voor naadloze integratie met LLM runtimes zoals vLLM.
3. Prestatievalidatie: Uitgebreide experimenten die aantonen dat de I/O-overhead en de end-to-end runtime aanzienlijk worden verminderd in vergelijking met state-of-the-art baselines.

Experimentele Resultaten

De auteurs evalueerden PTStore op de ALCF Polaris HPC testbed (560 nodes, A100 GPU's) met behulp van twee extractive QA-workloads: WikiQA (lange context) en SQUAD (hoog volume aan vragen). Het gebruikte LLM was Mistral-7B-instruct-V2.

Baselines

• vLLM Vanilla: Standaard vLLM zonder cross-request prefix sharing.
• vLLM Prefix: vLLM met lokale prefix sharing (binnen een node).
• EvoStore: Een gedistribueerde tensor store die incrementele opslag en RDMA gebruikt, maar geen lokale prefix-replicatie heeft.
• PTStore: Het voorgestelde systeem met gedistribueerd bewustzijn en lokale replicatie.

Bevindingen

• Zwakke Schaalbaarheid (8–32 GPU's): PTStore presteerde aanzienlijk beter dan EvoStore en vLLM Prefix. Terwijl EvoStore leed onder hoge RDMA I/O-overhead bij het ophalen van remote prefixes, mitigeerde de lokale replicatie van PTStore dit, wat resulteerde in een "detached advantage" in Time to First Token (TTFT).
• Sequentielengte Schaalbaarheid (1k–8k tokens):
  • Voor korte sequenties (1k) was de lokale caching van vLLM competitief.
  • Naarmate de sequentielengte toenam, groeide het voordeel van PTStore. Bij 8k tokens was PTStore bijna 2x sneller dan de prefix caching van vLLM en 20% sneller dan EvoStore.
  • De prestatiekloof werd groter bij langere contexten omdat de kosten van herberekening of remote I/O de voordelen van lokale-only caching overtroffen.
• Efficiëntiewinst: Op lange passage Q&A datasets voerde PTStore inferenties 5–6 keer efficiënter uit dan baselines die geen geheugen over nodes aggregeren en de KV-caches opnieuw moeten genereren.

Betekenis en Claims

Het artikel claimt dat PTStore een kritiek gat in de schaalbare LLM-inferentie-serving opvult: het onvermogen van huidige systemen om KV-cache prefixes efficiënt te hergebruiken over gedistribueerde nodes. Door incrementele opslag te combineren om redundantie te minimaliseren, geconsolideerde metadata voor snelle queries, en een replicatiestrategie om de lokatie te optimaliseren, maakt PTStore het volgende mogelijk:

• Ordes van grootte expansie van de effectieve KV-cache grootte door het geheugen over de cluster te aggregeren.
• Significante reductie in TTFT, met name voor long-context workloads waar herberekening kostbaar is.
• Schaalbaarheid die de communicatie-bottlenecks en metadata-synchronisatieproblemen vermijdt die eerdere gedistribueerde benaderingen teisteren.

De auteurs positioneren PTStore als een fundamentele stap richting schaalbare AI-inferentie, waarbij zij opmerken dat toekomstig werk zich zal richten op dynamische geheugenbalancering, ML-gebaseerde eviction-policies en bredere benchmarking tegen systemen zoals LMCache en Mooncake op real-world conversationele en code-completion traces.