Batching de inferência de LLM em 1 GPU

01 · CONTEXTO E REQUISITOS

O que a entrevista pede

Este é um problema de entrevista que circula em listas de perguntas de engenharia de IA: projetar um sistema de batching de inferência para uma única GPU, que aceita até cem entradas por lote enquanto os usuários esperam de forma síncrona. A lista da Exponent registra o tema como feito na Anthropic.

A resposta-modelo da Exponent abre pelos requisitos. Funcionais: o usuário envia uma entrada e espera o resultado, o sistema agrupa até cem requisições individuais numa única GPU e devolve cada resultado a quem pediu quando termina. Não funcionais: minimizar o tempo ocioso entre lotes, devolver erro quando um lote falha e escalar para várias GPUs. As entidades centrais são Request, Batch e Result.

02 · POR QUE O LOTE IMPORTA

A memória, e não o cálculo, limita a vazão

Segundo a Anyscale, a inferência de LLM é limitada pela E/S de memória, e não pelo cálculo. Por isso a vazão é determinada em grande parte pelo tamanho do lote que cabe na memória de alta largura de banda da GPU.

O artigo do vLLM mostra para onde vai essa memória numa GPU A100 de 40 GB servindo um modelo de 13 bilhões de parâmetros: aproximadamente 65% para os pesos, que ficam estáticos, e perto de 30% para os estados dinâmicos dos pedidos. Esses estados são os tensores de chave e valor da atenção, o cache KV. O restante da memória guarda outros dados, como ativações.

O cache KV pesa. No OPT de 13 bilhões de parâmetros, o cache de um único token ocupa 800 KB. Como o OPT gera sequências de até 2048 tokens, o cache de um pedido pode chegar a 1,6 GB, e mesmo reservando toda a memória disponível ao cache KV caberiam só algumas dezenas de pedidos.

A conta de guardanapo da Anyscale aponta na mesma direção: limitando as sequências a 512 tokens, cabem no máximo cerca de 28 por lote; com 2048 tokens, o lote fica limitado a 7. É um limite superior, pois não sobra espaço para os cálculos intermediários.

03 · BATCHING ESTÁTICO

O lote que espera pelo mais lento

O batching estático é a abordagem tradicional: o tamanho do lote permanece constante até a inferência terminar. Ele rende bem quando todas as saídas têm o mesmo tamanho. A Anyscale dá o exemplo de uma classificação em que cada saída tem um único token; com entradas também do mesmo tamanho, cada lote estático atinge a melhor utilização possível da GPU.

Um chatbot não pode assumir comprimentos fixos nem de entrada nem de saída, e a variação na geração pode causar subutilização enorme das GPUs.

1. No diagrama, o cliente fala com a API, a API alimenta a fila de pedidos e um montador de lotes entrega o trabalho ao trabalhador da GPU.

2. No batching estático, o tamanho do lote permanece constante até a inferência terminar.

3. Com isso, a GPU fica subutilizada até a última sequência do lote terminar de gerar.

4. Pedidos que terminaram antes dos outros não conseguem voltar ao cliente, e os recém-chegados precisam esperar o lote atual terminar por completo.

04 · CONTINUOUS BATCHING

Escalonar por iteração, não por pedido

O continuous batching também é chamado de dynamic batching ou de batching com iteration-level scheduling. Segundo a Anyscale, o artigo Orca é o primeiro a tratar esse problema, até onde ela sabe: em vez de esperar que todas as sequências do lote terminem, o tamanho do lote é determinado a cada iteração. O Orca propõe ainda o selective batching, que aplica o batching só a um conjunto selecionado de operações.

1. O endpoint coloca os pedidos recém-chegados no pool de requisições, e o escalonador seleciona do pool os pedidos que rodam a seguir.

2. O escalonador chama o motor de execução para rodar uma única iteração do modelo sobre os pedidos selecionados.

3. Como o escalonador recebe um retorno a cada iteração, ele detecta a conclusão de um pedido e devolve seus tokens ao cliente imediatamente.

4. Um pedido recém-chegado pode ser considerado depois de apenas uma iteração do modelo, reduzindo bastante o atraso de fila.

5. Quando uma sequência termina, outra pode ser inserida no lugar dela, com mais utilização de GPU do que no batching estático.

6. O prefill tem padrão computacional diferente da geração e não se agrupa facilmente com ela; os frameworks tratam isso com o hiperparâmetro waiting_served_ratio.

Nos benchmarks da Anyscale, feitos numa única GPU A100 de 40 GB com o modelo OPT-13B da Meta, o continuous batching chega a até 23 vezes mais vazão (com vLLM, que soma otimizações de memória) e a 8 vezes sobre o batching ingênuo, tanto no Ray Serve quanto no text-generation-inference da Hugging Face. Os sistemas de continuous batching também melhoram a latência mediana.

05 · MEMÓRIA E ESCALONAMENTO

O cache KV decide quantos pedidos cabem

O cache KV de um pedido cresce e encolhe conforme o modelo gera tokens, e seu comprimento não é conhecido de antemão. Os sistemas anteriores pré-alocavam um bloco contíguo do tamanho máximo do pedido. O perfil do artigo do vLLM mostra que apenas 20,4% a 38,2% da memória do cache KV guardava estados reais de tokens nesses sistemas.

1. Numa A100 de 40 GB servindo um modelo de 13 bilhões de parâmetros, aproximadamente 65% da memória vai para os pesos e perto de 30% para os estados dinâmicos dos pedidos.

2. O PagedAttention permite que o cache KV fique não contíguo, alocado em blocos de tamanho fixo, e a alocação passa a ocorrer sob demanda em vez de antecipada.

3. Segundo a Anyscale, o blog do vLLM limita o desperdício a menos de 4%, só no último bloco, e essa economia vira lote maior.

4. Quando os blocos físicos livres acabam, o vLLM seleciona sequências para despejar e transfere o cache KV delas para a CPU.

5. Na alternativa de recomputação, o cache KV é recalculado quando as sequências preemptadas voltam a ser escalonadas.

6. O vLLM adota a política primeiro a chegar, primeiro a ser servido, e os pedidos mais recentes são os primeiros a serem preemptados.

No lado do escalonador, o Orca define o max batch size como o maior número de pedidos possível num lote. Aumentá-lo troca mais vazão por mais latência, com retornos decrescentes, e o operador pode ajustar esse botão para maximizar a vazão dentro do orçamento de latência.

O escalonador do Orca também respeita a memória: seleciona no máximo max batch size pedidos pela ordem de chegada e reserva espaço para chaves e valores quando o pedido é escalonado pela primeira vez. Uma implementação ingênua pode cair num deadlock, em que o escalonador não consegue iniciar uma iteração para nenhum pedido porque não há espaço para a próxima chave e valor.

06 · TRADE-OFFS E FECHAMENTO

O que dizer ao fechar a entrevista

O continuous batching continua limitado pelo espaço disponível na GPU. À medida que o sistema satura, há menos oportunidades de injetar pedidos novos imediatamente quando eles chegam. O artigo do vLLM, por sua vez, relata melhorar a vazão de LLMs populares em 2 a 4 vezes com o mesmo nível de latência, frente a sistemas como FasterTransformer e Orca.

A Anyscale trata o custo de servir como consequência direta da eficiência com que se serve numa dada latência, e por isso mede vazão e latência em conjunto.