Batching continu

Le batch continu maximise l'utilisation du GPU en réorganisant dynamiquement le lot à chaque étape de génération. À mesure que les requêtes se terminent, de nouvelles requêtes rejoignent immédiatement le lot au lieu d'attendre la fin de l'ensemble. Le GPU reste saturé et le débit reste élevé.

Pour les déploiements en production, utilisez transformers serve. Il s’appuie sur ContinuousBatchingManager et expose un point de terminaison HTTP compatible OpenAI.

generate_batch

Le batch continu est pris en charge via generate_batch(). Transmettez une liste de prompts tokenisés et récupérez les résultats pour l’ensemble une fois leur traitement terminé. generate_batch gère la planification en interne et bloque jusqu’à l’achèvement de toutes les requêtes.

Pour les cas d’usage de service et de diffusion en continu, utilisez directement ContinuousBatchingManager pour gérer les requêtes.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import ContinuousBatchingConfig, GenerationConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B",
    attn_implementation="flash_attention_2",
    device_map="auto",
    dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")

prompts = [
    "What's up?",
    "Name a cat breed.",
    "Write a detailed history of quantum mechanics.",
]
inputs = [tokenizer.encode(p) for p in prompts]

generation_config = GenerationConfig(
    max_new_tokens=64,
    eos_token_id=tokenizer.eos_token_id,
)

outputs = model.generate_batch(inputs=inputs, generation_config=generation_config)

for request_id, output in outputs.items():
    text = tokenizer.decode(output.generated_tokens, skip_special_tokens=True)
    print(f"[{request_id}] {text}")

ContinuousBatchingManager

ContinuousBatchingManager exécute un thread en arrière-plan et vous permet de soumettre des requêtes et de récupérer les résultats de manière indépendante. À chaque étape de génération, il vérifie les requêtes terminées et planifie de nouvelles requêtes pour rejoindre le lot. Cela est utile pour le streaming, le service en temps réel ou la soumission de requêtes à mesure qu’elles arrivent.

Utilisez continuous_batching_context_manager() pour démarrer et arrêter le gestionnaire en toute sécurité. L’exemple ci-dessous contient des entrées de longueurs variables. Dès que le prompt le plus court est terminé, il quitte le lot tandis que les prompts plus longs continuent de générer. Avec un lot statique, vous devriez les rembourrer tous à la même longueur. Le lot continu libère le prompt terminé afin que vous puissiez commencer immédiatement le traitement du prompt suivant.

with model.continuous_batching_context_manager(generation_config=generation_config) as manager:
    manager.add_request(
        input_ids=tokenizer.encode("Write a detailed history of quantum mechanics."),
        request_id="long",
        max_new_tokens=512,
    )
    manager.add_request(
        input_ids=tokenizer.encode("What's up?"),
        request_id="short_0",
        max_new_tokens=32,
    )
    manager.add_request(
        input_ids=tokenizer.encode("Name a cat breed."),
        request_id="short_1",
        max_new_tokens=32,
    )

    for result in manager:
        text = tokenizer.decode(result.generated_tokens, skip_special_tokens=True)
        print(f"[{result.request_id}] {text}")

Vous pouvez également appeler init_continuous_batching() pour gérer le cycle de vie vous-même.

manager = model.init_continuous_batching(generation_config=generation_config)
manager.start()

# submit and retrieve requests...

Arrêt du gestionnaire

Le manager exécute un thread en arrière-plan et détient des ressources distribuées. L’arrêt s’effectue en deux étapes afin que vous puissiez choisir comment traiter le travail en cours.

Appel stop() pour arrêter le thread en arrière-plan. Par défaut, le gestionnaire cesse d’accepter de nouvelles soumissions et attend la fin des demandes en file d’attente et actives avant que le thread ne se termine.

manager.stop()

Passez hard_stop=True pour abandonner immédiatement le travail en attente. Les demandes en file d’attente et actives échouent avec une RuntimeError au lieu de se terminer.

manager.stop(hard_stop=True)

Une fois que stop est appelé, add_request() et add_requests() ignorent les nouvelles soumissions et affichent un avertissement. Vous pouvez toujours rappeler start pour exécuter une autre session de génération avec le même gestionnaire.

Appelez destroy() pour libérer les ressources distribuées. destroy arrête d’abord le gestionnaire s’il est encore en cours d’exécution, et le gestionnaire ne peut pas être redémarré par la suite. Utilisez-le lorsque vous avez terminé avec le lotissement continu pour la durée de vie du processus.

manager.destroy()

continuous_batching_context_manager() gère ce processus. Il appelle stop à la sortie et destroy, sauf si vous passez persistent_manager=True pour mettre le gestionnaire en cache sur le modèle pour la session suivante.

Ajout de requêtes

add_request() soumet une requête unique. Fournissez un request_id ou laissez le gestionnaire en générer un automatiquement.

manager.add_request(input_ids=input_ids, request_id="my_request")

add_requests() soumet un lot en une seule fois. Il trie automatiquement les entrées afin de maximiser les correspondances du cache de préfixe lorsque le partage de blocs est activé.

manager.add_requests(inputs=inputs)

Annuler une requête avec cancel_request().

manager.cancel_request(request_id="my_request")

Paramètres d'échantillonnage par requête

Activez per_request_processors pour appliquer temperature, top_k et top_p de manière indépendante pour chaque requête au sein d'un même passage avant, afin de permettre des paramètres d'échantillonnage différents pour différentes requêtes (par exemple, des sorties créatives à haute température par opposition à des sorties précises à basse température).

cb_config = ContinuousBatchingConfig(per_request_processors=True)

# each request gets its own sampling parameters
manager.add_request(input_ids=inputs_a, temperature=0.9, top_p=0.95)
manager.add_request(input_ids=inputs_b, temperature=0.1, top_k=10)

Chaque paramètre de GenerationConfig doit avoir une valeur non par défaut afin de créer le processeur de logits associé à l'exécution. Par exemple, définissez temperature sur une valeur autre que None ou 1 pour prendre en charge le contrôle de la température par requête. Les requêtes avec une température de 1 peuvent toujours être créées par la suite.

Récupération des résultats

Itérez sur le manager pour recevoir les résultats à mesure qu'ils arrivent.

for result in manager:
    print(tokenizer.decode(result.generated_tokens, skip_special_tokens=True))

get_result() récupère le résultat suivant de la file de sortie. Passez un request_id pour filtrer une requête spécifique. Si le résultat suivant dans la file ne correspond pas, il est remis en file et la méthode renvoie None.

# next available result
result = manager.get_result()

# filter for a specific request
result = manager.get_result(request_id="my_request")

Streaming

Définissez streaming=True sur une requête, puis utilisez request_id_iter() pour itérer sur les sorties partielles au fur et à mesure que les jetons sont générés.

from transformers.generation.continuous_batching import RequestStatus

manager.add_request(input_ids=input_ids, request_id="streamed", streaming=True)

for chunk in manager.request_id_iter(request_id="streamed"):
    token = tokenizer.decode(chunk.generated_tokens[-1:], skip_special_tokens=True)
    print(token, end="", flush=True)
    if chunk.status == RequestStatus.FINISHED:
        break

ContinuousBatchingConfig

ContinuousBatchingConfig contrôle le cache KV, l’ordonnancement, les graphes CUDA, l’utilisation de la mémoire et bien plus encore. Transmettez-le à generate_batch ou à init_continuous_batching() via l’argument continuous_batching_config, indépendamment de GenerationConfig. Les deux configurations décrivent des aspects distincts. GenerationConfig est centrée sur le modèle et contient les paramètres d’échantillonnage et d’arrêt, tandis que ContinuousBatchingConfig est centrée sur le matériel et contient les paramètres de mémoire et d’ordonnancement.

Définir continuous_batching_config sur un GenerationConfig est obsolète, émet un FutureWarning et sera supprimé dans la v5.19.

Par défaut, max_batch_tokens est 8192, limité par la mémoire GPU disponible et jamais inférieur à 256, tandis que num_blocks remplit la mémoire restante. Utilisez le tableau ci-dessous pour vous aider à choisir les fonctionnalités appropriées.

from transformers.generation import ContinuousBatchingConfig

cb_config = ContinuousBatchingConfig(
    max_memory_percent=0.8,
    page_size=256,           # KV cache page size in tokens
    scheduler_type="fifo",        # "fifo" or "prefill_first"
)

outputs = model.generate_batch(
    inputs=inputs,
    generation_config=generation_config,
    continuous_batching_config=cb_config,
)

Taille de page du cache KV

page_size définit le nombre de jetons de KV qu’une page d’attention complète contient pour une couche. Il doit être d’au moins 4, sinon l’initialisation lève une ValueError. La valeur par défaut 256 convient à la plupart des charges de travail.

Le paramètre block_size est obsolète et remplacé par page_size.

Une taille de page plus importante réduit le suivi mais gaspille de l’espace lorsqu’une séquence s’arrête en milieu de page. Des pages plus petites gaspillent moins d’espace mais nécessitent un suivi plus important. Vous n’avez généralement pas besoin d’ajuster page_size. Si votre charge de travail comporte de nombreuses requêtes courtes, essayez une valeur plus petite. S’il comporte quelques requêtes très longues, essayez une valeur plus grande. Consultez Gestion de la mémoire pour comprendre comment les pages, les blocs et les fragments s’articulent.

cb_config = ContinuousBatchingConfig(page_size=128)

Taille du lot de préremplissage

max_batch_tokens définit le budget de jetons pour un seul passage avant, c’est-à-dire le nombre maximal de jetons de requête que le modèle traite simultanément. Un budget plus élevé permet d’insérer davantage de jetons de prompt dans chaque préremplissage, ce qui augmente le débit de préremplissage et réduit le temps jusqu’au premier jeton sur les charges de travail riches en prompts. L’ordonnanceur répartit les prompts qui dépassent le budget sur plusieurs étapes. Consultez chunked prefill pour savoir comment les prompts trop volumineux sont gérés.

Par défaut, max_batch_tokens est 8192. La valeur est limitée par la mémoire GPU disponible afin que les tenseurs d’entrée par lot n’évincent pas le cache KV, et elle ne descend jamais en dessous de 256. Sur un GPU avec peu de mémoire libre, la limite l’abaisse en dessous de 8192.

max_batch_tokens et num_blocks puisent dans le même budget mémoire, ils sont donc en concurrence l’un avec l’autre. Un budget de jetons plus important alloue des tampons d’entrée plus volumineux et laisse moins de blocs de cache KV pour les requêtes concurrentes ou longues. Augmentez max_batch_tokens pour accroître le débit de préremplissage si vous disposez de plus de mémoire, et réduisez-le pour libérer de la mémoire afin d’ajouter des blocs KV ou pour éviter les erreurs de mémoire insuffisante.

cb_config = ContinuousBatchingConfig(max_batch_tokens=16384)

Limites de lot et de planification

max_requests_per_batch limite le nombre de requêtes exécutées dans un seul passage avant. Le modèle produit un tenseur de logits dimensionné par le nombre de jetons du lot et la taille du vocabulaire, et il est converti en fp32 pour l’échantillonnage, ce qui double son empreinte mémoire. Chaque requête n’a besoin que d’une prédiction de jeton suivant, donc une limite plus faible maintient ce tenseur plus petit et évite les erreurs de mémoire insuffisante sur les lots dominés par la préremplissage avec de grands vocabulaires. Par défaut, il est défini au nombre de prompts soumis, avec une valeur de repli de 1024, puis limité pour s’adapter à max_batch_tokens et num_blocks.

cb_config = ContinuousBatchingConfig(max_requests_per_batch=256)

safety_margin réserve une fraction de la capacité libre du cache pour les requêtes actives. Lorsque la capacité libre passe sous le seuil de safety_margin (une fraction de la taille totale du cache), le planificateur cesse d’admettre de nouveaux prefills et ne poursuit que le décodage des requêtes déjà en cours. Cette approche privilégie l’achèvement du travail actif par rapport au démarrage de nouvelles requêtes, ce qui protège la latence de décodage et retarde le offloading. La valeur doit se situer entre 0 et 1, où 0 désactive la marge.

cb_config = ContinuousBatchingConfig(safety_margin=0.15)

La valeur par défaut dépend de l’ordonnanceur. Pour l’ordonnanceur FIFO, elle est de 0.15 et pour prefill_first, elle est de 0.0. Consultez Admission pour comprendre comment la marge interagit avec l’ordonnancement.

Probabilités logarithmiques

ContinuousBatchingConfig renvoie la log-probabilité de chaque jeton généré lorsque return_logprobs=True. Cela est utile pour le RL, où les logprobs constituent une entrée de certaines boucles d’entraînement.

cb_config = ContinuousBatchingConfig(return_logprobs=True)

# generate_batch()

for request_id, output in outputs.items():
    for token_id, log_prob in zip(output.generated_tokens, output.logprobs):
        token = tokenizer.decode([token_id])
        print(f"{token} | logprob: {log_prob}")

Graphes CUDA

Les graphes CUDA éliminent le surcoût de dispatch CPU en enregistrant le graphe d’exécution GPU une seule fois, puis en le rejouant pour des lots dont les formes correspondent. Activez-les explicitement avec use_cuda_graph=True.

cb_config = ContinuousBatchingConfig(use_cuda_graph=True)

Lorsqu’il est actif, le gestionnaire remplit les longueurs des requêtes et des paires clé-valeur jusqu’à des intervalles fixes afin que les formes se répètent et que les graphes soient réutilisés. Des valeurs plus petites de q_padding_interval_size et kv_padding_interval_size réduisent le calcul gaspillé sur le remplissage, mais cela signifie qu’il y a plus de formes uniques que le graphe doit enregistrer et stocker, ce qui coûte davantage de mémoire.

cb_config = ContinuousBatchingConfig(
    use_cuda_graph=True,
    q_padding_interval_size=64,
    kv_padding_interval_size=16384,
)

Batching asynchrone

Le lot asynchrone superpose la planification CPU du lot suivant au calcul GPU du lot actuel. Il nécessite des graphes CUDA et double approximativement la VRAM utilisée pour les tenseurs d'entrée.

cb_config = ContinuousBatchingConfig(
    use_cuda_graph=True,
    use_async_batching=True,
)

Compilation

default_compile_level applique torch.compile aux passes avant du modèle. La compilation échange un coût de mise en route unique contre des passes avant plus rapides lors de la génération. Des niveaux plus élevés exécutent des optimisations plus agressives, ce qui améliore le débit mais allonge la mise en route. Maintenez le niveau bas pour les tests et l’itération des benchmarks, et augmentez-le pour les charges de travail de service de longue durée où le coût de mise en route est amorti sur de nombreuses requêtes.

Le niveau varie de 0 à 3. Le niveau 0 est la valeur par défaut et ignore entièrement la compilation.

cb_config = ContinuousBatchingConfig(default_compile_level=1)

Le niveau fournit une configuration CompileConfig par défaut aux chemins d'exécution varlen et decode. Il ne s'applique qu'à un chemin qui n'a pas de configuration explicite, de sorte que varlen_compile_config et decode_compile_config ont la priorité lorsqu'ils sont définis. Sous FlashAttention, le chemin varlen ignore la compilation car max_length_k déclenche des recompilations fréquentes, si bien que le niveau n'affecte que le chemin decode dans ce cas.

Chemin rapide de décodage

Lorsqu’un lot ne contient que des requêtes de décodage (un jeton de requête par séquence), le gestionnaire peut dispatcher vers le noyau flash_attn_with_kvcache au lieu du noyau à longueur variable. C’est plus rapide que le chemin varlen car le noyau lit et écrit le cache KV paginé sur place via un tableau de blocs, plutôt que de passer par une mise à jour manuelle. Consultez Paged attention pour les détails au niveau du noyau.

Le chemin rapide est dimensionné par max_blocks_per_request, qui détermine la taille de la table de blocs par requête. Par défaut, cette valeur est inférée automatiquement. Si max_prompt_length et max_generated_length sont définis sur le gestionnaire, la table de blocs est dimensionnée pour accueillir la longueur de séquence maximale. Sinon, une valeur par défaut de repli (32 blocs par requête) est utilisée.

Définissez max_blocks_per_request sur une valeur spécifique pour dimensionner explicitement la table de blocs. Cela est utile lorsque vous connaissez la longueur de séquence maximale par requête et souhaitez limiter le coût mémoire de la table de blocs.

cb_config = ContinuousBatchingConfig(max_blocks_per_request=64)

Définissez max_blocks_per_request=0 pour désactiver le chemin rapide et forcer le passage de chaque lot par le noyau varlen. Cela restaure le comportement par défaut antérieur et est utile lorsque le chemin rapide n’est pas disponible pour votre implémentation d’attention (le gestionnaire le désactive également automatiquement lorsque le noyau sous-jacent ne peut pas être utilisé).

cb_config = ContinuousBatchingConfig(max_blocks_per_request=0)

Le chemin rapide repose sur le noyau flash_attn_with_kvcache, qui est disponible pour deux combinaisons de périphérique et d’implémentation de l’attention.

Pour toute autre combinaison, ou lorsque le noyau ne peut pas être importé, le gestionnaire bascule sur le chemin varlen. Il n’affiche un avertissement que lorsque vous définissez explicitement max_blocks_per_request.

La pleine attention prend en charge les tables de blocs de noyau, mais l’attention glissante ne le fait pas. Si le modèle comporte des couches d’attention glissante, y compris les modèles hybrides qui utilisent également la pleine attention, l’initialisation force max_blocks_per_request à 0, et toute valeur non nulle définie est écrasée. Le gestionnaire enregistre alors que le chemin rapide de décodage n’est pas disponible car certains types d’attention du modèle ne prennent pas en charge les tables de blocs de noyau.

Déchargement sur CPU

Le déchargement sur CPU copie les blocs de cache KV évacués vers un tampon CPU épinglé préalloué lorsque le cache KV du GPU est plein. Une fois de l’espace disponible dans le cache, le gestionnaire recopie les blocs sur le GPU et reprend la requête sans recalculer son prompt ni les jetons générés.

Définissez cpu_offload_space sur l'espace d'échange CPU en GiB. La valeur par défaut, 0.0, désactive le déchargement sur CPU.

cb_config = ContinuousBatchingConfig(cpu_offload_space=8.0)

Par défaut, cpu_offload_space_safety_threshold=0.8 limite l'espace demandé à 80 % de la RAM système disponible lorsque psutil est installé. Définissez cpu_offload_space=None pour dimensionner le pool d'échange à partir du seuil de sécurité.

Délai d'attente du parallélisme de tenseurs

Sous le parallélisme de tenseurs, le gestionnaire crée un groupe de communication CPU pour coordonner les soumissions, les annulations et l’arrêt des requêtes entre les rangs. cpu_group_timeout limite la durée pendant laquelle une opération collective sur ce groupe peut bloquer avant que le processus ne plante. Si un rang se fige, le délai d’attente empêche les autres d’attendre indéfiniment.

Définissez un délai d’attente plus long pour les charges de travail qui émettent des collectives peu fréquentes, ou passez None pour le désactiver.

cb_config = ContinuousBatchingConfig(cpu_group_timeout=600.0)

Mise en cache des préfixes

Lorsque plusieurs requêtes partagent un préfixe commun, comme un prompt système, le gestionnaire réutilise leurs blocs de cache KV au lieu de les recalculer. allow_block_sharing est défini par défaut sur True. Le partage de préfixe n’est actif que si tous les allocateurs de type d’attention prennent en charge le partage, si bien que seuls les modèles à attention complète l’utilisent. Les modèles purement glissants et hybrides le désactivent même lorsque allow_block_sharing est True.

cb_config = ContinuousBatchingConfig(
    allow_block_sharing=True,  # default
)

Choisir un backend d'attention

Le batch continu prend en charge FlashAttention, SDPA et "paged|eager". Définissez attn_implementation lors du chargement du modèle.

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B",
    attn_implementation="flash_attention_2",
    device_map="auto",
    dtype=torch.bfloat16,
)

Le batch continu s’exécute plus rapidement avec FlashAttention qu’avec eager ou SDPA, principalement parce que FlashAttention n’a pas besoin de masque d’attention. Si le modèle prend en charge FlashAttention, un modèle chargé avec "eager" ou "sdpa" bascule sur FlashAttention 3, ou sur FlashAttention 2 si FlashAttention 3 n’est pas installé, et affiche un avertissement. Si aucune des deux versions n’est installée, le backend ne change pas.

Pour conserver eager ou SDPA, par exemple lors de tests, définissez auto_switch_to_flash=False dans ContinuousBatchingConfig. Chaque fois que le backend reste "eager", le continuous batching exécute "paged|eager" à la place.

Parallélisme de tenseurs

Pour les modèles trop volumineux pour tenir sur un seul GPU, répartissez les poids sur plusieurs appareils au moyen du parallélisme de tenseurs. Définissez le nombre d’appareils avec DistributedConfig(tp_size=N). Le lotissement continu lit la taille du parallélisme de tenseurs depuis le modèle afin de dimensionner le cache KV paginé par fragment. Consultez Tensor parallelism pour la liste des architectures prises en charge et le fonctionnement de la répartition.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig
from transformers.generation import ContinuousBatchingConfig, GenerationConfig

distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-32B",
    attn_implementation="flash_attention_2",
    distributed_config=distributed_config,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-32B")

inputs = [tokenizer.encode(p) for p in ["What's up?", "Name a cat breed."]]
generation_config = GenerationConfig(max_new_tokens=64, eos_token_id=tokenizer.eos_token_id)

outputs = model.generate_batch(inputs=inputs, generation_config=generation_config)

Lancez le script avec torchrun en définissant --nproc-per-node sur le nombre de GPU sur lesquels vous souhaitez répartir le sharding.

torchrun --nproc-per-node 4 cb_tp.py

La taille du parallélisme des tenseurs doit diviser le num_key_value_heads du modèle (vérifiez la configuration du modèle). Sinon, le cache pagelé lève une erreur au démarrage, il faut donc choisir une valeur appropriée pour --nproc-per-node.

Ne définissez pas device_map avec distributed_config. Les deux sont en conflit car device_map place des modules entiers sur des GPU spécifiques, tandis que le parallélisme de tenseurs répartit ces mêmes paramètres sur tous les GPU.

Attention à fenêtre glissante

Les modèles avec attention à fenêtre glissante (Mistral, Gemma 2) fonctionnent avec le lotissement continu. Les couches glissantes utilisent SlidingAttentionCacheAllocator, qui ne conserve qu’une fenêtre fixe de jetons récents en recyclant un petit ensemble de pages. Le type d’attention provient des layer_types ou de sliding_window du modèle, et non de ContinuousBatchingConfig.

from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("google/gemma-2-2b")
config.sliding_window = 4096

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-2b",
    config=config,
    attn_implementation="sdpa",
    device_map="auto",
    dtype=torch.bfloat16,
)

Prefix caching et le decode fast path sont désactivés dès lors que l’attention glissante est présente, y compris dans les modèles hybrides qui comportent également des couches d’attention complète.

Étapes suivantes

• L’article de blog Continuous batching aborde la mise en cache KV, le préremplissage par morceaux et la planification dynamique, avec des chiffres de référence de performance.
• Pour une analyse plus approfondie du fonctionnement du système de batch continu, consultez le document Architecture du batch continu.