Resumo Técnico: PersistentKV

Declaração do Problema

O atendimento de modelos de linguagem de grande escala (LLM) autorregressivos é cada vez mais limitado pelo movimento do cache Key-Value (KV) em vez da multiplicação de matrizes densas. Embora os sistemas modernos de atenção paginada (ex: vLLM, FlashInfer) tenham reduzido com sucesso a fragmentação do cache KV e otimizado os kernels de decodificação nativos paginados, não existe um único "melhor" kernel para todos os cenários de atendimento.

Especificamente, existe uma tensão entre:

1. Decodificação de contexto longo com baixa atividade: Onde um único kernel otimizado pode subutilizar GPUs de consumo devido à insuficiência de trabalho paralelo.
2. Comprimentos de sequência mistos: Onde o agendamento de muitos lançamentos de comprimento exato gera alto overhead, enquanto lotes (batches) grosseiramente preenchidos introduzem ineficiência.

O artigo investiga se um motor de decodificação de tabela de blocos nativo, combinado com um agendamento consciente de páginas, pode melhorar a latência de atendimento em GPUs de consumo nos regimes onde kernels de atenção paginada de qualidade de produção já são fortes.

Metodologia

Os autores propõem o PersistentKV, um sistema que compreende um motor de decodificação de tabela de blocos nativo e uma política de agendamento consciente de páginas.

1. Motor de Decodificação de Tabela de Blocos Nativo

O PersistentKV é projetado para Grouped-Query Attention (GQA), onde $G = H_q/H_{kv}$ cabeças de consulta compartilham uma cabeça KV.

• Layout: Ele executa diretamente sobre layouts de cache KV paginados nativos (tabelas de blocos) sem reempacotar as páginas K,V em tensores contíguos.
• Divisão de Sequência: Para expor paralelismo suficiente para cargas de trabalho de baixa atividade, a dimensão da sequência é dividida em $S$ intervalos disjuntos. Cada divisão computa um estado de softmax online parcial. Um kernel de mesclagem separado combina esses estados exatamente.
• Limites Locais por Linha: Diferente de abordagens baseadas em baldes (buckets) que utilizam comprimentos preenchidos (padded), o PersistentKV utiliza os comprimentos reais de sequência por linha ($N^*_i$) para determinar contagens de tiles e intervalos de divisão. Divisões vazias escrevem estados neutros e saem precocemente, evitando leituras de KV preenchidas.
• Fila de Trabalho Compacta: Para lidar com lotes irregulares (requests com comprimentos variados) sem excesso de lançamentos de kernel, o agendador emite uma grade 1D de tarefas apenas para intervalos de linha–cabeça-KV–divisão não vazios. Isso evita o "fan-out de lançamento" do bucketing de comprimento exato, preservando o paralelismo de sequência.

2. Política Adaptativa Calibrada

O sistema emprega uma política calibrada manualmente para selecionar a rota de execução com base no tamanho do lote ($B$) e no comprimento máximo de contexto ($N_{max}$):

• $B=1$ (Contexto Longo): Utiliza o PersistentKV com baldes de comprimento e divisão de sequência para expor CTAs (Compute Thread Arrays) suficientes.
• $B \ge 8$ (Contexto Longo): Utiliza o PersistentKV com a fila de trabalho compacta para amortizar o overhead de dois kernels (divisão + mesclagem).
• $B=4$ (Fronteira): Direciona para o FlashInfer. Neste tamanho de lote, o overhead de divisão e mesclagem no PersistentKV supera os benefícios, causando uma regressão no tempo de execução (wall-clock time).
• Caso contrário: Utiliza FlashInfer.

3. Protocolo de Avaliação

• Hardware: NVIDIA RTX 3060 (28 SMs, 12GB VRAM).
• Baselines: FlashInfer 0.2.5 (primário), vLLM, TensorRT-LLM.
• Traços: Traços sintéticos (Bucketed, Bimodal, Uniform, Zipf) projetados para isolar comprimento de contexto, tamanho de lote ativo e irregularidade de comprimento.
• Métricas: Correção (tolerância FP16 contra FlashInfer), tempo de evento CUDA (trabalho de GPU enfileirado) e tempo de parede sincronizado (latência end-to-end incluindo planejamento do host).

Principais Contribuições

1. Motor GQA de Tabela de Blocos Nativo: Um motor de decodificação que expõe trabalho de paralelismo de sequência para decodificação de contexto longo de baixa atividade sem reempacotar caches KV.
2. Agendamento Consciente de Página: Implementação de agendamento de comprimento de sequência local por linha e uma fila de trabalho compacta que alinha fronteiras de divisão e prefetch com comprimentos reais de requisição, evitando grades retangulares vazias.
3. Metodologia Adaptativa: Um framework de avaliação rigoroso usando sementes de calibração/held-out, verificações de correção e contadores de carga de trabalho derivados do agendador para demonstrar que a atribuição de trabalho é uma variável decisiva junto com a matemática de atenção.

Resultados

Na semente mantida (20260623) com limiares fixos derivados de uma semente de calibração (20260622):

• Requisição Única ($B=1$): O PersistentKV com baldes de comprimento atinge uma melhoria de 1.399× no throughput de parede sincronizado sobre o FlashInfer em traços de contexto longo com baldes.
• Lote Alto ($B=8$):
  • Bimodal: Melhoria de 1.265×.
  • Uniform: Melhoria de 1.082×.
  • Tipo Zipf: Melhoria de 1.063×.
• Caso de Fronteira ($B=4$): A política adaptativa seleciona corretamente o FlashInfer, evitando uma regressão de desempenho onde o PersistentKV teria um desempenho inferior.
• Overhead de Agendamento: A fila de trabalho compacta reduz o fan-out de lançamento de 16.00 lançamentos/passo (buckets exatos) para 2.00 lançamentos/passo em traços bimodal $B=8$, enquanto reduz o tráfego de estado de mesclagem em aproximadamente 37%.

Significância e Alegações

O artigo faz uma alegação modesta e focada em sistemas: O PersistentKV identifica um nicho concreto de sistemas para o agendamento adaptativo de decodificação consciente de página.

• Não é um Kernel Universal: Os autores afirmam explicitamente que o PersistentKV não é universalmente mais rápido; o FlashInfer permanece superior para pequenos lotes ativos onde o overhead de divisão-mesclagem domina.
• A Atribuição de Trabalho Importa: Os resultados demonstram que "a atribuição de trabalho, não apenas a matemática da atenção, é uma variável decisiva do sistema de atendimento".
• Contribuição de Nível de Workshop: O artigo posiciona-se como uma contribuição de nível de workshop que fornece uma metodologia reproduzível e ablações. Destaca que, em regimes de contexto longo e baixa atividade em GPUs de consumo, dividir a sequência e usar uma fila de trabalho compacta pode superar significativamente um único kernel altamente otimizado.
• Direções Futuras: Os autores observam que um sistema de produção exigiria um modelo de custo aprendido para os limiares, integração em um runtime de atendimento completo (incluindo MLP e amostragem) e testes em hardware da geração Hopper com diferentes características de desempenho.