Ultimo aggiornamento: 2026-06-22

La memoria dell'attenzione

L'attenzione standard dei transformer confronta ogni token con ogni altro token, costruendo una matrice dei punteggi prima della softmax. Quella matrice è quadrata nella lunghezza della sequenza, quindi la sua memoria cresce con il quadrato del contesto — il famoso costo quadratico dell'attenzione. Questo calcolatore stima la dimensione di quella matrice materializzata a partire dalla lunghezza della sequenza, dal numero di teste di attenzione, dalla dimensione del batch e dai byte usati per ogni punteggio. È precisamente la memoria che FlashAttention evita di memorizzare.

Il termine quadratico

Per una sequenza di $s$ token, l'attenzione forma una matrice di punteggi $s \times s$: la riga $i$, colonna $j$ contiene quanto il token $i$ si attende al token $j$. Memorizzarla richiede una memoria proporzionale a $s^2$. Questo è diverso dai pesi del modello o dalla cache chiave-valore, che crescono in modo lineare con il contesto — la matrice dei punteggi cresce con il suo quadrato, quindi con contesti lunghi diventa il buffer intermedio dominante e la ragione per cui l'attenzione ingenua esaurisce la memoria.

La formula

Ogni testa costruisce la propria matrice dei punteggi e ogni sequenza nel batch porta con sé la propria copia, quindi la memoria in byte è

dove $B$ è la dimensione del batch, $H$ il numero di teste di attenzione, $s$ la lunghezza della sequenza ed $e$ i byte per punteggio memorizzato. Dividendo per $10^9$ si ottengono i gigabyte. Il quadrato su $s$ è ciò che rende costoso il contesto lungo: il numero di teste e il batch moltiplicano solo in modo lineare.

Che cos'è questa memoria

Il valore misura la matrice dei punteggi materializzata che l'attenzione standard scrive in memoria tra il prodotto query-chiave e la somma sui valori pesata dalla softmax. È il buffer di attivazione che un'implementazione ingenua deve mantenere. Se persista o meno dipende dal contesto: in un semplice passaggio in avanti è transitorio e può essere riutilizzato tra gli strati, mentre i framework di addestramento possono conservare una copia per strato per il passaggio all'indietro, a meno che non intervengano il checkpointing o un kernel fuso.

FlashAttention

FlashAttention calcola lo stesso output senza mai scrivere la matrice $s \times s$ completa. Percorre le chiavi e i valori in piccoli blocchi, mantenendo statistiche della softmax aggiornate progressivamente, così da aver bisogno di un solo blocco in memoria alla volta, il che trasforma la memoria da quadratica a lineare nella lunghezza della sequenza. La matrice che questo calcolatore dimensiona è esattamente ciò che FlashAttention rinuncia a memorizzare — quindi il risultato qui riportato è una buona approssimazione della memoria che un kernel di attenzione fuso risparmia a una data lunghezza di contesto.

Esempio svolto

Consideriamo una singola sequenza di 4.096 token con 32 teste di attenzione in precisione a 16 bit:

circa 1,07 GB per la matrice dei punteggi di un singolo strato. Raddoppiando il contesto a 8.192 token, il termine quadratico prende il sopravvento: la stessa espressione dà circa 4,29 GB, quattro volte tanto per il doppio della lunghezza. Il lato chiave-valore dello stesso budget di attenzione è trattato nel Calcolatore della dimensione della cache KV, mentre la parte relativa ai pesi del modello nel Calcolatore della VRAM di inferenza per LLM.