Notizie IA Logo

AITalk

Notizie ed analisi sull'Intelligenza Artificiale

PagedAttention e RadixAttention, parliamo di KV Cache

ResearchGenerative AIApplications
Video

pagedattention-radixattention.jpg

Nel pezzo precedente su AITalk avevamo aperto con un numero scomodo: Llama-3.1-70B, in precisione BF16, accumula circa 0,31 megabyte di cache per ogni token elaborato, il che significa che un contesto di 128.000 token costa già 40 gigabyte di memoria GPU, prima ancora di aver caricato una singola richiesta concorrente. Avevamo raccontato tre risposte a questo problema, TurboQuant, OSCAR ed EpiCache, tre modi diversi di rendere quei dati più piccoli, agendo sul numero di bit per valore o su quali porzioni di conversazione vale la pena tenere in memoria.

Ma comprimere i dati risolve solo metà del problema. È come riuscire a impacchettare i faldoni di un archivio in scatole più sottili senza però cambiare il modo in cui l'archivio è organizzato: se gli scaffali restano assegnati in blocco a ogni pratica, indipendentemente da quanto sia effettivamente voluminosa, e se ogni nuova pratica identica alla precedente viene ricopiata da zero invece di essere recuperata, l'efficienza guadagnata sul peso dei singoli fascicoli si perde comunque per strada. Questo articolo racconta le due tecniche che hanno affrontato proprio questo secondo problema, non quanto pesa un token nella cache, ma come lo si alloca e come lo si riusa. Si chiamano PagedAttention e RadixAttention, e non sono alternative alla compressione: operano su un livello diverso, e nei sistemi di produzione più maturi convivono con essa.

Due problemi indipendenti, non uno solo

Anche disponendo di una cache già compressa, un motore di inferenza si scontra con due inefficienze che non hanno nulla a che fare l'una con l'altra. La prima riguarda l'allocazione della memoria. Il metodo tradizionale, quello adottato dai primi sistemi di serving, riservava per ogni richiesta in arrivo un blocco contiguo di memoria GPU dimensionato sulla lunghezza massima di contesto che il modello poteva teoricamente raggiungere. Il problema è che il sistema non ha modo di sapere in anticipo quanto sarà lunga la risposta effettiva, per cui il margine di sicurezza finiva quasi sempre per essere sovradimensionato. Ne derivano due forme di spreco distinte. C'è la frammentazione interna, quando una richiesta prenota spazio per migliaia di token futuri ma ne genera solo poche decine, lasciando inutilizzata la maggior parte della memoria assegnata. E c'è la frammentazione esterna, quando richieste di lunghezza diversa terminano in momenti diversi lasciando buchi sparsi nella memoria GPU: la somma dello spazio libero può essere anche abbondante, ma nessun frammento è abbastanza grande e contiguo da ospitare una nuova richiesta corposa. Il risultato pratico, misurato dai ricercatori che per primi hanno affrontato il problema, era che i sistemi di allora sprecavano tra il 60 e l'80 per cento della memoria allocata, un numero che vale la pena tenere a mente quando si valuta l'entità del guadagno successivo.

Il secondo problema è concettualmente distante dal primo e riguarda il calcolo, non la memoria. Nei carichi di lavoro reali le richieste raramente sono indipendenti tra loro. Migliaia di utenti che interrogano lo stesso assistente condividono lo stesso prompt di sistema, una conversazione a più turni ripropone a ogni scambio l'intera storia precedente, un agente che ragiona in loop continua ad appendere nuovi passaggi a un contesto che in larga parte resta identico al giro prima. In tutti questi casi il motore di inferenza si ritrova a ricalcolare da zero, durante la fase di prefill, gli stessi identici tensori chiave e valore che aveva già prodotto pochi istanti prima per un'altra richiesta, o per lo stesso utente al turno precedente. È un lavoro puramente ridondante, e in scenari di chat prolungata o di sistemi RAG può arrivare a rappresentare la quota più pesante del tempo di risposta percepito dall'utente.

Queste due inefficienze sono nate da esigenze diverse e hanno ricevuto soluzioni diverse. PagedAttention affronta la prima, RadixAttention affronta la seconda.

PagedAttention: la memoria come un sistema operativo

L'idea alla base di PagedAttention, introdotta dal gruppo che ha dato vita a vLLM, prende in prestito quasi di peso un concetto vecchio di decenni: la memoria virtuale a pagine dei sistemi operativi. Chi ha mai aperto il cofano di un computer sa che i programmi non ottengono un blocco fisico contiguo di RAM, ma un'illusione di contiguità costruita sopra pagine sparse gestite da una tabella di traduzione. PagedAttention applica esattamente la stessa logica alla cache di attenzione.

Il meccanismo si articola in quattro passaggi. Anzitutto, la cache di ogni sequenza viene divisa in blocchi logici di dimensione fissa, tipicamente sedici o trentadue token ciascuno, invece di essere trattata come un unico contenitore monolitico. In secondo luogo il sistema mantiene una tabella delle pagine, una mappa che traduce ogni blocco logico nella sua reale collocazione fisica sparsa nella memoria della GPU: durante il calcolo dell'attenzione il motore consulta questa tabella per raccogliere i blocchi necessari, e per il modello la sequenza appare continua anche se fisicamente non lo è affatto. Il terzo passaggio è la crescita su richiesta: la memoria non viene più riservata in blocco all'inizio, ma allocata un blocco alla volta man mano che la generazione procede, così che una risposta di sessanta token occupi soltanto lo spazio di sessanta token, non quello previsto per un contesto massimo che forse non verrà mai raggiunto. Il quarto elemento, forse il più elegante, è la condivisione dei blocchi con copy-on-write: se più richieste iniziano con lo stesso prompt, come capita quotidianamente con un system prompt condiviso da migliaia di utenti, i blocchi fisici corrispondenti vengono referenziati in comune anziché duplicati, e solo nel momento preciso in cui una richiesta diverge dalle altre il blocco interessato viene effettivamente copiato.

Il risultato di questa architettura è che la frammentazione crolla quasi a zero, con margini di spreco che nei benchmark pubblicati scendono sotto il 4 per cento contro il 60-80 per cento dei sistemi precedenti, permettendo di servire un numero molto più alto di richieste concorrenti sullo stesso hardware. È interessante notare che PagedAttention non tocca in alcun modo l'algoritmo di attenzione né altera gli output del modello: la sua innovazione è puramente architetturale, riguarda dove e come i dati vengono fisicamente collocati, non cosa il modello calcola. immagine1.jpg Immagine tratta dal paper su arxiv.org

RadixAttention: un albero che non dimentica

Se PagedAttention risolve il dove, RadixAttention risolve il cosa. L'idea, introdotta dal team di SGLang, parte da un'osservazione semplice quanto trascurata: nei sistemi tradizionali, una volta esaurita una richiesta, la sua cache veniva semplicemente scartata, come se ogni conversazione nascesse e morisse senza lasciare traccia utile per le successive. RadixAttention capovolge questa logica trasformando la cache in una struttura persistente e ricercabile, un radix tree, cioè un trie compresso in cui ogni prefisso unico di token viene memorizzato una sola volta, e richieste diverse che condividono lo stesso inizio percorrono lo stesso ramo dell'albero, biforcandosi soltanto nel punto esatto in cui i loro token cominciano a differire.

Immaginiamo tre utenti che aprono una conversazione con lo stesso messaggio di sistema, seguito ciascuno da una domanda diversa. Invece di conservare tre copie pressoché identiche dello stesso prefisso, l'albero lo mantiene una volta sola e crea tre foglie separate soltanto per l'ultima parte, quella davvero specifica di ciascuna domanda. Quando arriva una nuova richiesta, il motore percorre l'albero token per token cercando il prefisso più lungo già presente, recupera immediatamente i tensori chiave e valore corrispondenti senza ricalcolarli, e infine calcola soltanto il suffisso non ancora visto, inserendo il nuovo percorso nell'albero perché possa essere riutilizzato in futuro. Un aspetto spesso sottovalutato è che la struttura non può crescere all'infinito: quando la memoria GPU si satura, RadixAttention applica una politica di eviction basata su least recently used che elimina per prima le foglie meno usate di recente, proteggendo però i nodi interni condivisi da cui dipendono più richieste, in modo da preservare proprio i prefissi che rendono l'intero meccanismo redditizio.

Il beneficio principale non riguarda la memoria in sé, ma il Time to First Token, cioè quanto tempo passa prima che l'utente veda comparire la prima parola della risposta: se millenovecento token su duemila risultano già presenti nell'albero, il modello deve elaborare soltanto i cento rimanenti, con un impatto diretto sulla latenza percepita, particolarmente evidente in chatbot, sistemi RAG e workflow agentici dove i contesti crescono in modo incrementale invece di essere riscritti da capo ogni volta. È un meccanismo che ricorda, per chi frequenta la narrativa videoludica di nicchia, la struttura di 80 Days, il gioco narrativo di Inkle in cui migliaia di percorsi di viaggio condividono interi tratti di trama comune e si separano solo nei punti di scelta effettiva: il motore del gioco non riscrive da zero ogni possibile diramazione, ma la costruisce innestandola su un tronco narrativo condiviso. immagine2.jpg https://arxiv.org/html/2312.07104v1

Due strati, non due rivali

Un errore comune, alimentato in parte dal modo in cui i due progetti vengono spesso presentati in contrapposizione, è pensare a PagedAttention e RadixAttention come tecnologie concorrenti tra cui scegliere. In realtà rispondono a domande diverse. PagedAttention decide dove i blocchi di cache vivono fisicamente nella memoria GPU, RadixAttention decide se quei blocchi esistono già e possono essere riutilizzati. Il radix tree, in un sistema che integra entrambe le tecniche, punta semplicemente a blocchi che sono a loro volta gestiti dall'allocatore paginato: sono due strati sovrapposti dello stesso stack, non due strade alternative. immagine3.jpg

Tornando all'esempio dei tre utenti con lo stesso system prompt, senza RadixAttention il motore ricalcolerebbe quel prefisso tre volte, e senza PagedAttention ciascuna di quelle tre richieste riserverebbe comunque una porzione di memoria sovradimensionata. Combinando le due tecniche, il prefisso viene calcolato una sola volta, immagazzinato in blocchi paginati compatti, e riutilizzato da ogni richiesta che lo condivide.

Vale la pena precisare che RadixAttention non è l'unica strada verso il riuso automatico dei prefissi. Anche vLLM, il progetto che ha introdotto PagedAttention, supporta oggi il prefix caching automatico, ma con una struttura dati differente: invece di un radix tree, utilizza un meccanismo di hashing a catena, in cui ogni blocco completato riceve un hash calcolato a partire dall'hash del blocco genitore, dai token contenuti nel blocco stesso e da eventuali metadati aggiuntivi. Poiché ogni hash dipende da quello precedente, l'intera catena rappresenta univocamente il prefisso che porta a quel blocco, e se un'altra richiesta produce esattamente la stessa sequenza di token genera la stessa catena di hash, trovando così immediatamente i blocchi già calcolati. Nella pratica, per la maggior parte delle applicazioni, la differenza tra le due implementazioni resta più architetturale che funzionale: entrambi i motori arrivano allo stesso risultato, saltare il prefill dei prefissi già visti, seppure con strutture dati diverse alle spalle.

Oltre le due tecniche: dove va la gestione della cache

PagedAttention e RadixAttention hanno affrontato i due problemi fondativi, ma i contesti continuano a crescere, e con loro sono emerse esigenze che nessuna delle due tecniche, da sola, era stata pensata per coprire.

La prima direzione è la memoria gerarchica. Trattare la GPU come l'unico livello di cache disponibile diventa insostenibile quando i contesti superano le centinaia di migliaia di token: progetti come LMCache e Mooncake organizzano ormai la cache su più livelli, con i blocchi più caldi nella memoria ad alta banda della GPU, quelli usati di recente nella RAM dell'host, e quelli più vecchi ma potenzialmente ancora utili spostati su storage distribuito o dischi NVMe, recuperati automaticamente quando servono di nuovo. È un'architettura che ricorda da vicino, per chi ha letto Zafón, il cimitero dei libri dimenticati: nulla viene distrutto definitivamente, ma i volumi meno consultati vengono spostati in scaffali sempre più remoti, pronti a essere ripescati quando qualcuno ne ha di nuovo bisogno.

La seconda direzione è il routing consapevole della cache. In un'infrastruttura distribuita su più nodi, un normale bilanciatore di carico che smista le richieste in modo puramente round robin rischia di mandare due turni della stessa conversazione a due GPU diverse, vanificando ogni vantaggio del prefix caching anche quando il prefisso condiviso esiste. Sistemi come il router di llm-d affrontano il problema instradando le richieste verso la replica che già possiede la cache del prefisso richiesto, bilanciando non solo il carico ma anche la località dei dati.

La terza direzione, meno discussa ma rilevante per chi opera infrastrutture multi-tenant, riguarda la sicurezza. Il prefix caching introduce un potenziale canale laterale: se un prefisso viene servito visibilmente più in fretta perché già presente in cache, un utente malintenzionato che osserva tempi di risposta insolitamente bassi potrebbe dedurre che un altro utente ha inviato in precedenza lo stesso prompt, senza che il contenuto della risposta venga mai esposto direttamente. La contromisura più diffusa è il cache salting, che include nell'hash usato per identificare i blocchi anche un elemento specifico del tenant, così che richieste identiche provenienti da clienti diversi generino comunque chiavi di cache differenti, eliminando la possibilità di un hit incrociato tra tenant pur mantenendo intatti i benefici di latenza all'interno dello stesso tenant.

Compressione e gestione, due facce della stessa medaglia

Tornando al punto da cui eravamo partiti nell'articolo precedente, la fotografia completa della gestione moderna della KV cache emerge solo mettendo insieme i due fronti. La compressione, di cui TurboQuant, OSCAR ed EpiCache sono espressioni recenti, riduce la dimensione del problema agendo su quanti bit servono per rappresentare ogni token e su quali porzioni di storia conversazionale valga la pena conservare. La gestione, di cui PagedAttention e RadixAttention sono i capisaldi, ne ottimizza la struttura, decidendo dove quei dati vivono fisicamente e se possono essere riutilizzati invece di essere ricalcolati. Un sistema di produzione allo stato dell'arte, oggi, integra entrambi gli approcci: comprime ciò che deve restare in memoria e organizza in modo efficiente lo spazio e il calcolo attorno a quella memoria compressa.

Per chi progetta o valuta un'infrastruttura di inferenza, questo significa che la scelta non è più soltanto quale modello servire, ma come servirlo: padroneggiare il funzionamento di PagedAttention e RadixAttention, oggi, è diventato importante quanto conoscere l'architettura del modello che quella cache sta effettivamente supportando.


Nota tecnica: i dati citati in questo articolo provengono dai paper originali di vLLM e SGLang e dalla documentazione dei relativi progetti open source; non sono stati riprodotti in modo indipendente da questa redazione, e come sempre in questo campo i numeri di laboratorio non garantiscono automaticamente le stesse prestazioni in scenari di produzione con carichi di lavoro reali.