Resumo Técnico: Requests of a Feather Must Flock Together

Declaração do Problema

A inferência de Modelos de Linguagem Grandes (LLM) é fundamentalmente limitada pela memória durante a fase de decodificação auto-regressiva. Cada token gerado exige atenção ao cache de Chave-Valor (KV) de todos os tokens anteriores, necessitando de uma varredura completa do cache KV na memória da GPU a cada passo. Embora trabalhos anteriores se concentrassem em maximizar o tamanho do lote para melhorar a utilização da GPU, este artigo identifica um compromisso crítico e negligenciado: a homogeneidade de prefixos.

Em cargas de trabalho do mundo real (por exemplo, prompting com poucos exemplos, prompts de sistema, RAG), solicitações frequentemente compartilham prefixos significativos. Motores de inferência existentes (como vLLM e SGLang) utilizam árvores radiais ou estruturas de hash para detectar esses prefixos compartilhados e reutilizar caches KV, principalmente para reduzir a pegada de memória. No entanto, os agendadores atuais priorizam a maximização do tamanho do lote ou a minimização de evições de memória sem considerar o impacto da heterogeneidade de prefixos nos padrões de acesso à memória.

Os autores demonstram que lotes heterogêneos (contendo solicitações com prefixos diferentes) interrompem a localidade espacial e temporal nos acessos à DRAM da GPU, levando a uma utilização inadequada da largura de banda e à redução do rendimento. Por outro lado, lotes homogêneos de prefixos (onde todas as solicitações compartilham um prefixo comum) permitem a pré-busca eficiente de hardware e a reutilização de cache. O artigo argumenta que os agendadores existentes falham em reconhecer que um lote moderadamente pequeno e homogêneo pode superar um lote maior e heterogêneo. Além disso, os mecanismos usados para detectar o compartilhamento de prefixos (por exemplo, travessias de árvores radiais) incorrem em sobrecarga de CPU proibitiva, frequentemente comparável ou superior ao tempo de execução da GPU.

Metodologia

O artigo apresenta o Feather, um agendador consciente de prefixos projetado para otimizar o compromisso entre o tamanho do lote e a homogeneidade de prefixos. O Feather integra dois componentes principais:

1. Árvore de Hash em Blocos (CHT)

Para abordar a alta sobrecarga de CPU das travessias tradicionais de árvores radiais, o Feather emprega uma estrutura de dados leve chamada Árvore de Hash em Blocos.

• Mecanismo: Em vez de comparações token a token, as solicitações são representadas como vetores de hashes cumulativos sobre blocos de tokens de tamanho fixo.
• Eficiência: O CHT mantém um "conjunto de trabalho" de hashes para o lote ativo e um min-heap de solicitações em espera ordenadas por sua "contagem faltante" (o número de blocos ainda não cobertos pelo lote ativo).
• Complexidade: Isso permite que o agendador identifique a melhor solicitação candidata para adicionar a um lote em tempo $O(\log W)$ (onde $W$ é o número de solicitações em espera), uma melhoria significativa em relação à complexidade $O(W \times T)$ ou $O(E)$ das travessias de árvores radiais em profundidade.
• Operações: O CHT suporta inserção eficiente, encontrar a melhor solicitação, adicionar uma solicitação ao lote ativo e remover solicitações concluídas, utilizando atualizações preguiçosas do heap e manutenção incremental da ponta para minimizar a sobrecarga.

2. Política de Agrupamento por Aprendizado por Reforço (RL)

O Feather usa uma política baseada em RL para determinar o ponto de parada ideal para a formação de lotes.

• O Compromisso: Adicionar mais solicitações aumenta a utilização da GPU (maior tamanho do lote), mas pode diluir a homogeneidade de prefixos (reduzindo a localidade de memória).
• Estado e Ação: O agendador observa o tamanho atual do lote, a perda na profundidade do prefixo compartilhado ($\Delta$) e o número de solicitações em espera que compartilham a nova profundidade de prefixo ($w$). Ele decide se deve Adicionar outra solicitação ou Parar e despachar o lote.
• Políticas: Os autores avaliam três políticas: uma heurística estática, um Bandito Contextual (UCB) e Q-Learning. O objetivo é maximizar a recompensa, definida como o rendimento de decodificação.
• Adaptabilidade: A política de RL aprende a formar lotes que são "justos o suficiente" — grandes o suficiente para utilizar a GPU, mas homogêneos o suficiente para maximizar a localidade de memória. Ela se adapta dinamicamente à intensidade da carga de trabalho e aos padrões de compartilhamento de prefixos.

Contribuições Principais

1. Identificação de um Compromisso Crítico: O artigo estabelece que a homogeneidade de prefixos é um determinante mais forte do rendimento de decodificação do que o tamanho do lote sozinho. Demonstra que até mesmo uma única solicitação divergente em um lote pode interromper a localidade e reduzir o rendimento em quase 2x.
2. Quantificação da Sobrecarga de CPU: Os autores revelam que os mecanismos de detecção de prefixos mais avançados (como travessias de árvores radiais no SGLang) incorrem em sobrecargas de CPU que podem atingir 50–90% da latência total, efetivamente negando os benefícios da reutilização de prefixos.
3. Agendador Feather: O design e a implementação do Feather, que integra a estrutura de dados CHT com uma política de agrupamento baseada em RL. O Feather é integrado tanto no vLLM quanto no SGLang.
4. Agnosticismo de Hardware: Diferentemente de kernels de atenção conscientes de prefixos que exigem ajuste específico de hardware, o Feather opera na camada de agendamento, tornando-o compatível com kernels de atenção existentes e configurações de hardware.

Resultados

A avaliação foi conduzida em GPUs NVIDIA RTX 6000 Ada e A100 usando modelos como Llama 3 8B, Qwen e LongChat 13B em diversas cargas de trabalho (L-Eval, LongBench).

• Ganhos de Rendimento: O Feather alcança 2–10× maior rendimento de ponta a ponta em comparação com agendadores existentes (FCFS do vLLM, LPM/DFS-W do SGLang) em cargas de trabalho com prefixos compartilhados variando de 1K a 10K tokens.
• Sensibilidade a Prefixos: Em experimentos com 800 solicitações divididas em 8 grupos de prefixos, processá-las em 8 lotes homogêneos de 100 solicitações resultou em um rendimento quase 2× maior do que processá-las em lotes maiores e heterogêneos.
• Redução de Sobrecarga: O CHT do Feather reduz a sobrecarga de detecção de prefixos em até 1000× em comparação com abordagens baseadas em DFS para sequências longas. A sobrecarga de agendamento permanece abaixo de 1% do tempo de execução da GPU, enquanto as abordagens de base podem exceder o tempo de execução da GPU.
• Robustez: O Feather desempenha-se de forma equivalente ao status quo (FCFS) quando as cargas de trabalho não possuem compartilhamento de prefixos, garantindo nenhuma degradação de desempenho em cenários não compartilhados.
• Comparação com Kernels: O Feather supera ou iguala kernels de atenção especializados conscientes de prefixos (como PAT), apesar de não modificar o próprio kernel, provando que otimizações no nível de agendamento podem ser tão eficazes quanto as no nível de kernel.

Significado

O artigo afirma que maximizar o tamanho do lote não é mais suficiente para um serviço eficiente de LLM. O principal gargalo nas fases modernas de decodificação é a largura de banda de memória, que é fortemente influenciada pela localidade de acesso. Ao deslocar o foco de "maximizar o tamanho do lote" para "otimizar a homogeneidade de prefixos", o Feather libera ganhos de desempenho significativos sem exigir modificações de kernel específicas de hardware.

O trabalho destaca que o custo de encontrar prefixos compartilhados (sobrecarga de CPU) tem sido um gargalo oculto nos sistemas atuais. O CHT leve do Feather e sua política adaptativa de RL demonstram que é possível equilibrar dinamicamente a utilização de computação e a localidade de memória, levando a melhorias substanciais tanto no rendimento quanto na latência para cargas de trabalho de inferência de LLM do mundo real. Os autores planejam tornar o Feather de código aberto contribuindo com ele para os repositórios vLLM e SGLang.