KV cache: come un modello ricorda i token che ha già letto

La KV cache spiegata senza giri di parole: cosa si salva davvero, quanta memoria occupa, perche GQA e paged attention hanno cambiato i costi.
Immagine realizzata con IA
Le immagini presenti nel post possono essere realizzate o elaborate con strumenti di Intelligenza Artificiale. Questa scelta nasce anche dalla necessità di prevenire le continue contestazioni e richieste economiche improprie che ricevo relative a fotografie utilizzate che sono sempre state acquistate con regolare licenza.
Contenuti dell'articolo

Un modello che scrive una risposta non la pensa tutta insieme: la costruisce un token alla volta. Ogni token nuovo deve guardare tutti quelli che lo precedono. Se a ogni passo rifacesse da capo i conti sui token già letti, il costo esploderebbe. La KV cache esiste per non ripetere quel lavoro.

Il problema: ricalcolare tutto a ogni token

La generazione è autoregressiva. Prendi il prompt, produci un token, lo accodi al testo, ripeti. A ogni giro l’attenzione confronta il token appena nato con tutti i precedenti, per capire a cosa somiglia e da dove prendere informazione.

Senza cache, quel confronto si paga rifacendo le proiezioni su tutta la sequenza. Al passo numero cento ricalcoli key e value per tutti i cento token, anche se sono identici a quelli del passo novantanove: stessi pesi, stessi input, stessi risultati. È lavoro buttato e cresce in modo quadratico con la lunghezza del testo.

Fai il conto su una conversazione normale. Un contesto da 2.000 token e una risposta da 500: senza cache il modello rifà le proiezioni di 2.000 token per cinquecento volte. Con la cache le fa una volta sola, all’inizio e poi solo per il token nuovo.

Cosa si salva dentro la cache

In ogni layer del Transformer l’attenzione produce tre matrici: query, key e value. La query serve al token appena arrivato, per chiedere “a cosa devo badare”. Key e value servono a rispondere, quindi devono restare disponibili per tutti i token futuri. Per questo si conservano solo K e V e la cache si chiama così. Se ti serve il meccanismo completo, l’ho raccontato in query, key e value.

La struttura è la stessa per ogni blocco: per layer, per testa, due matrici grandi quanto la sequenza per la dimensione della testa. Non è una struttura fissa come i pesi: cresce a ogni token aggiunto e si svuota quando la richiesta finisce.

Il funzionamento a regime è semplice. Arriva il token nuovo, si calcolano la sua query e le sue key e value, si accodano le due nuove matrici e la query fa il prodotto scalare con tutte le key in memoria. Un token di calcolo, tutto il resto in lettura.

Il conto della memoria

Qui iniziano i guai, perché la cache non è gratis. Prendi un modello da 7 miliardi di parametri con 32 layer, 32 teste e dimensione della testa 128, quindi 4.096 dimensioni per layer. Per ogni token servono due matrici da 32 teste per 128 valori, su 32 layer: 262.144 numeri per token.

In precisione fp16, due byte per numero, significa circa mezzo megabyte per token. Su un contesto di 8.000 token sono circa 4 GB di memoria video. Su 32.000 token si arriva a 16 GB. Se servi sedici richieste insieme, moltiplica tutto per sedici.

Quanta memoria occupa la KV cache: da 4 GB con 8.000 token a 64 GB con 128.000 token, su un modello da 7 miliardi in fp16
La cache cresce in modo lineare con il testo: mezzo megabyte per token

Il punto è che i pesi del modello da 7 miliardi in fp16 occupano circa 14 GB. Con un contesto lungo e un po’ di concorrenza, la cache supera i pesi. Chi dimensiona le GPU per l’inferenza lo sa bene: non è raro che il collo di bottiglia sia la memoria occupata dai token già letti, non il modello.

Prefill e decode: due lavori diversi

L’inferenza ha due fasi che si comportano agli antipodi. La prima si chiama prefill: il prompt arriva tutto insieme e il modello lo processa in parallelo, con matrici larghe. La GPU sta al massimo del calcolo, si dice che il carico è compute-bound.

La seconda è il decode: si genera un token per volta. Le matrici sono strettissime, una riga alla volta e il tempo se ne va leggendo pesi e cache dalla memoria. Qui la GPU spesso non è satura: aspetta i dati. Il carico è memory-bandwidth-bound.

Questa asimmetria spiega due cose che noti usando i modelli. Il primo token di una risposta arriva dopo una pausa che sembra eterna, perché il prompt va processato. Poi i token escono veloci. E spiega perché i server mettono più richieste in parallelo: nel decode la potenza di calcolo libera va riempita con altro lavoro.

Condividere le teste: GQA e MQA

La prima contromisura è non avere una cache per ogni testa. Nella multi-query attention, MQA, tutte le teste delle query condividono un’unica testa per key e value. Su 32 teste la cache si riduce di trentadue volte, a costo di una piccola perdita di qualità.

La via di mezzo si chiama grouped-query attention, GQA. Le teste delle query restano tante, quelle di key e value si raggruppano: con 32 query e 8 teste KV la cache si divide per quattro. È la scelta di molti modelli recenti, da LLaMA 3 a Mistral.

Il risultato è che a parità di contesto e di scheda video puoi servire molte più richieste. Qualità quasi identica, memoria molto più bassa: è uno di quei compromessi che conviene sempre e infatti è diventato lo standard.

Paged attention e batch continuo

Attenzione a non confondere le due ottimizzazioni con un trucco di benchmark. La cache è memoria vera, allocata e liberata in tempo reale: gestirla male significa o sprecare spazio o bloccare nuove richieste. Per questo i server di inferenza dedicano più codice alla gestione della cache che al modello stesso.

Il secondo problema è l’allocazione. Se prenoti un blocco di memoria contiguo grande quanto la lunghezza massima possibile, butti via spazio per ogni richiesta corta. E con richieste di lunghezza imprevedibile la frammentazione cresce fino a rendere impossibile aggiungerne altre.

La soluzione arrivata da vLLM nel 2023 si chiama paged attention e copia l’idea della memoria virtuale dei sistemi operativi. La cache si spezza in blocchi da sedici token, sparsi dove c’è posto, con una tabella che tiene le mappe. Nessun blocco contiguo da prenotare, spreco ridotto all’osso, molte più richieste in parallelo.

A questo si aggiunge il batch continuo: le richieste entrano ed escono dal gruppo a ogni passo, invece di aspettare che tutto il gruppo finisca. Una risposta breve libera subito il suo posto. Insieme, queste due tecniche hanno cambiato il costo dell’inferenza più di molti modelli nuovi.

Quando la cache si butta

Se il contesto è enorme e la memoria finita, una strada è dimenticare. La finestra scorrevole tiene solo gli ultimi N token e scarta i più vecchi: la memoria resta costante, il modello continua a funzionare, ma perde tutto quello che è uscito dalla finestra. Mistral 7B usava 4.096 token.

C’è però un dettaglio scoperto sperimentando: i primi token della sequenza ricevono sempre molta attenzione, anche quando sembrano inutili. Si chiamano attention sink. StreamingLLM li tiene sempre in cache e il modello continua a comportarsi bene anche molto oltre la finestra addestrata.

C’è poi il caso opposto: conversazioni lunghe in cui serve ricordare tutto. Lì la finestra non basta e si accetta il costo in memoria, magari riducendo la precisione della cache o servendo meno utenti in parallelo. Non esiste una scelta giusta in assoluto, esiste il compromesso che regge il carico.

Le altre strade sono comprimere. Si può quantizzare la cache a 8 bit o a 4, perdendo un po’ di precisione, oppure spostare le parti vecchie su CPU e disco, accettando un accesso più lento. Ogni soluzione sposta il problema, nessuna lo elimina.

Prompt caching e KV cache non sono la stessa cosa

Qui nasce la confusione più comune. Il prompt caching è un servizio commerciale: se la parte iniziale della richiesta è identica a una già vista, il fornitore riusa il prefill e ti fa pagare meno i token di input.

La KV cache è una struttura interna: matrici in memoria video, per singola richiesta, durante la generazione. Una si vede in fattura, l’altra si vede nel consumo di VRAM. Il nome inganna, ma sono due cose diverse che lavorano sullo stesso pezzo di calcolo.

Cosa cambia per chi usa i modelli

La prima conseguenza è sul contesto lungo. Caricare centomila token di documenti occupa diversi gigabyte di cache per una sola richiesta: bello, costoso e con poca concorrenza disponibile.

La seconda è sui tempi. Il primo token paga il prefill di tutto il prompt, i successivi no. Se un’applicazione sembra lenta all’inizio e poi fluida, non è un bug: è l’architettura che lavora così.

La terza è sulla scelta del modello. Due modelli con la stessa qualità possono costare molto diverso in esercizio se uno usa GQA e l’altro no, perché a parità di scheda il primo serve più utenti. La lunghezza si paga due volte: una in token, una in memoria.

Domande frequenti

Che differenza c’è tra KV cache e memoria del modello?

I pesi sono fissi e si caricano una volta. La KV cache è temporanea, cresce con il testo della singola richiesta e si libera quando la richiesta finisce. Stessa memoria video, scopi diversi.

Perché non si mette in cache anche la query?

Perché la query serve solo al token che stai generando, una volta sola. Conservarla non farebbe risparmiare nulla: il risparmio sta nelle key e nelle value, che vengono rilette da tutti i token successivi.

La KV cache si può disattivare?

Sì e in addestramento è normale farlo perché il calcolo è già parallelo su tutta la sequenza. In generazione disattivarla significa rifare tutto ogni volta: più lento di ordini di grandezza.

Il prompt caching riduce la KV cache?

No. Riduce il calcolo del prefill e il prezzo dei token di input. La cache in memoria video serve comunque per generare la risposta. Sono due risparmi diversi, che si sommano.

Perché il primo token arriva lento?

Perché il prompt va processato tutto insieme prima di poterne generare l’erede. Con contesti lunghi quella fase dura secondi. Dopo, il decode procede un token per volta ed è molto più rapido.

La lunghezza del contesto ha un limite tecnico?

Sì, ed è la memoria. Ogni token in più nella cache costa spazio e riduce quanti utenti puoi servire insieme. I modelli con contesti enormi funzionano, ma il costo in hardware cresce in fretta.

Per approfondire: i Transformer, query, key e value, anatomia di un blocco e le leggi di scala.

Newsletter settimanale con le ultime novità sull’intelligenza artificiale

Corso gratuito sull’Intelligenza Artificiale, dalle nozioni base ai concetti più complessi

Condividi:

Potrebbero interessarti