Chunked Prefill

Il prefill suddiviso in blocchi è una tecnica di distribuzione dei modelli che divide un prompt di input lungo in gruppi più piccoli di token e li elabora in più passaggi, alternandoli alla generazione di token per altre richieste. Senza questa tecnica, un motore di inferenza che accetta un prompt lungo dedica un intero passaggio a quel prompt, e tutte le risposte già in streaming per gli altri utenti si fermano finché l’elaborazione non termina. L’articolo SARATHI del 2023 ha introdotto i prefill a blocchi; il successivo Sarathi-Serve ha costruito su di essi una pianificazione «senza arresti». Dynamic SplitFuse di DeepSpeed-FastGen usa la stessa idea; oggi il prefill a blocchi è abilitato per impostazione predefinita in vLLM e disponibile in SGLang e TensorRT-LLM, i motori che servono molti modelli linguistici di grandi dimensioni (LLM) con pesi aperti.

Prefill, decodifica e interruzioni della generazione#

L’inferenza autoregressiva dei Transformer ha due fasi:

• Prefill: il modello elabora in parallelo tutti i token del prompt, calcola le rispettive chiavi e valori e li memorizza nella cache KV. Il prefill esegue molte operazioni per ogni byte di pesi letto, quindi in genere è limitato dalla capacità di calcolo.
• Decodifica: a ogni passaggio, il modello genera un token per ciascuna richiesta, leggendo ogni volta l’intera cache. La decodifica esegue pochi calcoli per byte letto, quindi in genere è limitata dalla larghezza di banda della memoria.

Un passaggio su una GPU che contiene un prefill di 20.000 token richiede molto più tempo di un passaggio di decodifica, quindi tutte le richieste del batch devono aspettare. Sarathi-Serve chiama queste pause interruzioni della generazione; per gli utenti con prompt brevi, si manifestano come picchi nella latenza tra token.

Come funziona il prefill suddiviso in blocchi#

Lo scheduler assegna a ogni passaggio un budget di token, impostato in vLLM tramite max_num_batched_tokens. Aggiunge prima un token per ogni decodifica in corso, poi riempie il budget rimanente con blocchi dei prompt in attesa. Se un prompt non entra nel budget, continua nel passaggio successivo, estendendo la propria cache KV blocco dopo blocco finché non è stato elaborato per intero e può iniziare la decodifica. Mescolare le due fasi in un unico passaggio consente anche di sfruttare meglio la GPU, perché i blocchi di prefill possono utilizzare la capacità di calcolo lasciata libera dalla decodifica:

La suddivisione in blocchi modifica soltanto il momento in cui vengono elaborati i token del prompt. Ogni blocco presta attenzione alle chiavi e ai valori memorizzati nella cache di tutti i blocchi precedenti, oltre che ai propri token, applicando una maschera causale; il risultato equivale quindi a un prefill in un solo passaggio, come mostra un livello di attenzione in PyTorch:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
prompt_len, chunk = 1000, 256
q, k, v = torch.randn(3, 1, 8, prompt_len, 64).unbind(0)  # one layer: 8 heads, 64 dims each

# One-shot prefill: all prompt tokens attend causally in a single pass
full = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# Chunked prefill: each chunk attends to the cached keys and values of earlier chunks plus itself
outputs = []
for start in range(0, prompt_len, chunk):
    end = min(start + chunk, prompt_len)
    mask = torch.ones(end - start, end, dtype=torch.bool).tril(diagonal=start)  # causal over absolute positions
    q_chunk, k_cached, v_cached = q[..., start:end, :], k[..., :end, :], v[..., :end, :]
    outputs.append(F.scaled_dot_product_attention(q_chunk, k_cached, v_cached, attn_mask=mask))
chunked = torch.cat(outputs, dim=2)

print(f"{len(outputs)} chunks, same result as one pass: {torch.allclose(full, chunked, atol=1e-5)}")

Stampa 4 chunks, same result as one pass: True: il prompt di 1.000 token viene elaborato in blocchi di 256, 256, 256 e 232 token e gli output corrispondono a quelli del passaggio singolo.

Tecniche di distribuzione correlate#

• Batching continuo: aggiunge e rimuove richieste a ogni passaggio, senza aspettare che termini un intero batch. Il prefill a blocchi consente di pianificare i prompt lunghi all’interno di questo ciclo.
• PagedAttention: archivia la cache KV in pagine di dimensione fissa, così un prompt elaborato in più passaggi può ampliare la propria cache senza riservare in anticipo memoria per tutta la sua lunghezza.
• Memorizzazione nella cache dei prompt: riutilizza chiavi e valori memorizzati nella cache per un prefisso di prompt ripetuto, eliminando del tutto il prefill di quella parte.
• Prefill e decodifica disaggregati: esegue le due fasi su pool di GPU separati e trasferisce la cache KV tra loro, così ciascun pool può essere dimensionato e ottimizzato per il proprio collo di bottiglia.
• Decodifica speculativa: accelera la fase di decodifica generando in bozza più token e verificandoli in un unico passaggio.

Applicazioni#

Il prefill a blocchi è particolarmente utile quando prompt lunghi e brevi condividono lo stesso server. Un assistente di generazione aumentata dal recupero (RAG) che riempie la propria finestra di contesto con un rapporto di 50 pagine può elaborarlo in blocchi mentre le chat degli altri utenti continuano a essere trasmesse in streaming. I modelli visione-linguaggio (VLM) aggiungono centinaia o migliaia di token immagine a ogni prompt, perciò le pipeline di telecamere che inviano fotogrammi da descrivere producono proprio questo tipo di carico misto. L’interfaccia LLM di Ultralytics invia testo e immagini a qualsiasi endpoint compatibile con OpenAI, compreso un server vLLM self-hosted, dove le richieste di una pipeline YOLO26 per descrivere i fotogrammi in cui rileva persone vengono elaborate a blocchi insieme al traffico degli altri utenti.

Ottimizzazione del budget di token#

Il budget comporta un compromesso tra due metriche di latenza. Un budget ridotto lascia poco spazio al prefill a ogni passaggio, mantenendo bassa la latenza tra token per gli utenti in streaming e ritardando il tempo al primo token per i nuovi prompt; un budget elevato produce l’effetto opposto. Ogni blocco successivo di un prompt lungo rilegge inoltre la cache KV in crescita dei blocchi precedenti, perciò blocchi molto piccoli generano traffico aggiuntivo in memoria. Gli operatori ottimizzano il budget valutando la latenza di inferenza mediana e di coda, il tempo al primo token e la velocità di elaborazione con lunghezze di prompt realistiche. La guida all’ottimizzazione di vLLM indica 2.048 token come budget che privilegia la latenza tra token e consiglia più di 8.192 token per ottenere la massima velocità di elaborazione con modelli più piccoli su GPU di grandi dimensioni.