Resumo Técnico: GQLA – Atenção Latente de Consulta em Grupo para Decodificação de Modelos de Linguagem de Grande Escala Adaptativa a Hardware

1. Declaração do Problema

A decodificação moderna de Modelos de Linguagem de Grande Escala (LLM) autoregressiva é fundamentalmente limitada pelo tráfego de cache de Chaves-Valores (KV), pois cada token gerado exige a leitura de todo o histórico de chaves e valores em cache da memória fora do chip. Embora arquiteturas recentes como a Atenção Latente de Múltiplas Cabeças (MLA), utilizada no DeepSeek-V2/V3, alcancem redução de estado da arte no cache KV ao comprimir conjuntamente chaves e valores em um latente de baixo posto, elas sofrem de três desvantagens de hardware acopladas quando implantadas em hardware de inferência diverso:

1. Acoplamento de Hardware: Os pesos treinados da MLA expõem apenas um único caminho de decodificação (uma forma de Atenção de Múltiplas Consultas "absorvida" ou MQA). Este caminho é perfeitamente ajustado para a relação de computação-largura de banda de GPUs de ponta como a NVIDIA H100 (crista $\approx$ 295 FLOPs/byte). No entanto, em GPUs restritas a exportação ou de consumo como a H20 (que mantém a largura de banda, mas reduz a computação em $\sim$7x, baixando a crista para $\approx$ 37 FLOPs/byte), este caminho torna-se severamente limitado pela computação, resultando em baixo rendimento.
2. Perda de Paralelismo de Tensores (TP): A forma MQA absorvida canaliza todas as cabeças de consulta através de um único KV latente compartilhado. Esta estrutura impede o paralelismo de tensores ao longo do eixo das cabeças, pois cada dispositivo deve replicar todo o latente, limitando a escalabilidade.
3. Ganho Nulo de Previsão de Múltiplos Tokens (MTP): O MTP, que gera múltiplos tokens por etapa para aumentar a intensidade aritmética, empurra a decodificação já limitada pela computação da H20 ainda mais além de sua crista de eficiência, resultando em ganho zero de rendimento. Por outro lado, na H100, o MTP ultrapassa a crista, reduzindo a eficiência do ganho.

2. Metodologia: Atenção Latente de Consulta em Grupo (GQLA)

Os autores propõem a Atenção Latente de Consulta em Grupo (GQLA), uma modificação arquitetônica mínima da MLA que preserva a compressão latente conjunta de baixo posto, enquanto expõe dois caminhos de decodificação algebricamente equivalentes sobre os mesmos parâmetros treinados.

Arquitetura

A GQLA indexa as matrizes de projeção ascendente ($W_{UK}, W_{UV}$) por $g$ grupos em vez de replicá-las através de todas as $h_q$ cabeças de consulta. Este design habilita dois caminhos de execução distintos:

1. Caminho MQA-Absorvido (Amigável à Memória):

   • Mecanismo: Idêntico à MLA. As projeções ascendentes são absorvidas nas projeções de consulta e saída. Todas as $h_q$ cabeças atendem diretamente ao cache latente compacto ($cKV$).
   • Cache: Armazena apenas o latente ($r_{kv}$) e uma chave RoPE compartilhada ($d^R_h$).
   • Hardware Alvo: Otimizado para hardware classe H100, onde a intensidade aritmética ($\approx$ 242 FLOPs/byte para $s_q=1$) situa-se logo abaixo da crista memória-computação.

2. Caminho GQA (Amigável à Computação):

   • Mecanismo: O latente é expandido em $g$ grupos de chaves/valores. A decodificação executa Atenção de Consulta Agrupada (GQA) padrão sem expansão latente por etapa.
   • Cache: Armazena chaves e valores expandidos por grupo ($2gd_h$) mais a chave RoPE compartilhada.
   • Hardware Alvo: Otimizado para hardware classe H20. Ao ajustar a contagem de grupos ($g$) e o número de tokens de consulta por etapa ($s_q$), a intensidade aritmética pode ser fixada na crista da H20 ($\approx$ 37 FLOPs/byte).
   • Escalabilidade: Suporta até paralelismo de tensores de 8 vias sem redundância ao longo do eixo do grupo.

Seleção de Caminho

O tempo de execução seleciona o caminho com base no hardware alvo sem retreinamento ou kernels personalizados:

• H100: Utiliza o caminho MQA-absorvido ($s_q=1$).
• H20: Utiliza o caminho GQA (ex: $g=8, s_q=2$ ou $g=4, s_q=1$).
• MTP: Na H20, o caminho GQA com $s_q=2$ recupera ganhos de rendimento próximos ao ideal, enquanto a MLA não obtém ganhos.

TransGQLA: Conversão de Modelos Pré-treinados

Para evitar o pré-treinamento do zero, os autores introduzem o TransGQLA, uma extensão do pipeline TransMLA. Ele converte um checkpoint GQA pré-treinado em um modelo GQLA através de uma única modificação direcionada:

• Fusão de Cabeças: Em vez de fundir cabeças agrupadas em um único latente e replicar projeções ascendentes através de todas as cabeças de consulta (como no TransMLA), o TransGQLA mantém as projeções ascendentes indexadas por grupo.
• Resultado: O modelo mantém estruturalmente o caminho de execução GQA, enquanto habilita o caminho MQA-absorvido via operação de absorção padrão. Isso preserva o paralelismo de tensores ao longo do eixo do grupo.

GQLA Esparso

O design estende-se à atenção esparsa de granularidade fina. Como a configuração canônica GQLA ($h_q/g = 16$) corresponde à dimensão do bloco Tensor-Core MMA ($m=16$), a GQLA esparsa pode preservar o caminho GQA em hardware classe H20. Em contraste, a MLA esparsa está estruturalmente travada no caminho MQA absorvido devido a restrições de blocos, herdando os mesmos problemas de sobrecarga de computação.

3. Contribuições Principais

1. Identificação das Desvantagens da MLA: O artigo identifica que o design de caminho único da MLA acopla-o a proporções classe H100, sacrifica o paralelismo de tensores no eixo das cabeças e produz ganhos nulos de MTP em GPUs de consumo.
2. Arquitetura GQLA: Uma variante mínima da MLA que expõe dois caminhos algebricamente equivalentes (MQA-absorvido e GQA) sobre os mesmos pesos. Uma única configuração ($h_q=128, g=8$) com uma cabeça de MTP fixa as linhas de teto tanto da H100 quanto da H20 simultaneamente.
3. Pipeline TransGQLA: Um método de conversão que transforma checkpoints GQA pré-treinados em modelos GQLA com mudanças arquitetônicas mínimas, mantendo o paralelismo de tensores e permitindo uma transição suave para a forma latente.
4. Análise de Linha de Teto: Uma verificação teórica mostrando que os pesos GQLA podem atingir rendimento próximo ao pico tanto na H100 quanto na H20 alternando caminhos, enquanto a MLA é ineficiente na H20.

4. Resultados

Os autores avaliaram o TransGQLA no modelo LLaMA-3-8B:

• Preservação de Capacidade: A conversão dos pesos GQA para a forma latente GQLA (sem treinamento adicional) resultou em uma perda de apenas $\sim$9,7 pontos médios em benchmarks de raciocínio de senso comum (MMLU, ARC, etc.) em comparação com o modelo pré-treinado original de 15T tokens.
• Compressão de Cache: No caminho MQA-absorvido, o cache KV por token foi comprimido para 28,125% da linha de base GQA original (uma redução de 71,875%).
• Projeções de Rendimento:
  • H100: O caminho MQA-absorvido atinge 354K tokens/seg.
  • H20: O caminho GQA (com $g=8, s_q=2$) atinge 221K tokens/seg, uma melhoria de 3,4x sobre a MLA no mesmo dispositivo (que fica presa em 65K tokens/seg devido à limitação por computação).
• Pré-treinamento Contínuo: Os autores projetam que o pré-treinamento contínuo (seguindo a trajetória TransMLA) recuperará a capacidade perdida para dentro de 0,5 pontos do modelo original após 30B tokens, mantendo a compressão do cache KV.

5. Significado e Alegações

O artigo afirma que a GQLA oferece um princípio de design prático e agnóstico a hardware para inferência de LLM. Ao expor múltiplos caminhos de decodificação algebricamente equivalentes sobre um único conjunto de pesos treinados, a GQLA permite que um modelo se adapte a diferentes restrições de hardware (rico em computação vs. rico em memória) no momento da implantação, sem retreinamento ou kernels personalizados.

Os autores posicionam a GQLA como uma solução para o problema de "acoplamento de hardware", permitindo implantação eficiente tanto em aceleradores de ponta (H100) quanto em hardware de consumo/restrito a exportação (H20). Eles enfatizam que esta abordagem suporta Previsão de Múltiplos Tokens e Paralelismo de Tensores em hardware onde a MLA falha em fazê-lo eficientemente.

Limitações Reconhecidas:

• As alegações de hardware baseiam-se em modelagem de Linha de Teto; benchmarks completos em nível de kernel em hardware H20/H100 real estão pendentes.
• Os resultados de pré-treinamento contínuo para TransGQLA estão atualmente em andamento e ainda não finalizados no artigo.
• A avaliação está atualmente limitada ao backbone LLaMA-3-8B e a benchmarks de raciocínio de senso comum.
• A GQLA esparsa é descrita estruturalmente, mas ainda não foi testada contra métodos de atenção esparsa existentes.