Zipage: Maintain High Request Concurrency for LLM Reasoning through Compressed PagedAttention

O artigo apresenta o Zipage, um motor de inferência de LLM que utiliza a técnica Compressed PagedAttention para superar gargalos de memória e permitir alta concorrência em tarefas de raciocínio, alcançando mais de 2,1 vezes de aceleração com apenas 5% de perda de desempenho em relação à abordagem de KV cache completo.

O essencial

Imagine que você tem um chef de cozinha genial (o Modelo de Linguagem, ou LLM) que é incrível criando receitas complexas, resolvendo problemas de matemática ou escrevendo códigos. Mas há um problema: a cozinha desse chef é muito pequena.

Quando o chef pensa em uma resposta longa, ele precisa anotar tudo o que já pensou em um caderno (chamado de KV Cache) para não esquecer o contexto. Se o cliente pedir uma resposta muito longa, o caderno enche rapidinho. Como a cozinha é pequena, assim que o caderno enche, o chef tem que parar de atender novos clientes ou começar a jogar páginas do caderno fora para fazer espaço. Isso é lento e faz o chef esquecer detalhes importantes, estragando a resposta.

O papel "Zipage" apresenta uma solução inteligente para esse problema, permitindo que a cozinha atenda muitos clientes ao mesmo tempo sem ficar bagunçada ou lenta.

Aqui está a explicação passo a passo, usando analogias do dia a dia:

1. O Problema: O Caderno Infinito

Normalmente, para cada pergunta longa, o sistema precisa guardar tudo o que foi dito. É como se cada cliente tivesse um caderno gigante. Se você tem 100 clientes, você precisa de 100 cadernos gigantes. A memória do computador (a geladeira da cozinha) acaba rápido, e você só pode atender poucos clientes de cada vez.

2. A Solução Mágica: "Compressed PagedAttention"

Os autores criaram uma técnica chamada Compressed PagedAttention. Pense nela como um sistema de arquivamento inteligente e comprimido.

• PagedAttention (Páginas): Em vez de um caderno contínuo, o sistema usa "blocos" (como folhas soltas). Se você precisa de mais espaço, o sistema apenas adiciona uma nova folha. Isso já é eficiente, mas ainda ocupa muito espaço se a conversa for longa.
• Compressão (O Truque): Aqui está a mágica. O sistema sabe que, em uma conversa longa, nem tudo é importante. Algumas partes são apenas "encheção de linguiça" ou repetições. O Zipage olha para o caderno, identifica as partes menos importantes e as comprime (joga fora o que não é essencial), mantendo apenas o "ouro" (as informações cruciais).
• O Resultado: O caderno de cada cliente nunca fica maior que um tamanho fixo (digamos, 4 páginas). Mesmo que a conversa dure horas, o sistema só guarda o que realmente importa. Isso libera espaço na memória para atender muitos mais clientes simultaneamente.

3. O Gerente de Cozinha: "Zipage"

Com essa técnica, eles criaram um novo "motor" de inferência chamado Zipage. Pense no Zipage como um gerente de cozinha superorganizado que faz três coisas incríveis:

• Agendamento Híbrido (A fila inteligente): O gerente sabe que alguns clientes só querem um café rápido (respostas curtas) e outros querem um banquete (respostas longas). Ele mistura as filas para que os clientes rápidos não fiquem presos atrás dos lentos, e os lentos não bloqueiem a cozinha inteira. Ele usa o espaço vazio deixado pelos rápidos para atender mais gente.
• Cache de Prefixo (A receita compartilhada): Se 10 clientes pedem "Como fazer um bolo de cenoura?", o gerente não escreve a receita do zero 10 vezes. Ele escreve uma vez e compartilha a "página inicial" para todos. O Zipage faz isso mesmo comprimindo o resto da conversa, economizando muito tempo e espaço.
• Compressão Assíncrona (Cozinhar enquanto espera): Normalmente, quando o sistema precisa comprimir o caderno, ele para tudo para fazer isso. O Zipage faz isso "nas costas" (assincronamente). Enquanto o chef está escrevendo a próxima palavra, o ajudante está organizando e comprimindo as páginas antigas no fundo. Nada para, tudo flui.

4. Os Resultados: Mais Rápido e Mais Inteligente

O papel mostra que, ao usar o Zipage:

• Velocidade: O sistema ficou mais de 2 vezes mais rápido do que os sistemas atuais em tarefas de raciocínio (como matemática e código).
• Qualidade: Mesmo jogando fora informações menos importantes, a qualidade da resposta manteve-se em 95% do nível original. É como se você tivesse um resumo perfeito que não perde nenhum ponto crucial da história.
• Concorrência: O sistema consegue atender muito mais pessoas ao mesmo tempo sem travar.

Resumo em uma frase

O Zipage é como um sistema de gerenciamento de memória que transforma um caderno infinito e pesado em uma série de cartões de visita compactos e inteligentes, permitindo que a inteligência artificial resolva problemas complexos para milhares de pessoas ao mesmo tempo, sem esquecer nada importante e sem travar a cozinha.

É uma evolução que torna os modelos de IA mais acessíveis, rápidos e capazes de lidar com o mundo real, onde as conversas são longas e as pessoas são muitas.

Resumo técnico

Resumo Técnico: Zipage

1. O Problema

Com o avanço dos Grandes Modelos de Linguagem (LLMs) focados em raciocínio (como em matemática e codificação), o paradigma de geração mudou para exigir sequências de saída extremamente longas antes de produzir a resposta final. Isso cria um gargalo crítico de memória:

• Gargalo de Memória (KV Cache): Durante a fase de decodificação, o cache de Chave-Valor (KV) cresce linearmente com o comprimento da sequência. Em cenários de alta concorrência (múltiplos usuários simultâneos), a memória da GPU esgota-se rapidamente, limitando o número de requisições que podem ser atendidas.
• Limitações das Soluções Atuais:
  • Métodos de evicção de KV cache existentes (como SnapKV, PyramidKV) geralmente focam na compressão do input ou não são integrados a motores de inferência modernos, perdendo recursos essenciais como continuous batching e prefix caching.
  • Métodos que integram evicção em motores (como R-KV, G-KV) muitas vezes usam evicção granular por página (bloco), o que pode descartar informações críticas e degradar a qualidade do raciocínio.
  • Soluções como KV-Compress comprimem apenas a entrada, mas interrompem o cache de prefixo e aumentam drasticamente o custo de pré-preenchimento (prefilling).

2. Metodologia: Compressed PagedAttention

Os autores propõem o Compressed PagedAttention, uma abordagem que combina o gerenciamento de memória baseado em páginas (PagedAttention, usado no vLLM) com uma estratégia de evicção de KV cache token a token e flexível.

Componentes Principais:

• Compressed PagedAttention:

  • Define um limite máximo de blocos ($N_{max}$) que uma requisição pode ocupar durante a decodificação.
  • Quando uma requisição atinge esse limite, é acionada uma operação de compressão que evoca os tokens menos importantes do KV cache e reorganiza os mantidos nos primeiros $N_{max}-1$ blocos, liberando o bloco restante para uso futuro.
  • A evicção é baseada em pontuação de importância (usando atenção, redundância e scores globais), garantindo que apenas informações críticas sejam descartadas.

• Estratégia de Agendamento Híbrido (Hybrid Scheduling):

  • Para evitar subutilização de memória e bloqueios, o Zipage implementa um agendamento inteligente.
  • Permite que requisições com sequências curtas (que não atingem o limite de blocos) sejam executadas sem alocar slots de consulta para compressão, aumentando a concorrência.
  • Gerencia a preempção de requisições de forma seletiva para evitar deadlocks e garantir que requisições com prefixos compartilhados não sejam bloqueadas.

• Prefix Caching Adaptado:

  • A compressão tradicional quebraria a estrutura de prefixos compartilhados. O Zipage modifica o processo para comprimir apenas em blocos de destino específicos, preservando os blocos compartilhados (prefixos) entre múltiplas requisições. Isso reduz drasticamente o custo de prefilling e o uso de memória.

• Decodificação e Compressão Assíncronas:

  • Reconhecendo que a compressão consome recursos de GPU, o sistema executa a compressão de forma assíncrona em relação à decodificação.
  • Requisições que não precisam de compressão continuam a decodificar enquanto outras são comprimidas em segundo plano, maximizando a utilização da GPU e evitando gargalos de latência.

• Otimizações de Kernel (Flash e Lightning Redundancy):

  • Para resolver o alto custo computacional da pontuação de redundância (que originalmente era $O(N^2 \times b^2)$), os autores propõem o Lightning Redundancy Score ($O(N \times b^2)$).
  • Esta técnica calcula a similaridade apenas dentro de blocos e utiliza memória compartilhada e registradores para reduzir a complexidade de memória, acelerando significativamente o processo de compressão.

3. Contribuições Chave

1. Arquitetura Zipage: Desenvolvimento de um motor de inferência de alta concorrência baseado em Compressed PagedAttention.
2. Gerenciamento de Memória Híbrido: Integração de evicção token a token com PagedAttention, permitindo limites de memória fixos por requisição sem sacrificar a qualidade do raciocínio.
3. Preservação de Prefixo: Uma estratégia inovadora para manter o prefix caching mesmo durante a compressão dinâmica, algo que métodos anteriores não conseguiam fazer eficientemente.
4. Otimizações de Baixo Nível: Implementação de kernels GPU personalizados (via Triton) para scores de redundância e compressão, reduzindo a sobrecarga computacional para níveis quase imperceptíveis.

4. Resultados Experimentais

Os testes foram realizados em modelos de raciocínio (Qwen3 e DeepSeek-R1) em tarefas de matemática (AMC 23, AIME 24) e codificação (LiveCodeBench).

• Desempenho de Throughput (TPS):
  • O Zipage alcançou um aceleramento de mais de 2.1x em comparação com motores de KV cache completo (Full KV) em tarefas de raciocínio matemático.
  • Superou significativamente o vLLM e o Nano-vLLM em todos os cenários de carga, especialmente em cargas de trabalho mistas e de longa duração.
• Qualidade do Modelo (Pass@1):
  • O Zipage manteve ~95% do desempenho de um motor Full KV (usando um orçamento de cache de 2048 tokens), demonstrando que a evicção agressiva não degrada significativamente a capacidade de raciocínio.
  • Em orçamentos maiores (4096), o desempenho aproximou-se ainda mais do Full KV.
• Eficiência de Recursos:
  • A técnica de compressão assíncrona e o Lightning Redundancy Score reduziram o tempo de compressão para cerca de 10% do tempo total, permitindo que a decodificação ocorresse sem interrupções significativas.
  • O uso de prefix caching no Zipage resultou em ganhos massivos de throughput em cenários com muitos prompts compartilhados.

5. Significado e Impacto

O trabalho Zipage representa um avanço crucial para a viabilidade industrial de LLMs de raciocínio em produção.

• Escalabilidade: Permite que provedores de serviços de IA atendam a um número muito maior de usuários simultâneos sem precisar de hardware exponencialmente maior, resolvendo o problema de escassez de memória em sequências longas.
• Viabilidade Econômica: Ao reduzir o custo de inferência por token (aumentando o TPS) sem sacrificar a qualidade da resposta, torna os modelos de raciocínio complexos mais acessíveis comercialmente.
• Inovação em Sistemas: Demonstra que é possível combinar técnicas de compressão de dados (geralmente vistas como perda de informação) com gerenciamento de memória de página de alta performance, criando um novo padrão para motores de inferência de próxima geração.

Em resumo, o Zipage resolve o dilema entre concorrência (quantas requisições) e qualidade (precisão do raciocínio) em LLMs, tornando-se uma solução essencial para a próxima geração de aplicações de IA generativa.