MISA: Mixture of Indexer Sparse Attention for Long-Context LLM Inference

Resumo Técnico: MISA (Mistura de Atenção Esparsa com Indexador)

Declaração do Problema
À medida que os modelos de linguagem de grande escala (LLMs) evoluem para processar contextos de centenas de milhares a milhões de tokens, os mecanismos de atenção densa tornam-se computacionalmente proibitivos. Embora técnicas de atenção esparsa tenham surgido para mitigar isso, a Atenção Esparsa DeepSeek (DSA) representa atualmente o estado da arte para atenção esparsa fina durante a inferência. A DSA emprega um "indexador" aprendido que pontua cada token de prefixo usando um grande conjunto de $H_I$ cabeças (por exemplo, 64 no DeepSeek-V3.2) para selecionar os top-$k$ tokens para o mecanismo de atenção principal.

Embora o design de múltiplas cabeças do indexador da DSA seja crucial para a expressividade — permitindo que diferentes cabeças se especializem em padrões distintos de relevância (recenticidade, sintaxe, semântica) —, ele cria um gargalo significativo. O custo computacional de pontuar cada token com todas as $H_I$ cabeças escala como $O(H_I L)$, tornando o indexador o custo dominante durante a inferência de contexto longo. Melhorias existentes como HISA (busca hierárquica de bloco para token) e IndexCache (reutilização de índice entre camadas) reduzem custos operando no eixo do token ou no eixo da camada, mas ainda exigem que todas as $H_I$ cabeças permaneçam ativas dentro do kernel, não abordando diretamente o gargalo do eixo das cabeças.

Metodologia: MISA
Os autores propõem MISA (Mistura de Atenção Esparsa com Indexador), uma substituição direta para o indexador da DSA que trata as $H_I$ cabeças do indexador como um conjunto de Mistura de Especialistas (MoE). Em vez de ativar todas as cabeças para cada consulta, a MISA roteia um subconjunto dependente da consulta de $h \ll H_I$ cabeças ativas.

O mecanismo central envolve duas etapas:

1. Roteador Leve Agrupado por Blocos: O prefixo é partitionado em $M = \lceil L/B \rceil$ blocos. O roteador computa afinidades por cabeça usando estatísticas de nível de bloco agrupadas (chaves de indexação com média agrupada) em vez de tokens individuais. Ele agrega essas pontuações para estimar a importância de cada cabeça para a consulta atual e seleciona as top-$h$ cabeças. Esta etapa opera em $M \ll L$ chaves, resultando em uma sobrecarga negligenciável de $O(H_I M)$.
2. Pontuação Esparsa de Tokens: Apenas as $h$ cabeças ativas selecionadas realizam a pesada pontuação em nível de token no prefixo completo. Os tokens top-$k$ finais são selecionados com base nessas pontuações reduzidas.

Esta abordagem reduz a complexidade do indexador por consulta de $O(H_I L)$ para $O(hL + H_I M)$. Como $M \ll L$ e $h \ll H_I$, as economias computacionais são substanciais, preservando a diversidade total do conjunto do indexador (todas as cabeças permanecem disponíveis; o roteador simplesmente escolhe quais consultar).

Variante Hierárquica (MISA†)
Para fechar ainda mais a lacuna entre eficiência e a seleção exata do indexador denso da DSA, os autores introduzem MISA†, um pipeline de grosso para fino:

• Etapa Grossa: A passagem roteada por MoE seleciona um conjunto de candidatos ampliado de $k'$ tokens (por exemplo, $k'=8192$).
• Etapa Fina: O indexador original da DSA (usando todas as $H_I$ cabeças) re-pontua apenas este conjunto de candidatos para recuperar os tokens top-$k$ finais.
  Isso permite que a MISA† recupere as seleções do indexador denso quase exatamente, enquanto ainda se beneficia da computação reduzida da etapa grossa.

Principais Contribuições

1. Roteamento no Eixo das Cabeças: O artigo identifica os produtos cabeça-token por token como o custo dominante na DSA e introduz o roteamento no eixo das cabeças como um eixo de eficiência complementar às hierarquias existentes no eixo do token.
2. Arquitetura MISA: Um indexador roteado por MoE de substituição direta que ativa apenas $h$ cabeças por consulta via um roteador leve agrupado por blocos, reduzindo a complexidade sem sacrificar a diversidade do conjunto de cabeças.
3. Extensão MISA†: Uma variante hierárquica que usa a passagem roteada para podar o espaço de busca, seguida por uma reclassificação completa da DSA, alcançando uma recuperação quase perfeita das seleções do indexador denso.
4. Eficiência e Desempenho: O método alcança acelerações significativas sem exigir treinamento adicional.

Resultados Experimentais
Avaliado no DeepSeek-V3.2 ($H_I=64$) e GLM-5 ($H_I=32$) sem treinamento adicional:

• LongBench: A MISA (com $h=8$) iguala o indexador denso da DSA dentro de 0,20 pontos médios no DeepSeek-V3.2 e supera-no no GLM-5. A MISA† reduz ainda mais a lacuna para dentro de 0,1 pontos no DeepSeek-V3.2. Ambas superam as linhas de base Block-Sparse e HISA em média.
• Agulha no Palheiro (NIAH): A MISA e a MISA† preservam mapas de calor de recuperação totalmente verdes até contexto de 128K, igualando o desempenho do indexador denso da DSA. Em contraste, o Block-Sparse mostra lacunas de precisão, e a HISA mostra degradações menores em profundidades extremas.
• Recuperação de Tokens: A MISA† recupera mais de 92% dos tokens selecionados pelo indexador completo da DSA por camada (medido por IoU).
• Aceleração: Uma implementação de kernel TileLang alcança aproximadamente uma aceleração de 3,82× sobre o kernel original do indexador da DSA em uma única GPU NVIDIA H200 para contextos longos.

Significado e Alegações
O artigo afirma que o roteamento no eixo das cabeças do indexador é um eixo prático e complementar de eficiência para atenção esparsa fina. Ao desacoplar a decisão de roteamento (quais cabeças usar) da seleção de tokens (a quais tokens atender), a MISA demonstra que é possível reduzir drasticamente o custo computacional do indexador (em 8× no DeepSeek-V3.2 e 4× no GLM-5) enquanto mantém a qualidade de recuperação e a expressividade do indexador denso completo. Os autores posicionam isso como uma melhoria fundamental que pode ser combinada ortogonalmente com qualquer esquema existente de atenção esparsa em nível de token.

Limitações
Os autores reconhecem que seus experimentos de velocidade medem a latência do kernel em vez da latência de ponta a ponta do modelo e que o método não reduz o volume de acesso à memória ao cache KV. Além disso, os resultados são obtidos inserindo a MISA em modelos pré-treinados sem ajuste fino conjunto do roteador, sugerindo potencial para ganhos adicionais através do treinamento.