Chunked Prefill

Le préremplissage par blocs est une technique de service de modèles qui divise un long prompt d’entrée en petits groupes de jetons et les traite en plusieurs étapes, entre lesquelles s’insère la génération de jetons pour d’autres requêtes. Sans cette technique, un moteur d’inférence qui accepte un long prompt consacre une étape entière à celui-ci, et toutes les réponses déjà diffusées aux autres utilisateurs s’interrompent jusqu’à ce qu’il soit terminé. L’article SARATHI de 2023 a introduit le préremplissage par blocs, puis son prolongement Sarathi-Serve a conçu une planification « sans interruption » à partir de cette approche. Dynamic SplitFuse de DeepSpeed-FastGen repose sur la même idée. Le préremplissage par blocs est désormais activé par défaut dans vLLM et disponible dans SGLang et TensorRT-LLM, les moteurs qui servent de nombreux grands modèles de langage (LLM) à poids ouverts.

Préremplissage, décodage et interruptions de génération#

L’inférence autorégressive par Transformer comporte deux phases :

• Préremplissage : le modèle traite en parallèle tous les jetons du prompt, calcule leurs clés et leurs valeurs, puis les stocke dans le cache KV. Le préremplissage effectue de nombreuses opérations par octet de poids lu ; il est donc généralement limité par le calcul.
• Décodage : le modèle génère un jeton par étape pour chaque requête et lit l’intégralité du cache à chaque fois. Le décodage effectue peu d’opérations arithmétiques par octet lu ; il est donc généralement limité par la bande passante mémoire.

Une étape GPU contenant un préremplissage de 20 000 jetons dure bien plus longtemps qu’une étape de décodage ; toutes les requêtes du lot doivent donc attendre. Sarathi-Serve appelle ces pauses des interruptions de génération. Elles se manifestent par des pics de latence entre les jetons pour les utilisateurs dont les propres prompts sont courts.

Fonctionnement du préremplissage par blocs#

Le planificateur attribue à chaque étape un budget de jetons, défini dans vLLM par max_num_batched_tokens. Il ajoute d’abord un jeton pour chaque décodage en cours, puis remplit le budget restant avec des blocs de prompts en attente. Un prompt qui ne tient pas dans le budget se poursuit à l’étape suivante : son cache KV s’allonge bloc après bloc jusqu’à ce que le prompt soit entièrement traité et que son décodage commence. Le mélange des deux phases dans une même étape améliore également l’utilisation du GPU, car les décodages laissent des capacités de calcul inoccupées que les blocs de préremplissage peuvent exploiter :

Le découpage ne change que le moment où les jetons du prompt sont traités. Chaque bloc utilise les clés et les valeurs mises en cache de tous les blocs précédents ainsi que ses propres jetons, sous un masque causal. Le résultat est donc identique à celui d’un préremplissage en une seule passe, comme le montre une couche d’attention dans 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)}")

Le programme affiche 4 chunks, same result as one pass: True : le prompt de 1 000 jetons est traité en blocs de 256, 256, 256 et 232 jetons, et les sorties correspondent à celles d’une seule passe.

Techniques de service apparentées#

• Traitement par lots continu : ajoute et retire des requêtes à chaque étape au lieu d’attendre la fin de tout un lot. Le préremplissage par blocs permet de planifier les longs prompts dans cette boucle.
• PagedAttention : stocke le cache KV dans des pages de taille fixe, afin qu’un prompt prérempli sur plusieurs étapes puisse étendre son cache sans réserver d’emblée la mémoire nécessaire à sa longueur totale.
• Mise en cache des prompts : réutilise les clés et les valeurs mises en cache d’un préfixe de prompt répété, ce qui élimine complètement le préremplissage de cette partie.
• Préremplissage et décodage désagrégés : exécute les deux phases sur des groupes de GPU distincts et transfère le cache KV entre eux, afin de pouvoir dimensionner et régler chaque groupe en fonction de son propre goulot d’étranglement.
• Décodage spéculatif : accélère la phase de décodage en proposant plusieurs jetons, puis en les vérifiant en une seule passe.

Applications#

Le préremplissage par blocs est surtout utile lorsque de longs et de courts prompts partagent un serveur. Un assistant de génération augmentée par récupération (RAG) qui remplit sa fenêtre de contexte avec un rapport de 50 pages peut le préremplir par blocs pendant que les conversations des autres utilisateurs continuent d’être diffusées. Les modèles vision-langage (VLM) ajoutent des centaines ou des milliers de jetons d’image à chaque prompt ; les pipelines de caméras qui envoient des images pour description produisent donc exactement ce mélange. L’interface LLM d’Ultralytics envoie du texte et des images à tout point de terminaison compatible avec OpenAI, y compris un serveur vLLM auto-hébergé, où les requêtes d’un pipeline YOLO26 visant à décrire les images dans lesquelles il détecte des personnes sont préremplies par blocs en parallèle du trafic des autres utilisateurs.

Régler le budget de jetons#

Le budget met en balance deux mesures de latence. Un petit budget laisse peu de place au préremplissage à chaque étape : il réduit la latence entre les jetons pour les utilisateurs qui reçoivent les réponses en continu et retarde le délai avant le premier jeton pour les nouveaux prompts. Un grand budget produit l’effet inverse. Chaque bloc suivant d’un long prompt relit également le cache KV grandissant des blocs précédents, de sorte que des blocs très petits entraînent davantage de trafic mémoire. Les opérateurs règlent le budget en fonction de la médiane et de la queue de la latence d’inférence, du délai avant le premier jeton et du débit, sur des longueurs de prompts réalistes. Le guide d’optimisation vLLM cite 2 048 jetons comme budget favorisant la latence entre les jetons et recommande plus de 8 192 jetons pour obtenir un débit maximal avec des modèles plus petits sur de grands GPU.