Resumo Técnico: Atenção Global Hierárquica (HGA)

Definição do Problema

O principal gargalo para a implantação de Grandes Modelos de Linguagem (LLMs) pré-treinados de contexto longo não é apenas a complexidade computacional $O(n^2)$ da atenção causal densa, mas o requisito de memória do cache de Chave/Valor (K/V). Na atenção densa, o cache K/V para todos os tokens históricos deve residir na memória do acelerador (VRAM). Para modelos grandes quantizados (ex: FP8), os próprios pesos do modelo frequentemente consomem a maior parte do orçamento de VRAM, deixando espaço insuficiente para armazenar o cache K/V para contextos longos (ex: 32K–64K tokens). Essa limitação impede a inferência "out-of-the-box" de checkpoints pré-treinados em hardware de classe de consumo ou de GPU única sem retreinamento ou modificação arquitetural.

Métodos existentes de atenção esparsa frequentemente exigem treinamento do zero, introduzem parâmetros de calibração ou aproximam o próprio output da atenção, potencialmente degradando a qualidade das representações aprendidas. Há uma necessidade de um mecanismo de substituição direta (drop-in) que permita que checkpoints pré-treinados existentes utilizem recuperação esparsa e eficiente em termos de memória, preservando a exatidão do cálculo final da atenção.

Metodologia: Atenção Global Hierárquica (HGA)

HGA é um patch de nível de sistema projetado como um substituto direto (drop-in replacement) para a atenção causal densa em transformers pré-treinados. Ele preserva os parâmetros originais do checkpoint ($W_Q, W_K, W_V, W_O$) e não requer retreinamento ou calibração.

Mecanismo Central

O HGA desloca o problema de recuperação de longo alcance da materialização densa token-a-token para um processo de roteamento hierárquico de dois níveis, seguido de atenção exata ao nível de token:

1. Particionamento: A sequência de entrada é dividida em blocos (chunks) de tamanho fixo (ex: $C=64$ tokens). Cada bloco é subdividido em grupos.
2. Roteamento Hierárquico:
   • Nível 1 (Roteamento de Bloco): O bloco atual consulta uma tabela compacta de "resumos de chaves de blocos" derivados de blocos anteriores.
   • Nível 2 (Roteamento de Grupo): Dentro dos blocos selecionados, a consulta roteia para "resumos de chaves de grupo" específicos para refinar a seleção.
   • Visibilidade Determinística: A política de roteamento combina a recuperação baseada em conteúdo com janelas de visibilidade fixas: o bloco atual, um número configurável de blocos "sumidouro" (sink) iniciais e blocos locais recentes estão sempre visíveis.
3. Resumos Sensíveis ao RoPE: Para garantir a compatibilidade com Embeddings de Posição Rotativos (RoPE), o HGA constrói resumos usando uma estratégia mista: pares de RoPE de alta frequência são rotacionados por token antes da média, enquanto pares de baixa frequência são calculados pela média no espaço de chave bruta e rotacionados na posição representativa do bloco. Esses resumos são computados a partir das chaves projetadas existentes sem novas camadas lineares.
4. Atenção Exata ao Token: Uma vez que os blocos ou grupos relevantes são selecionados, o roteador busca os vetores K/V originais ao nível de token da memória RAM do host (ou NVMe) para a VRAM. O output final da atenção é computado usando o softmax padrão sobre este conjunto limitado e exato de tokens. Nenhum valor de resumo sintético é usado no output final.

Arquitetura de Memória

O HGA introduz uma abstração de armazenamento K/V em camadas:

• Hot (Quente): Pesos do modelo, resumos de blocos e janelas sempre visíveis residem na VRAM.
• Warm (Morna): Um cache LRU limitado na VRAM mantém blocos roteados recentemente.
• Cold (Fria): Todo o histórico de tokens do cache K/V reside na RAM do host (ou NVMe).
  Isso desacopla o consumo de memória da GPU do comprimento total do contexto; o uso de VRAM depende primariamente dos pesos do modelo e do conjunto de trabalho roteado, não do comprimento total da sequência.

Principais Contribuições

1. Compatibilidade Direta (Drop-In Compatibility): Um mecanismo para substituir a atenção densa em transformadores GQA/MoE sem modificar pesos, normas ou introduzir parâmetros de calibração.
2. Recuperação Exata de Tokens: Ao contrário de métodos que aproximam os outputs de atenção, o HGA aproxima apenas a etapa de recuperação, garantindo que a atenção final seja computada exatamente sobre os tokens selecionados.
3. Política de Roteamento de Dois Níveis: Uma estratégia que combina sumidouros determinísticos/janelas locais com recuperação baseada em conteúdo do contexto médio, utilizando resumos de blocos e de grupos.
4. Sistema de Armazenamento em Camadas: Uma implementação que permite que todo o histórico K/V resida na RAM do host, possibilitando a inferência de contexto longo em hardware com VRAM limitada (ex: uma única RTX 5090).
5. Desempenho Zero-Shot: Capacidade demonstrada de executar modelos pré-treinados (Qwen3-30B-A3B) em contextos de 32K–64K sem ajuste fino (fine-tuning).

Resultados Experimentais

O artigo apresenta benchmarks através de um SmallLM de 40M de parâmetros e um modelo Qwen3 de 30B de parâmetros.

• Qualidade (Gap de Perda):
  • Em um SmallLM de 40M com cópia direta de pesos (sem ajuste fino), o modelo roteado atinge um gap de perda de +0.018 nats comparado à atenção densa em 8K tokens.
  • Através de comprimentos de contexto de 4K a 64K, o gap de perda permanece dentro de 0.01–0.02 nats em relação à atenção densa ao utilizar níveis de esparsidade de aproximadamente 3% a 12.5%.
  • Em 64K tokens, uma esparsidade de 6.25% mantém um gap de ~0.01 nats.
• Precisão de Recuperação:
  • Um teste "Needle-in-a-Haystack" no Qwen3-30B-A3B em 64K tokens alcançou uma taxa de sucesso de 100% (3/3 profundidades) usando roteamento de nível de grupo com ~1.9% de esparsidade, sem qualquer ajuste fino.
• Desempenho:
  • No SmallLM de 40M, uma implementação fundida em Triton alcançou um aceleração de 2.72× no treinamento e 2.43× de aceleração no prefill a 12K tokens em comparação com uma linha de base densa.
  • No Qwen3-30B, o modelo com patch roda interativamente a 32K de contexto em uma única RTX 5090 (32GB VRAM), uma configuração onde o armazenamento denso de K/V é impossível.
• Ajuste Fino (Fine-Tuning): O ajuste fino de um modelo de 0.6B com atenção roteada resultou em um throughput de ~917 tokens/seg (vs. ~570 para o denso) com um custo de roteamento negligível (+0.015 nats) em comparação com a linha de base densa.

Significância e Alegações

O artigo afirma que o HGA fornece uma solução prática de nível de sistema para o gargalo de memória de LLMs de contexto longo. Sua significância reside em:

• Preservação do Conhecimento Pré-treinado: Ao evitar o retreinamento e usar K/V exato de tokens para o output final, o HGA permite que checkpoints existentes se beneficiem da atenção esparsa sem degradar suas representações aprendidas.
• Eficiência de Hardware: Permite a implantação de modelos grandes e quantizados em hardware de consumo para contextos que excedem em muito a capacidade de VRAM necessária para caches densos.
• Atribuição de Erro: Os autores observam que o pequeno gap de qualidade restante (0.01–0.02 nats) não escala rapidamente com o comprimento do contexto. Eles sugerem que este gap é provavelmente dominado pela extrapolação de codificação posicional de contexto longo, em vez do algoritmo de roteamento em si, implicando que o HGA é um mecanismo de recuperação robusto que pode ser otimizado através da melhoria das representações posicionais.

Os autores posicionam o HGA não como um redesenho da arquitetura do transformer, mas como uma camada de roteamento hierárquica compatível que torna a inferência de contexto longo viável sob as restrições de hardware atuais.