Loteamento contínuo

O continuous batching maximiza a utilização da GPU ao reprogramar dinamicamente o lote a cada etapa de geração. À medida que as solicitações são concluídas, novas solicitações são adicionadas imediatamente, em vez de aguardar a conclusão de todo o lote. A GPU permanece totalmente ocupada e a taxa de transferência permanece alta.

Para implantações em produção, use transformers serve. Ele é construído sobre o ContinuousBatchingManager e expõe um endpoint HTTP compatível com a OpenAI.

generate_batch

O loteamento contínuo é suportado por meio de generate_batch(). Passe uma lista de prompts tokenizados e receba os resultados de todos eles quando estiverem concluídos. generate_batch gerencia o agendamento internamente e bloqueia até que todas as solicitações sejam concluídas.

Para casos de uso de atendimento e streaming, use ContinuousBatchingManager diretamente para gerenciar as solicitações.

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 executa uma thread em segundo plano e permite enviar solicitações e recuperar resultados de forma independente. A cada etapa de geração, ele verifica as solicitações finalizadas e agenda novas para se juntarem ao lote. Isso é útil para streaming, atendimento em tempo real ou para enviar solicitações conforme chegam.

Use continuous_batching_context_manager() para iniciar e encerrar o gerenciador com segurança. O exemplo abaixo contém entradas de comprimento variável. Assim que o prompt mais curto for concluído, ele sai do lote enquanto os prompts mais longos continuam a gerar. Com o loteamento estático, você teria que preencher todos até o mesmo comprimento. O loteamento contínuo libera o prompt concluído, permitindo que você comece a processar o próximo prompt imediatamente.

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}")

Você também pode chamar init_continuous_batching() para gerenciar o ciclo de vida por conta própria.

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

# submit and retrieve requests...

Encerrando o gerenciador

O manager executa uma thread em segundo plano e mantém recursos distribuídos. O encerramento ocorre em duas etapas, permitindo que você escolha o que fazer com o trabalho em andamento.

Chame stop() para interromper a thread em segundo plano. Por padrão, o gerenciador para de aceitar novos envios e aguarda a conclusão das solicitações na fila e das solicitações ativas antes de a thread ser encerrada.

manager.stop()

Passe hard_stop=True para abandonar o trabalho pendente imediatamente. As solicitações enfileiradas e ativas falham com um RuntimeError em vez de serem concluídas.

manager.stop(hard_stop=True)

Quando stop é chamado, add_request() e add_requests() descartam novos envios e registram um aviso. Você ainda pode chamar start novamente para executar outra sessão de geração com o mesmo gerenciador.

Chame destroy() para liberar os recursos distribuídos. destroy primeiro interrompe o gerenciador, caso ele ainda esteja em execução, e o gerenciador não pode ser reiniciado em seguida. Use-o quando terminar de usar o loteamento contínuo para o tempo de vida do processo.

manager.destroy()

continuous_batching_context_manager() gerencia esse processo. Ele chama stop na saída e destroy, a menos que você passe persistent_manager=True para fazer o cache do gerenciador no modelo para a próxima sessão.

Adicionando solicitações

add_request() envia uma única solicitação. Forneça um request_id ou deixe que o gerador crie um automaticamente.

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

add_requests() envia um lote de uma só vez. Ele ordena as entradas automaticamente para maximizar os acertos no cache de prefixo quando o compartilhamento de blocos está habilitado.

manager.add_requests(inputs=inputs)

Cancele uma solicitação com cancel_request().

manager.cancel_request(request_id="my_request")

Parâmetros de amostragem por solicitação

Ative per_request_processors para aplicar temperature, top_k e top_p de forma independente para cada solicitação dentro da mesma passagem de frente, permitindo parâmetros de amostragem diferentes para solicitações distintas (por exemplo, saídas criativas com temperatura alta versus saídas precisas com temperatura baixa).

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)

Cada parâmetro em GenerationConfig deve ter um valor não padrão para que o logits processor associado seja criado em tempo de execução. Por exemplo, defina temperature com um valor diferente de None ou 1 para suportar controle de temperatura por solicitação. Solicitações com temperatura 1 ainda podem ser criadas posteriormente.

Recuperando resultados

Itere sobre o manager para receber os resultados conforme eles chegam.

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

get_result() recupera o próximo resultado da fila de saída. Passe um request_id para filtrar uma solicitação específica. Se o próximo resultado na fila não corresponder, ele é reencaminhado para a fila e o método retorna None.

# next available result
result = manager.get_result()

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

Streaming

Defina streaming=True em uma solicitação e, em seguida, use request_id_iter() para iterar sobre as saídas parciais à medida que os tokens são gerados.

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 controla o KV cache, o agendamento, os CUDA graphs, o uso de memória e muito mais. Passe-o para generate_batch ou init_continuous_batching() com o argumento continuous_batching_config, separadamente de GenerationConfig. As duas configurações descrevem coisas diferentes. GenerationConfig é centrada no modelo e contém parâmetros de amostragem e de parada, enquanto ContinuousBatchingConfig é centrada no hardware e contém parâmetros de memória e de agendamento.

Definir continuous_batching_config em um GenerationConfig está obsoleto, emite um FutureWarning e será removido na v5.19.

Por padrão, max_batch_tokens é 8192, limitado pela memória de GPU disponível e nunca inferior a 256, enquanto num_blocks preenche a memória restante. Use a tabela abaixo para ajudá-lo a escolher os recursos apropriados.

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,
)

Tamanho da página do cache KV

page_size define quantos tokens de KV uma página de atenção completa armazena para uma camada. Ele deve ser no mínimo 4, caso contrário a inicialização levanta um ValueError. O valor padrão 256 funciona para a maioria das cargas de trabalho.

O parâmetro block_size está obsoleto e foi substituído por page_size.

Um tamanho de página maior reduz o controle administrativo, mas desperdiça espaço quando uma sequência termina no meio da página. Páginas menores desperdiçam menos espaço, mas exigem mais controle administrativo. Geralmente, não é necessário ajustar page_size. Se sua carga de trabalho tiver muitas solicitações curtas, tente um valor menor. Se tiver algumas solicitações muito longas, tente um valor maior. Consulte Gestão de memória para entender como páginas, blocos e fragmentos se encaixam.

cb_config = ContinuousBatchingConfig(page_size=128)

Tamanho do lote de prefill

max_batch_tokens define o orçamento de tokens para uma única passagem de frente, ou seja, o número máximo de tokens de consulta que o modelo processa de uma vez. Um orçamento maior agrupa mais tokens de prompt em cada pré-preenchimento, o que aumenta a taxa de transferência do pré-preenchimento e reduz o tempo até o primeiro token em cargas de trabalho com muitos prompts. O agendador divide prompts que excedem o orçamento entre as etapas. Consulte chunked prefill para saber como prompts de tamanho excessivo são tratados.

Por padrão, max_batch_tokens é 8192. O valor é limitado pela memória de GPU disponível para que os tensors de entrada por lote não comprometam o cache KV, e nunca fica abaixo de 256. Em uma GPU com pouca memória livre, o limite reduz o valor para menos de 8192.

max_batch_tokens e num_blocks utilizam o mesmo orçamento de memória, portanto, há um compromisso entre eles. Um orçamento de tokens maior aloc buffers de entrada maiores e deixa menos blocos de cache KV para solicitações concorrentes ou longas. Aumente max_batch_tokens para aumentar o throughput de prefill quando houver mais memória disponível, e reduza-o para liberar memória para mais blocos KV ou para evitar erros de memória insuficiente.

cb_config = ContinuousBatchingConfig(max_batch_tokens=16384)

Limites de lote e agendamento

max_requests_per_batch limita quantas solicitações são executadas em uma única passagem de frente. O modelo produz um tensor de logits dimensionado pelo número de tokens do lote e pelo tamanho do vocabulário, e ele é convertido para fp32 para a amostragem, o que dobra a pegada de memória. Cada solicitação precisa apenas de uma previsão de próximo token, portanto, um limite menor mantém esse tensor menor e evita erros de falta de memória em lotes com predominância de pré-preenchimento e vocabulários grandes. Por padrão, ele é definido como o número de prompts enviados, com um valor de reserva de 1024, e em seguida limitado para caber em max_batch_tokens e num_blocks.

cb_config = ContinuousBatchingConfig(max_requests_per_batch=256)

safety_margin reserva uma fração da capacidade livre do cache para as solicitações ativas. Quando a capacidade livre fica abaixo de safety_margin (uma fração do tamanho total do cache), o agendador para de admitir novos prefills e continua apenas o decodificação das solicitações já em andamento. Isso prioriza a conclusão do trabalho ativo em vez de iniciar novas solicitações, o que protege a latência de decodificação e adia o offloading. O valor deve estar entre 0 e 1, sendo que 0 desativa a margem.

cb_config = ContinuousBatchingConfig(safety_margin=0.15)

O valor padrão depende do agendador. Para o agendador FIFO, é 0.15 e para prefill_first, é 0.0. Consulte Admissão para saber como a margem interage com o agendamento.

Probabilidades de log

ContinuousBatchingConfig retorna a log probabilidade de cada token gerado quando return_logprobs=True. Isso é útil para RL, em que as logprobs são uma entrada para alguns dos loops de treinamento.

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}")

Gráficos CUDA

Os gráficos CUDA eliminam a sobrecarga de despacho da CPU gravando o gráfico de execução da GPU uma única vez e reproduzindo-o para lotes com formas compatíveis. Ative-os explicitamente com use_cuda_graph=True.

cb_config = ContinuousBatchingConfig(use_cuda_graph=True)

Quando ativo, o manager preenche os comprimentos de query e KV até intervalos fixos, de modo que as formas se repitam e os grafos sejam reutilizados. Valores menores de q_padding_interval_size e kv_padding_interval_size reduzem o desperdício de computação no padding, mas isso significa que há mais formas únicas que o grafo precisa gravar e armazenar, o que consome mais memória.

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

Loteamento assíncrono

O loteamento assíncrono sobrepõe o agendamento de CPU do próximo lote à computação de GPU do lote atual. Ele requer CUDA graphs e praticamente dobra a VRAM usada para os tensores de entrada.

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

Compilação

default_compile_level aplica torch.compile às passadas de forward do modelo. A compilação troca um custo de aquecimento único por passadas de forward mais rápidas durante a geração. Níveis mais altos executam otimizações mais agressivas, o que melhora o throughput, mas prolonga o aquecimento. Mantenha o nível baixo para testes e iteração de benchmarks, e aumente-o para cargas de trabalho de atendimento de longa duração, em que o custo de aquecimento é compensado ao longo de muitas requisições.

O nível varia de 0 a 3. O nível 0 é o padrão e ignora a compilação por completo.

cb_config = ContinuousBatchingConfig(default_compile_level=1)

O nível fornece um CompileConfig padrão para os caminhos de execução varlen e decode. Ele se aplica apenas a um caminho que não possui configuração explícita, portanto, varlen_compile_config e decode_compile_config têm precedência quando definidos. Sob FlashAttention, o caminho varlen ignora a compilação porque max_length_k provoca recompilações frequentes, de modo que, nesse caso, o nível afeta apenas o caminho decode.

Caminho rápido de decodificação

Quando um lote contém apenas solicitações de decodificação (um token de consulta por sequência), o gerenciador pode despachar para o kernel flash_attn_with_kvcache em vez do kernel de comprimento variável. Isso é mais rápido que o caminho varlen porque o kernel lê e grava o cache KV paginado in-place por meio de uma tabela de blocos, em vez de passar por uma atualização manual. Consulte Paged attention para detalhes no nível do kernel.

O caminho rápido é dimensionado por max_blocks_per_request, que dimensiona a tabela de blocos por solicitação. Por padrão, isso é inferido automaticamente. Se max_prompt_length e max_generated_length estiverem definidos no gerenciador, a tabela de blocos será dimensionada para acomodar o comprimento máximo da sequência. Caso contrário, um valor padrão de fallback (32 blocos por solicitação) é utilizado.

Defina max_blocks_per_request com um valor específico para dimensionar explicitamente a tabela de blocos. Isso é útil quando você conhece o comprimento máximo da sequência por solicitação e deseja limitar o custo de memória da tabela de blocos.

cb_config = ContinuousBatchingConfig(max_blocks_per_request=64)

Defina max_blocks_per_request=0 para desativar o caminho rápido e forçar cada lote a passar pelo kernel varlen. Isso recupera o comportamento padrão anterior e é útil quando o caminho rápido não está disponível para sua implementação de atenção (o gerenciador também o desativa automaticamente quando o kernel subjacente não pode ser usado).

cb_config = ContinuousBatchingConfig(max_blocks_per_request=0)

O caminho rápido depende do kernel flash_attn_with_kvcache, que está disponível para duas combinações de dispositivo e implementação de atenção.

Para qualquer outra combinação, ou quando o kernel não pode ser importado, o manager recorre ao caminho varlen. Ele registra um aviso apenas quando você define max_blocks_per_request explicitamente.

A atenção completa suporta tabelas de blocos de kernel, mas a atenção deslizante não. Se o modelo tiver camadas de atenção deslizante, incluindo modelos híbridos que também usam atenção completa, a inicialização força max_blocks_per_request para 0, e um valor diferente de zero definido por você é sobrescrito. Em seguida, o gerente registra que o caminho rápido de decodificação não está disponível porque alguns dos tipos de atenção do modelo não suportam tabelas de blocos de kernel.

Offloading de CPU

O offloading para CPU copia os blocos de cache KV evictados para um buffer de CPU fixo (pinned) pré-alocado quando a cache KV da GPU está cheia. Depois que o espaço na cache fica disponível, o gerenciador copia os blocos de volta para a GPU e retoma a requisição sem recalcular seu prompt e os tokens gerados.

Defina cpu_offload_space para o espaço de troca da CPU em GiB. O valor padrão, 0.0, desativa o offloading para a CPU.

cb_config = ContinuousBatchingConfig(cpu_offload_space=8.0)

Por padrão, cpu_offload_space_safety_threshold=0.8 limita o espaço solicitado a 80% da RAM do sistema disponível quando psutil está instalado. Defina cpu_offload_space=None para dimensionar o pool de swap a partir do limite de segurança.

Tempo limite de paralelismo de tensor

Sob paralelismo de tensor, o gerenciador cria um grupo de comunicação de CPU para coordenar o envio, o cancelamento e o encerramento de solicitações entre os ranks. cpu_group_timeout limita o tempo máximo que uma operação coletiva neste grupo pode bloquear antes de o processo ser encerrado. Se um rank travar, o timeout impede que os outros esperem indefinidamente.

Defina um tempo limite mais longo para cargas de trabalho que emitem coletivas com pouca frequência, ou passe None para desativá-lo.

cb_config = ContinuousBatchingConfig(cpu_group_timeout=600.0)

Cache de prefixo

Quando várias solicitações compartilham um prefixo comum, como um prompt de sistema, o gerenciador reutiliza os blocos de cache KV em vez de recalculá-los. allow_block_sharing tem o valor padrão True. O compartilhamento de prefixo fica ativo apenas quando todos os tipos de alocadores de atenção suportam o compartilhamento, de modo que apenas modelos com atenção completa o utilizam. Modelos puramente deslizantes e híbridos o desativam, mesmo quando allow_block_sharing é True.

cb_config = ContinuousBatchingConfig(
    allow_block_sharing=True,  # default
)

Escolha um backend de atenção

O Continuous batching suporta FlashAttention, SDPA e "paged|eager". Defina attn_implementation ao carregar o modelo.

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

O continuous batching é mais rápido com FlashAttention do que com eager ou SDPA, principalmente porque o FlashAttention não precisa de uma attention mask. Se o modelo suporta FlashAttention, um modelo carregado com "eager" ou "sdpa" passa a usar FlashAttention 3, ou FlashAttention 2 quando o FlashAttention 3 não está instalado, e registra um aviso. Se nenhuma das versões estiver instalada, o backend não é alterado.

Para manter o eager ou o SDPA, por exemplo, ao testar, defina auto_switch_to_flash=False em ContinuousBatchingConfig. Sempre que o backend permanecer "eager", o continuous batching executa "paged|eager" em vez disso.

Paralelismo de tensor

Para modelos grandes demais para caber em uma única GPU, distribua os pesos entre os dispositivos com paralelismo de tensor. Defina o número de dispositivos com DistributedConfig(tp_size=N). O loteamento contínuo lê o tamanho do paralelismo de tensor do modelo para dimensionar o cache KV paginado por shard. Consulte Paralelismo de tensor para a lista de arquiteturas suportadas e como o sharding funciona.

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)

Execute o script com torchrun, definindo --nproc-per-node para o número de GPUs em que você deseja distribuir os shards.

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

O tamanho do paralelismo de tensores deve dividir o num_key_value_heads do modelo (verifique a configuração do modelo). Caso contrário, o cache paginado gera um erro na inicialização, portanto, escolha um --nproc-per-node adequado.

Não defina device_map junto com distributed_config. Os dois conflitam, pois device_map posiciona módulos inteiros em GPUs específicas, enquanto a paralelização de tensores distribui esses mesmos parâmetros entre todas as GPUs.

Atenção com janela deslizante

Modelos com atenção de janela deslizante (Mistral, Gemma 2) funcionam com loteamento contínuo. As camadas deslizantes usam SlidingAttentionCacheAllocator, que mantém apenas uma janela fixa de tokens recentes, reciclando um pequeno conjunto de páginas. O tipo de atenção vem do layer_types ou sliding_window do modelo, e não do 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 e o decode fast path são desativados sempre que houver sliding attention, incluindo modelos híbridos que também possuem camadas de atenção completa.

Próximas etapas

• O post no blog sobre Continuous batching aborda o cache de KV, o prefill em blocos e o agendamento dinâmico, com números de benchmark de desempenho.
• Para uma análise mais aprofundada de como o sistema de continuous batching funciona, consulte o documento Arquitetura do continuous batching.