Zachary, Especialista em Algoritmos na Novita AI
Foco em Pesquisa: Aceleração de Inferência
Introdução
No último ano, desde o surgimento do H2O, houve uma proliferação de artigos de pesquisa sobre esparsidade KV. No entanto, um desafio significativo em aplicações do mundo real é a grande lacuna entre a pesquisa acadêmica e a implementação prática. Por exemplo, frameworks como vLLM adotam PagedAttention, uma técnica de memória baseada em paginação que é incompatível com a maioria dos algoritmos de esparsidade, ou onde esses algoritmos têm desempenho inferior ao PagedAttention. Tais problemas impedem que os algoritmos de esparsidade sejam utilizados efetivamente em ambientes de produção.
Ao revisar artigos acadêmicos recentes do último ano na área de esparsidade KV, fizemos modificações no framework vLLM com base na esparsidade KV. Essas modificações foram alinhadas com os recursos de otimização inerentes do framework vLLM, como Continuous Batching, FlashAttention e PagedAttention. Por fim, ao comparar o framework modificado com os frameworks de inferência state-of-the-art (SOTA), usando os modelos, parâmetros e hardware mais comumente implantados em produção, alcançamos uma aceleração de 1.5x no desempenho da inferência.
Antes de discutir esparsidade KV, é essencial abordar a característica das Massive Activations em LLMs. Em LLMs, um pequeno número de valores de ativação é significativamente mais ativo do que outros, às vezes ultrapassando os menos ativos por um fator de 100.000 ou mais. Em outras palavras, um pequeno subconjunto de tokens desempenha um papel crítico. Essa observação forma a base para a aplicação de técnicas de esparsidade KV, onde reter apenas os tokens importantes pode melhorar o desempenho da inferência.
Dados de análise específicos do Llama2 e Llama3 apoiam ainda mais essa observação, conforme mostrado na Fig. 1 e Fig. 2.
Fig. 1. Magnitudes de Ativação (eixo z) no LLaMA2-7B. Os eixos x e y são as dimensões de sequência e características. Para este modelo específico, observamos que ativações com magnitudes massivas aparecem em duas dimensões de características fixas (1415, 2533), e dois tipos de tokens — o token inicial, e o primeiro ponto final (.) ou token de nova linha (\ ).
Ao contrário do Llama2, o Llama3 não exibe a característica de Massive Activations em todas as camadas. Nas camadas inferiores, as ativações são mais uniformemente distribuídas, enquanto as camadas intermediárias exibem padrões de atenção localizados. É apenas nas camadas superiores que o fenômeno de Massive Activations se torna proeminente. Essa observação explica por que modelos como Llama3 adotam técnicas de esparsidade hierárquica, que adaptam a esparsidade a diferentes camadas de acordo com seus padrões de ativação específicos.
Fig. 2. Padrões de atenção da geração aumentada por recuperação (retrieval-augmented generation) através das camadas no LlaMa (Touvron et al., 2023a;b) revelam que, nas camadas inferiores, o modelo exibe um modo de atenção de amplo espectro, distribuindo as pontuações de atenção uniformemente por todo o conteúdo. Nas camadas intermediárias, a atenção torna-se mais localizada dentro de cada documento, indicando uma agregação refinada de informações (formas triangulares vermelhas tracejadas nas camadas 6 e 10). Isso culmina nas camadas superiores, onde a “atenção massiva” se concentra em alguns tokens-chave (barras de atenção concentradas após a camada 18), extraindo eficientemente informações essenciais para as respostas.
Modelo do Sistema
Após introduzir a característica de Massive Activations em LLMs, agora temos uma compreensão básica dos princípios por trás da esparsidade KV. Em seguida, discutiremos como a esparsidade KV é implementada. No processo de inferência de modelos grandes, três grandes gargalos são comumente encontrados: capacidade de memória da GPU, poder computacional e largura de banda de I/O. Para evitar cálculos redundantes, o KV é frequentemente armazenado em cache durante a inferência, levando a um consumo significativo de memória da GPU. Enquanto isso, a carga computacional e a taxa de transferência de I/O tornam-se gargalos críticos nos estágios de Prefill e Decode, respectivamente.
Nossa modificação do framework vLLM é baseada na característica de Massive Activations dos LLMs, implementando esparsidade por camadas (ou seja, diferentes níveis de esparsidade para diferentes camadas). Essa abordagem reduz significativamente a sobrecarga associada ao cache de KV, acelerando assim o processo de inferência dos modelos LLM.
Conforme mostrado na Fig. 3, durante o processo de inferência, aplicamos esparsidade KV a diferentes camadas do modelo. Ao empregar uma estratégia de poda (pruning), eliminamos pares KV com pontuações mais baixas, mantendo aqueles com pontuações mais altas e mais próximos. Essa abordagem reduz o uso de memória, bem como a sobrecarga computacional e de I/O, resultando em inferência acelerada.
Fig. 3. O impacto da esparsidade KV no modelo
Resultados da Aceleração da Inferência
A principal preocupação da maioria dos pesquisadores é o impacto real da esparsidade KV no framework vLLM. Antes de nos aprofundarmos nos detalhes, primeiro apresentamos os resultados da avaliação de desempenho.
1. Desempenho da Inferência
Em aplicações reais de LLM, um comprimento de entrada/saída de 4000/500 é a configuração mais comumente observada. Além disso, a GPU RTX 4090 é uma das GPUs mais utilizadas na indústria. Com base nessas condições, realizamos testes de desempenho em lote comparativos. O grupo de controle é o vLLM 0.6.1.p2, enquanto o grupo experimental é o Novita AI Sparse 0.5.1 (uma versão do vLLM 0.5.1 modificada com esparsidade KV). Várias rodadas de testes comparativos de desempenho foram realizadas entre os dois, com métricas-chave incluindo TTFT (time-to-first-token, crítico para a experiência do usuário) e Throughput (a velocidade real de inferência).
Os resultados finais dos testes são mostrados na Tabela 1. Ao aplicar a esparsidade KV, conseguimos aumentar a taxa de transferência (throughput) do vLLM em aproximadamente 1.58x, mantendo um TTFT utilizável (P50 abaixo de 1 segundo). Além disso, em cenários com tamanhos de lote maiores, o vLLM 0.6.1.p2 atingiu seu limite de desempenho em um nível de concorrência de 10, enquanto o Novita AI Sparse 0.5.1 foi capaz de manter um desempenho TTFT estável mesmo em um nível de concorrência de 20. Isso demonstra a robustez da esparsidade KV em garantir a estabilidade de desempenho em ambientes de produção.
Tabela 1: Comparação de Desempenho entre vLLM 0.6.1.p2 e Novita AI Sparse 0.5.1
2. Precisão do Modelo
Como a esparsidade KV é um algoritmo de compressão com perdas, é igualmente importante avaliar o desempenho do modelo. Realizamos avaliações de desempenho com base no Llama3–8B, incluindo avaliações usando MMLU e outros benchmarks. Os resultados indicam que a perda de precisão permanece dentro de aproximadamente 3%.
Conforme mostrado na Tabela 2, testamos o desempenho do modelo usando ambos os benchmarks MMLU e humanities:
Tabela 2: Resultados da Avaliação de Desempenho do Llama3–8B usando MMLU e Benchmarks de Humanidades
Em segundo lugar, realizamos avaliações de tarefas de QA em cenários de texto longo comumente encontrados em aplicações práticas, com comprimentos de entrada variando de 7k a 30k. A taxa de compressão foi mantida em mais de 10 vezes, conforme mostrado na Tabela 3. Os resultados indicaram uma perda geral de precisão de aproximadamente 10%.
Tabela 3: Resultados da Avaliação de Tarefas de QA em Texto Longo com Comprimentos de Entrada de 7k a 30k
Técnicas-Chave
1. Esparsidade Específica por Camada e Paralelismo de Tensores (Tensor Parallelism)
O framework vLLM emprega uma estratégia de escalonamento no nível da iteração (Iteration-Level Scheduling), comumente chamada de Continuous Batching. Essa abordagem de escalonamento é caracterizada por não esperar que cada sequência em um lote termine a geração; em vez disso, facilita o escalonamento no nível da iteração, resultando em maior utilização da GPU em comparação com o batching estático. Como mencionado no início deste artigo, modelos como Llama3–8B possuem inerentemente características como esparsidade hierárquica. Consequentemente, nossas modificações baseadas em Continuous Batching encontraram uma série de desafios, conforme mostrado na Fig. 4.
Fig. 4. Completando sete sequências usando continuous batching. À esquerda, você vê o lote após uma única iteração, enquanto à direita, o lote após várias iterações. Assim que uma sequência emite um token de fim de sequência, introduzimos uma nova sequência em sua posição (por exemplo, sequências S5, S6 e S7). Esse método aumenta a utilização da GPU ao evitar a necessidade de a GPU esperar que todas as sequências terminem antes de iniciar uma nova.
O principal desafio que enfrentamos foi o gerenciamento de memória em diversas camadas. Dado o uso do vLLM de um modelo unificado de gerenciamento de memória paginada e a limitação da estratégia de fila para suportar apenas Full Cache ou Sliding Window, fizemos modificações na estrutura central do vLLM para permitir suporte a ambas as opções.
Esse ajuste concede a cada camada a flexibilidade de optar por Full Cache ou Sliding Window, permitindo distintos níveis de esparsidade entre camadas, conforme representado no diagrama fornecido.
O primeiro desafio a abordar é o problema de gerenciamento de memória entre diferentes camadas. Como o vLLM utiliza um modelo unificado de gerenciamento de memória paginada para todas as camadas, e a estratégia de fila permite apenas Full Cache ou Sliding Window, precisamos ajustar a estrutura subjacente do vLLM. Especificamente, isso envolve habilitar suporte para tanto Full Cache quanto Sliding Window, de modo que diferentes camadas possam selecionar independentemente sua estratégia preferida de gerenciamento de memória. Esse ajuste facilita, em última análise, graus variados de esparsidade entre diferentes camadas. O diagrama estrutural correspondente é mostrado na Figura 5.
Fig. 5. Estrutura de Gerenciamento de Memória por Camada no vLLM
No geral, os ajustes podem ser categorizados em três fases distintas:
Fase de Inicialização do Serviço: Durante o processo de inicialização do serviço, a memória alocada para as três camadas é inicializada principalmente através do cálculo do espaço restante, juntamente com uma inicialização hierárquica da estrutura Metadata. Em contraste, o vLLM original opera como um Full Cache completo, sem levar em conta a lógica de alocação de memória para diferentes camadas.
Fase de Preparação da Etapa: Durante o processo de inicialização das etapas no vLLM, especificamente nas fases que envolvem os métodos do gerenciador de blocos e do executor do modelo, diferentes estratégias de gerenciamento de tabela de blocos são empregadas para diferentes camadas, a fim de controlar o mecanismo de atualização KV. Conforme ilustrado na figura anterior, utilizamos o mecanismo Full Cache para as três primeiras camadas, armazenando todos os KVs, enquanto as camadas além da terceira empregam a estratégia Sliding Window para substituir automaticamente os KVs com as pontuações mais baixas. Essa abordagem permite a alocação flexível de espaço de memória adaptado às necessidades de cada camada, economizando espaço enquanto maximiza a eficácia do algoritmo de esparsidade.
Fase de Computação da Esparsificação: Esta fase ocorre tipicamente após a execução do modelo, onde operações de esparsidade são executadas com base nas pontuações de esparsidade calculadas. Uma consideração chave nesta etapa é a compatibilidade com tecnologias como Tensor Parallelism, CUDA graphs e GQA. Em implementações práticas de engenharia, o Tensor Parallelism requer atenção especial. Conforme ilustrado na figura abaixo, em um cenário de múltiplas GPUs em uma única máquina, é essencial empregar um mecanismo de particionamento apropriado para confinar as operações de esparsificação KV dentro de cada GPU, evitando assim a comunicação entre GPUs. Felizmente, a unidade de Attention emprega uma abordagem de particionamento baseada em ColumnLinear, tornando sua implementação relativamente direta.
Fig. 6. Processo de Três Estágios para Gerenciamento de Memória e Esparsificação no vLLM
A discussão acima concentrou-se principalmente nas modificações feitas no nível do framework para suportar esparsidade hierárquica no vLLM. Em seguida, direcionaremos nossa atenção para as otimizações implementadas no nível do CUDA.
2. Modificação da Atenção
As modificações no nível CUDA concentram-se principalmente nos aprimoramentos feitos na unidade de computação de Attention, especificamente em relação ao FlashAttention e ao PagedAttention. Embora o PagedAttention seja relativamente direto, enfatizaremos as modificações relacionadas ao FlashAttention.
Primeiro, apresentamos a fórmula de cálculo para a Atenção:
Conforme mostrado na fórmula, o processo de cálculo do softmax envolve o cálculo do produto de QK, e é um algoritmo de 2 passagens (exige duas passagens). No entanto, o alvo final do cálculo O pode ser alcançado usando o algoritmo FlashAttention com uma única passagem (exige uma passagem). A lógica de implementação do FlashAttention pode ser ilustrada pela Fig. 7 do artigo do FlashAttention2. Em resumo, envolve a travessia em partes de Q para realizar cálculos em bloco em KV, enquanto atualiza gradualmente os dados de O, P e rowmax até que o loop termine. Finalmente, dividir O por ℓ possibilita o cálculo em uma única passagem do FlashAttention.
Fig .7. Lógica de Implementação do FlashAttention2
Para nossa implementação de esparsidade KV, é importante notar que, durante o processo de cálculo do FlashAttention, os valores necessários para o softmax — ou seja, o rowmax/ℓ (o máximo global e os valores acumulados) já são calculados. No entanto, o FlashAttention não retorna os resultados desses dois valores.
Com uma compreensão fundamental do FlashAttention, podemos agora discutir as fórmulas de pontuação para esparsificação apresentadas em artigos como PYRAMIDKV:
onde [n − α, n] é o intervalo dos tokens de instrução. Em cada camada l e para cada cabeça h, os top k l tokens com as maiores pontuações são selecionados, e seus respectivos caches KV são retidos. Todos os outros caches KV são descartados e não serão utilizados em nenhum cálculo subsequente durante todo o processo de geração.
Especificamente, calculamos a pontuação de seleção do i-ésimo token para retenção no cache KV como S em cada cabeça de atenção h por:
Na fórmula acima, A representa a pontuação do softmax, enquanto S denota a soma das pontuações dos a valores softmax mais próximos. Conforme indicado, o softmax em si é um algoritmo de 2 passagens, que não pode ser integrado diretamente no FlashAttention. Isso também implica que a pontuação de esparsificação não pode ser implementada em 1 passagem.
Como mencionado anteriormente, embora o FlashAttention não produza as pontuações do softmax, ele calcula valores intermediários como rowmax/ℓ. Portanto, consultando o Online Softmax: algoritmo de 2 passagens, podemos utilizar as saídas de rowmax/ℓ para alcançar a pontuação específica S para esparsidade KV. Na prática, precisamos apenas fazer melhorias no FlashAttention ilustrado na Fig. 7 para atender a esse requisito.
Fig .7. Modificação do FlashAttention para Pontuação de Esparsidade KV
- Saída de rowmax/ℓ no FlashAttention: O FlashAttention incorpora as saídas de rowmax/ℓ. Como esses valores representam o máximo e os totais acumulados, seu tamanho de dados é mínimo, exercendo impacto insignificante no desempenho geral do FlashAttention. Ao reutilizar rowmax/ℓ, podemos contornar a necessidade de uma segunda passagem (1-pass), melhorando assim a eficiência computacional geral.
- Adição da Etapa Online Softmax: Introduzimos a etapa Online Softmax, que se baseia no FlashAttention calculando as pontuações do softmax usando a fórmula p=e((kq)−row_max)/ℓ (lógica de 1 passagem). Posteriormente, ao realizar uma soma por linha, podemos calcular a pontuação de esparsificação S.
Na implementação do Online Softmax, dois detalhes críticos requerem atenção particular:
- Ajustando a Ordem do Loop: Como S é a soma das pontuações baseada na dimensão k, é necessário ajustar o cálculo no Online Softmax de q*k para *k**q, em comparação com o mecanismo de atenção. Esse ajuste evita comunicação desnecessária entre diferentes blocos, transformando a soma das colunas em uma soma das linhas, permitindo que todos os cálculos sejam concluídos dentro dos registradores.
- Posições de Início e Fim de q: Devido à presença de codificação causal e à lógica de soma para os a mais próximos, não é necessário percorrer todos os valores de q com k. Referenciando os indicadores de máscara rotulados como 0/1 na Fig .7, podemos determinar se um bloco correspondente requer cálculo com base nas posições inicial e final de q, o que conserva significativamente recursos computacionais. Por exemplo, se o comprimento do prompt de entrada é 5000 e o comprimento dos a mais próximos é 256, o cálculo convencional exigiria operações 5000*5000. No entanto, ao registrar as posições inicial e final, apenas operações *5000**256 são necessárias, aumentando assim a eficiência computacional.
O PagedAttention realiza um trabalho semelhante ao do FlashAttention, com o objetivo comum de calcular a pontuação final S. Uma explicação detalhada do PagedAttention não é fornecida aqui.
Conclusão
Nossa versão modificada, Novita AI Sparse 0.5.1, baseada no vLLM 0.5.1, atualmente suporta principalmente modelos como Llama3–8B e Llama3–70B. Ela opera no modo de inferência usando CUDA graphs e é implantada principalmente em GPUs de consumo como a RTX 4090. Ainda há espaço significativo para otimização tanto em aspectos de engenharia quanto algorítmicos.
Apenas um aviso: Os seguintes recursos ainda não são suportados no vLLM-0.6.2:
- Chunked-prefill
- Prefix caching
- FlashInfer e outros backends não-FlashAttention
Para aqueles interessados no Novita AI Sparse 0.5.1, você pode visitar o Novita AI para experimentá-lo antes dos outros.
Referências:
[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
[2] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference
[3] SnapKV: LLM Knows What You are Looking for Before Generation
[4] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
[5] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference
[6] MiniCache: KV Cache Compression in Depth Dimension for Large Language Models
[7] Layer-Condensed KV Cache for Efficient Inference of Large Language Models
[8] TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding
[9] CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion
[10] KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation
