LLM Serving Optimization with Variable Prefill and Decode Lengths

Resumo Técnico: Otimização de Atendimento de LLM com Comprimentos Variáveis de Prefill e Decode

Formulação do Problema
Este artigo aborda o problema de escalonamento offline para a inferência de Grandes Modelos de Linguagem (LLMs) em um único trabalhador computacional com um orçamento fixo de memória de cache Key-Value (KV). O sistema lida com uma fila de $n$ requisições que chegam simultaneamente em $t=0$. Cada requisição $i$ é caracterizada por um comprimento de prompt heterogêneo (tamanho de prefill, $s_i$) e um comprimento de resposta (tamanho de decode, $o_i$).

O desafio central surge da natureza dinâmica do consumo de memória durante o decode autorregressivo. Enquanto a fase de prefill consome memória proporcional a $s_i$, a fase de decode consome memória que cresce linearmente a cada token gerado ($s_i + j$ para o $j$-ésimo token). Consequentemente, o número de requisições simultâneas que um trabalhador pode processar flutua dinamicamente com base no progresso de geração das requisições ativas. O objetivo é minimizar a Latência Total de Ponta a Ponta (TEL), definida como a soma dos tempos de conclusão para todas as requisições, sujeito à restrição de que o uso agregado de memória de qualquer lote em qualquer passo de tempo não exceda a capacidade $M$.

Os autores relaxam a suposição de tamanhos uniformes encontrada em trabalhos anteriores (Jaillet et al., 2025), considerando um cenário realista onde as cargas de trabalho misturam prompts de conversação curtos com tarefas de sumarização de documentos longos.

Descobertas Teóricas e Desafios
O artigo estabelece que relaxar a suposição de tamanho de entrada uniforme altera fundamentalmente a complexidade do problema de escalonamento:

1. NP-Dureza: O problema de escalonamento generalizado com $s_i$ e $o_i$ heterogêneos é provado como NP-difícil via redução do problema 3-Partição restrito.
2. Falha de Heurísticas Padrão: Políticas padrão, incluindo First-Come-First-Served (FCFS), Shortest-Output-First (MC-SF) e priorização baseada em Tamanho Total, mostram ter razões de aproximação ilimitadas neste cenário heterogêneo. Especificamente, uma requisição com um input grande mas output pequeno pode consumir memória significativamente de forma breve, enquanto uma requisição com um input pequeno mas output longo pode ocupar memória por um período estendido, criando trade-offs que heurísticas simples não conseguem resolver.

Metodologia: O Algoritmo Sorted-F
Para enfrentar esses desafios, os autores propõem o Sorted-F, um algoritmo de escalonamento projetado para alcançar uma garantia de aproximação de fator constante. O algoritmo opera em duas fases:

1. Construção de Lote (Fase 1): O algoritmo seleciona iterativamente lotes viáveis de requisições para formar uma lista ordenada. A seleção é guiada por uma nova métrica de qualidade, $F(X)$, definida para um conjunto de requisições $X$ como:
   $$F(X) = \frac{\sum_{r_i \in X} o_i}{|X|^2}$$
   Esta métrica equilibra o tamanho do lote ($|X|$) contra o custo agregado de decode subsequente ($\sum o_i$). Minimizar $F(X)$ é equivalente a maximizar a "densidade de throughput" (requisições concluídas por unidade de tempo médio de decode). Dentro de cada lote selecionado, as requisições são ordenadas por comprimento de saída não decrescente.
2. Execução (Fase 2): O sistema executa o escalonamento sobre passos de tempo discretos, admitindo requisições da lista ordenada no lote atual enquanto as restrições de viabilidade de memória forem satisfeitas, utilizando batching contínuo e mistura de prefill/decode.

Garantia de Aproximação
Os autores provam que o Sorted-F alcança uma razão de aproximação constante de no máximo 48 em relação ao escalonamento offline ótimo. A prova envolve uma série de transformações estruturais (separação, agrupamento e alinhamento) que convertem o escalonamento do Sorted-F e o escalonamento ótimo em formas comparáveis, demonstrando que o aumento da latência em cada etapa é limitado por um fator constante.

Implementações Práticas
Reconhecendo que a otimização combinatória exigida na Fase 1 é computacionalmente dispendiosa para $n$ grandes, o artigo propõe três solvers práticos:

• Programação Dinâmica Exata: Fornece uma solução exata para instâncias pequenas ($n \le 100$) com complexidade $O(n^2 M)$.
• Busca de Troca Local (Local Swap Search): Uma heurística para escalas moderadas ($n \le 500$) que melhora iterativamente uma solução greedy inicial via trocas de pares ($O(n^2)$).
• Seleção Greedy de Quantis: Uma heurística de tempo linear ($O(n)$) para sistemas de grande escala ($n \ge 1000$) que utiliza amostragem estatística para filtrar outliers e preencher a capacidade de forma voraz (greedy).

Adicionalmente, o artigo introduz o Sorted-LP, uma heurística que utiliza a solução de uma relaxação de Programação Linear (LP) do problema para ordenar as requisições, e o Sorted-F+Amin, uma variante adaptativa que lida com a incerteza de predição de comprimentos de saída ao refinar dinamicamente as estimativas durante a inferência.

Resultos Experimentais
Experimentos numéricos em cargas de trabalho mistas (combinando o dataset LMSYS-Chat e sumarização de documentos longos do Arxiv) demonstram:

• Redução de Latência: O Sorted-F (usando Local Swap Search) supera consistentemente o FCFS, MC-SF e baselines guiados por LP. Em uma mistura 80/20 de curto/longo com 2000 requisições, o Sorted-F reduziu a latência média em 4.87$\times$ em comparação ao FCFS e 2.09$\times$ em comparação ao MC-SF.
• Gap de Otimalidade: Em instâncias tratáveis, o gap empírico entre o Sorted-F e o limite inferior da relaxação LP é pequeno (aproximadamente 1.03–1.09), sugerindo que o algoritmo performa próximo do ótimo na prática, apesar do limite teórico frouxo.
• Sensibilidade à Carga de Trabalho: Os benefícios do Sorted-F são mais pronunciados em regimes heterogêneos e sobrecarregados, onde a memória é uma restrição vinculante. Em regimes homogêneos ou de baixa carga, a vantagem diminui.
• Desempenho Online: Uma variante de horizonte recorrente (Sorted-F-Online) testada sob chegadas de Poisson mostra melhorias significativas sobre o MC-SF e FCFS em regimes online sobrecarregados, embora ofereça menos benefício em regimes estáveis de baixa latência.

Significância e Alegações
O artigo alega que comprimentos de prefill variáveis não são um detalhe de modelagem menor, mas um fator fundamental que torna as heurísticas de escalonamento padrão ineficazes e o problema NP-difícil. A principal contribuição é o algoritmo Sorted-F, que fornece a primeira garantia de aproximação de fator constante para o escalonamento de LLM sob restrições heterogêneas de input/output. O trabalho une a lacuna entre limites teóricos de escalonamento e o atendimento prático de LLM, oferecendo um framework escalável e consciente da memória que reduz significativamente a latência em cenários de carga de trabalho mista sem exigir a predição perfeita dos comprimentos de saída. Os autores enfatizam que, embora o limite teórico seja 48, os resultados empíricos sugerem que o algoritmo está muito mais próximo da solução ótima em configurações realistas.