PagedAttention

PagedAttention è un algoritmo di attention e un gestore di memoria per il serving di modelli linguistici di grandi dimensioni (LLM), che archivia la cache KV di ogni richiesta in piccoli blocchi di dimensione fissa, collocati ovunque nella memoria della GPU, e li tiene traccia tramite una tabella di blocchi per richiesta, come fa un sistema operativo con le pagine della memoria virtuale. Woosuk Kwon, Zhuohan Li e i loro colleghi dell’Università della California a Berkeley e di università partner l’hanno presentata insieme al motore di serving vLLM in Efficient Memory Management for Large Language Model Serving with PagedAttention, presentato a SOSP 2023. Eliminare la frammentazione delle allocazioni di cache contigue consente a una GPU di mantenere simultaneamente molte più sequenze; l’articolo riporta un throughput da 2 a 4 volte superiore rispetto a FasterTransformer e Orca a parità di latenza.

Il problema di memoria che risolve#

Durante la generazione di testo, un transformer conserva i vettori key e value di ogni token precedente, evitando di ricalcolarli. Questa cache cresce di una voce per token e per livello; una singola sequenza del modello LLaMA da 13 miliardi di parametri può occupare fino a 1,7 GB. La lunghezza dell’output non è nota quando arriva una richiesta, quindi i precedenti sistemi di serving riservavano per ogni richiesta un’area contigua dimensionata per la lunghezza massima possibile. Gran parte dell’area restava vuota e gli spazi tra le aree erano troppo piccoli per essere riutilizzati. Il team vLLM ha misurato che questi sistemi sprecavano dal 60% all’80% della memoria della cache KV, limitando la dimensione del batch molto prima che si esaurisse la capacità di calcolo.

Come funziona la paginazione#

PagedAttention suddivide la cache in blocchi, ognuno dei quali contiene le key e i value di un numero fisso di token: per impostazione predefinita, 16 in vLLM. Una richiesta riceve un nuovo blocco solo quando il precedente è pieno, scegliendolo da qualsiasi posizione libera; la tabella dei blocchi registra l’ordine. Il kernel del meccanismo di attention segue questa tabella per raccogliere key e value blocco per blocco, perciò la cache di una sequenza non deve mai essere contigua. Lo spreco è limitato alla parte inutilizzata dell’ultimo blocco di ogni sequenza: meno del 4% secondo le misurazioni di vLLM.

Questo esempio di codice PyTorch memorizza le key di una richiesta in blocchi sparsi all’interno di un pool condiviso e le rilegge tramite la relativa tabella dei blocchi:

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

Stampa la tabella dei blocchi [41, 7, 19] e keys intact: True, quindi 8 slot inutilizzati per la richiesta paginata, rispetto ai 2.008 di un’area riservata per 2.048 token.

Condivisione dei blocchi tra sequenze#

Le tabelle dei blocchi permettono anche a più sequenze di puntare agli stessi blocchi fisici. Il campionamento parallelo, la ricerca a fascio e le richieste che iniziano con lo stesso prompt di sistema riutilizzano una singola copia del prefisso condiviso, con un conteggio dei riferimenti per ogni blocco. Quando una sequenza deve scrivere in un blocco ancora utilizzato da altre sequenze, il motore lo copia prima: un’operazione di copia su scrittura mutuata dai sistemi operativi.

Gli autori di vLLM riferiscono che questa condivisione riduce fino al 55% l’uso di memoria del campionamento parallelo e della ricerca a fascio, consentendo un throughput fino a 2,2 volte superiore. L’hashing dei blocchi completi estende l’idea alle diverse richieste: è così che vLLM implementa la memorizzazione nella cache dei prompt.

Tecniche correlate#

PagedAttention determina dove vengono conservate le key e i value memorizzati nella cache e si combina con metodi che modificano la quantità da archiviare o la velocità di lettura:

• Flash Attention: suddivide il calcolo dell’attention in blocchi, evitando di scrivere l’intera matrice dei punteggi nella memoria principale della GPU. I motori attuali eseguono kernel in stile FlashAttention che leggono dai blocchi paginati.
• Batching continuo: ammette una nuova richiesta ogni volta che un’altra termina; la paginazione restituisce immediatamente i blocchi della sequenza completata al pool libero, lasciando spazio alla nuova richiesta.
• Attention a query raggruppate: condivide le teste key e value tra le teste query, riducendo il contenuto di ogni blocco.
• Quantizzazione della cache FP8: memorizza key e value con precisione a 8 bit, dimezzando all’incirca la memoria della cache rispetto ai formati a 16 bit.

L’interfaccia LLM di Ultralytics inoltra gli argomenti delle richieste senza modificarli, quindi una chiamata Chat Completions con n=3 inviata tramite base_url a un server vLLM locale restituisce tre campioni che condividono i blocchi memorizzati nella cache del prompt; solo le continuazioni divergenti richiedono nuova memoria.