Riepilogo Tecnico: PTStore (Prefix Tensor Store)

Problematica

I carichi di lavoro per l'inferenza di Large Language Models (LLM) sono diventati il carico dominante nei data center per l'high-performance computing (HPC), superando l'addestramento per consumo energetico e richiesta di risorse. L'inferenza degli LLM consiste in due fasi: prefill (elaborazione del prompt di input in parallelo) e decode (generazione sequenziale dei token). Per evitare la ridondanza nel calcolo dei meccanismi di attenzione, i sistemi utilizzano una Key-Value (KV) cache per memorizzare i risultati intermedi.

Sebbene i runtime all'avanguardia (ad es. vLLM) ottimizzino la cache KV all'interno di un singolo GPU o nodo, essi affrontano limitazioni significative su scala:

1. Mancanza di riutilizzo cross-nodo: I sistemi esistenti spesso non riescono ad aggregare la memoria tra nodi di calcolo distribuiti. Se una richiesta su un nodo condivide un prefisso con una richiesta su un altro nodo, il secondo nodo tipicamente ricalcola il prefisso invece di riutilizzare i tensori memorizzati in cache.
2. Colli di bottiglia di metadati e latenza: Approcci che tentano la caching distribuita (ad es. LMCache, EvoStore) soffrono spesso di elevati overhead di I/O dovuti all'accesso alla memoria remota o a una complessa sincronizzazione dei metadati (ad es. scalare Radix-Attention oltre un singolo nodo).
3. Vincoli di memoria: La memoria di un singolo GPU è insufficiente per ampie finestre di contesto, e l'offloading verso la memoria host o gli SSD introduce una latenza che nega i benefici della caching.

La sfida principale è abilitare un riutilizzo scalabile e a bassa latenza dei prefissi della cache KV attraverso un gran numero di GPU distribuite su molti nodi di calcolo senza incorrere in proibitivi overhead di I/O o di metadati.

Metodologia: Architettura PTStore

PTStore (Prefix Tensor Store) è un archivio di tensori replicato e distribuito progettato per affrontare queste limitazioni distribuendo e replicando i prefissi popolari della cache KV. Il sistema impiega un modello client-server in cui ogni nodo di calcolo esegue un server che aggrega la memoria host locale e gli SSD per servire sia i client GPU locali che quelli remoti.

Principi di Design Chiave

1. Archiviazione Incrementale dei Tensori (Struttura di tipo Trie):

   • Invece di memorizzare blocchi KV completi, PTStore memorizza le differenze incrementali (tensori) tra un nuovo oggetto e il più lungo prefisso comune (LCP) degli oggetti precedentemente memorizzati.
   • Ciò consente ai prefissi di crescere in modo privo di ridondanze nel tempo in direzioni divergenti, similmente a un trie, ma implementato tramite granularità a livello di tensore.
   • Metadati Consolidati: Per evitare l'costosa traversata di un trie distribuito, PTStore utilizza una struttura di metadati piatta. I metadati di ogni oggetto contengono un elenco di ID di tensori unici. Un'operazione di caricamento itera attraverso questi ID per verificare la presenza locale nella cache di replicazione; se mancante, li recupera remotamente dal server "proprietario".

2. Caching Gerarchico Distribuito con Replicazione:

   • Cache Posseduta (Owned Cache): Memorizza i tensori incrementali di cui uno specifico server è responsabile.
   • Cache di Replicazione (Replication Cache): Memorizza copie di prefissi "caldi" (popolari) localmente sul server per migliorare la località di accesso.
   • Gestione del Trade-off: Il sistema gestisce una soglia configurabile tra le cache possedute e quelle replicate. Prioritizza lo scarto dei tensori replicati (che possono essere recuperati nuovamente) rispetto all'evizione dei tensori posseduti (che richiedono il salvataggio su storage più lento) per bilanciare la velocità di recupero rispetto alla capacità di archiviazione.

3. Evizione Sensibile ai Pattern di Accesso:

   • PTStore utilizza una politica di evizione basata sulla frequenza (adattata da GDSF) piuttosto che sul Least Recently Used (LRU), poiché le strutture dei prefissi implicano che i tensori iniziali siano accessi più frequentemente.
   • Tiene conto del compromesso tra dimensione e frequenza, assicurando che tensori piccoli e frequenti non spostino tensori più grandi e costosi da recuperare.

4. Consolidamento Consapevole di RDMA:

   • Per minimizzare la dispersione, gli incrementi aggiunti a un LCP vengono consolidati in una singola regione contigua sul server proprietario.
   • Le operazioni di carico utilizzano RDMA bulk (Remote Direct Memory Access) per recuperare segmenti sparsi in parallelo tramite una singola RPC, evitando l'overhead di copiare i dati in una regione contigua prima del trasferimento.

Contributi Chiave

1. Principi di Design: Un insieme di principi di alto livello per un repository distribuito che integra l'archiviazione incrementale dei tensori, i metadati consolidati e la replicazione dei prefissi.
2. Prototipo PTStore: Un prototipo di ricerca che implementa questi principi, caratterizzato da un'API a basso livello in C++ e un'interfaccia Python per un'integrazione fluida con i runtime di LLM come vLLM.
3. Validazione delle Prestazioni: Esperimenti estesi che dimostrano una significativa riduzione dell'overhead di I/O e del tempo di esecuzione end-to-end rispetto ai baseline allo stato dell'arte.

Risultati Sperimentali

Gli autori hanno valutato PTStore sull'ALCF Polaris HPC testbed (560 nodi, GPU A100) utilizzando due carichi di lavoro di estrazione QA: WikiQA (contesto lungo) e SQUAD (alto volume di domande). L'LLM utilizzato è stato Mistral-7B-instruct-V2.

Baseline

• vLLM Vanilla: vLLM standard senza condivisione del prefisso tra le richieste.
• vLLM Prefix: vLLM con condivisione del prefisso locale (all'interno di un nodo).
• EvoStore: Un archivio di tensori distribuito che utilizza l'archiviazione incrementale e RDMA ma manca della replicazione locale del prefisso.
• PTStore: Il sistema proposto con consapevolezza distribuita e replicazione locale.

Risultati

• Debole Scalabilità (8–32 GPU): PTStore ha superato significativamente EvoStore e vLLM Prefix. Mentre EvoStore soffriva di elevati overhead di I/O RDMA durante il recupero di prefissi remoti, la replicazione locale di PTStore ha mitigato questo problema, risultando in un "vantaggio distaccato" nel Time to First Token (TTFT).
• Scalabilità della Lunghezza della Sequenza (1k–8k token):
  • Per sequenze brevi (1k), la caching locale di vLLM era competitiva.
  • All'aumentare della lunghezza della sequenza, il vantaggio di PTStore cresceva. A 8k token, PTStore era quasi 2 volte più veloce della caching del prefisso di vLLM e il 20% più veloce di EvoStore.
  • Il divario di prestazioni si amplia con contesti più lunghi perché il costo della ricalcolazione o dell'I/O remoto supera i benefici della caching solo locale.
• Guadagni di Efficienza: Sui dataset di domanda/risposta su passaggi lunghi, PTStore ha eseguito le inferenze 5–6 volte più efficientemente rispetto ai baseline che non aggregano la memoria tra i nodi e richiedono la rigenerazione delle cache KV.

Significato e Claim

Il paper sostiene che PTStore affronti una lacuna critica nell'inferenza scalabile degli LLM: l'incapacità degli attuali sistemi di riutilizzare efficientemente i prefissi della cache KV attraverso nodi distribuiti. Combinando l'archiviazione incrementale per minimizzare la ridondanza, i metadati consolidati per query rapide e una strategia di replicazione per ottimizzare la località, PTStore abilita:

• Un'espansione di ordini di grandezza della dimensione effettiva della cache KV aggregando la memoria attraverso il cluster.
• Una significativa riduzione del TTFT, in particolare per i carichi di lavoro a contesto lungo dove la ricalcolazione è costosa.
• Scalabilità che evita i colli di bottiglia della comunicazione e i problemi di sincronizzazione dei metadati che affliggono i precedenti approcci distribuiti.

Gli autori posizionano PTStore come un passo fondamentale verso l'inferenza AI scalabile, notando che il lavoro futuro si concentrerà sul bilanciamento dinamico della memoria, politiche di evizione basate su ML e benchmarking più ampio contro sistemi come LMCache e Mooncake su tracce reali di conversazione e completamento del codice.