🤖 AI

TileMix: Tile-Centric Mixed-Precision Attention for LLM Inference Acceleration

O TileMix é um kernel de precisão mista centrado em tiles e livre de treinamento que acelera a inferência de LLMs de contexto longo ao rotear dinamicamente tiles de escores de atenção alinhados ao hardware entre caminhos FP16 e INT8 dentro de uma operação de atenção densa fundida, recuperando assim a precisão perdida por métodos de baixa precisão uniformes enquanto melhora o throughput.

Autores originais: Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang, Kewei Sha, Yunhe Feng

Publicado 2026-08-19

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang, Kewei Sha, Yunhe Feng

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Os grandes modelos de linguagem tornaram-se os motores por trás de uma nova geração de inteligência artificial, capazes de resumir livros inteiros, responder a perguntas complexas de documentos longos e manter conversas que abrangem milhares de palavras. Para fazer isso, esses modelos dependem de um mecanismo chamado atenção, que permite que eles pesem a importância de cada palavra em uma frase em relação a todas as outras palavras. Quando o texto é curto, esse processo é rápido. Mas conforme o contexto cresce para incluir capítulos inteiros ou contratos jurídicos, o custo computacional explode. O modelo deve calcular uma pontuação para cada par possível de palavras, criando uma grade massiva de dados que exige enormes quantidades de memória e poder de processamento. Esse gargalo tornou difícil executar essas ferramentas poderosas de forma eficiente em hardware padrão, especialmente ao lidar com as sequências longas exigidas para tarefas do mundo real, como análise jurídica ou revisão de prontuários médicos.

Pesquisadores tentaram resolver isso simplificando a matemática. Uma abordagem comum é reduzir a precisão dos números que o computador utiliza, mudando de cálculos de alta precisão para outros de baixa precisão mais rápidos. Isso é como mudar de medir ingredientes com uma balança de laboratório para usar uma colher de cozinha; é muito mais rápido, mas se você fizer isso para cada etapa de uma receita complexa, o prato final pode ter um sabor errado. Outra abordagem é pular cálculos inteiros, ignorando palavras que parecem sem importância. Embora isso economize tempo, corre-se o risco de cortar justamente as conexões que o modelo precisa para entender a história. O desafio tem sido encontrar uma maneira de usar a velocidade da matemática de baixa precisão sem sacrificar a precisidade da matemática de alta precisão, mantendo ao mesmo tempo a capacidade do modelo de ver o quadro completo.

Uma equipe de pesquisadores introduziu um novo método chamado TileMix, que aborda esse problema tratando o processo de atenção não como um bloco único e uniforme de trabalho, mas como uma coleção de pequenos blocos gerenciáveis (tiles). Imagine a enorme grade de pontuações de pares de palavras como um grande mosaico. Em vez de pintar todo o mosaico com um único tipo de tinta, o TileMix divide essa grade em pequenas seções quadradas alinhadas ao hardware. Para cada uma dessas seções, o sistema toma uma decisão de milésimos de segundo: este grupo específico de pares de palavras precisa do cálculo lento de alta precisão ou pode se dar bem com o de baixa precisão mais rápido? Essa decisão é tomada com base em um mapa pré-planejado que agrupa esses blocos, permitindo que o computador alterne entre alta e baixa precisão dentro da mesma operação sem parar para reorganizar os dados.

O núcleo desta inovação é um sistema de roteamento que atua como um controlador de tráfego para o processador do computador. Ele compacta essas decisões em um código conciso, essencialmente uma sequência de bits que diz ao processador qual caminho seguir para cada bloco. Se um bloco for marcado para alta precisão, o processador usa suas unidades matemáticas mais precisas. Se for marcado para baixa precisão, ele muda para suas unidades mais rápidas e eficientes em termos de energia. Crucialmente, ambos os caminhos atualizam um estado de memória compartilhado que monitora os resultados globais, garantindo que a saída final permaneça consistente. Isso permite que o sistema preserve a conectividade total do modelo — o que significa que nenhuma interação entre palavras é jamais ignorada — enquanto ainda colhe os benefícios de velocidade da matemática de baixa precisão para as partes do cálculo que podem tolerá-la.

Em seus experimentos, os pesquisadores testaram este método em vários modelos de linguagem de grande escala populares, incluindo LLaMA, Qwen e Vicina, usando sequências variando de 16.000 a 64.000 palavras. Eles compararam sua abordagem com o método padrão de alta precisão e uma versão que utilizava matemática de baixa precisão para tudo. Os resultados mostraram que usar matemática de baixa precisão para todo o cálculo frequentemente levava a uma queda significativa na precisão, fazendo com que o modelo perdesse detalhes ou falhasse em tarefas de recuperação de informações. No entanto, o TileMix foi capaz de recuperar a maior parte dessa qualidade perdida. Ao rotear cuidadosamente apenas grupos específicos de blocos para o caminho de alta precisão, o sistema manteve níveis de precisão muito próximos da linha de base de alta precisão total, enquanto simultaneamente alcançava velocidades de processamento muito maiores.

O estudo também explorou diferentes padrões para decidir quais blocos deveriam receber o tratamento de alta precisão. Eles descobriram que o arranjo importava; alguns padrões, que mantinham os cálculos de alta precisão em regiões espaciais específicas da grade de palavras, funcionavam melhor do que outros dependendo da tarefa. Por exemplo, certos layouts que preservavam a alta precisão para interações locais de palavras funcionaram melhor em tarefas de recall factual, enquanto outros foram mais robustos para respostas a perguntas gerais. Os pesquisadores demonstraram que este método pode ser aplicado sem retreinar os modelos, o que significa que pode ser implementado imediatamente em sistemas existentes. Eles também mostraram que a abordagem funciona bem com lotes de comprimentos variáveis e diferentes arquiteturas de modelos, sugerindo que é uma ferramenta flexível para melhorar a eficiência.

Fundamentalmente, o TileMix oferece um equilíbrio controlável entre velocidade e precisão. Ele sugere que o futuro do processamento de contexto longo eficiente não exige escolher entre ser rápido ou ser inteligente. Em vez disso, ao misturar inteligentemente os dois, os sistemas podem processar documentos longos a uma velocidade que se aproxima dos limites teóricos do hardware, sem as severas penalidades de precisão que vêm do uso de baixa precisão em todo lugar. Esta abordagem fornece um caminho prático para a implementação de modelos de linguagem de grande escala em cenários onde tanto a velocidade quanto a precisão são críticas, como na análise em tempo real de conjuntos de dados massivos ou ferramentas interativas que precisam compreender contextos longos e complexos instantaneamente.

Resumo Técnico: TileMix: Atenção de Precisão Mista Centrada em Tiles para Aceleração de Inferência de LLMs

Declaração do Problema

A inferência de contexto longo em Grandes Modelos de Linguagem (LLMs) enfrenta um gargalo significativo durante a fase de prefill devido à complexidade computacional quadrática ( $O(L^2)$ ) da autoatenção densa. As estratégias de aceleração existentes geralmente se enquadram em três categorias, cada uma com limitações:

Quantização de Baixa Precisão: Métodos como a quantização INT8 melhoram a eficiência aritmética e de memória, mas normalmente aplicam um caminho de precisão uniforme em toda a matriz de atenção. Isso frequentemente deixa o roteamento de precisão espacial fora do loop de atenção fundido, e a execução uniforme de baixa precisão pode degradar a qualidade do contexto longo.
Métodos Baseados em Esparsidade: Abordagens que selecionam interações de tokens ativos (ex: janelas deslizantes, padrões esparsos) reduzem a computação, mas alteram a conectividade densa do grafo de atenção, podendo perder interações críticas de tokens.
Atenção Fundida Consciente de IO: Kernels como o FlashAttention particionam a atenção em tiles alinhados ao hardware para otimizar o movimento de dados e fundir operações (computação de score, softmax online, agregação de valor). No entanto, esses kernels tradicionalmente mantêm um caminho de computação de score uniforme, perdendo a oportunidade de rotear regiões espaciais específicas para diferentes níveis de precisão dentro de uma única execução.

O desafio central abordado é como integrar precisão heterogênea (FP16 e INT8) em um único kernel de atenção fundido sem quebrar a conectividade densa de tokens ou incorrer em um overhead significativo de metadados, garantindo que diferentes caminhos aritméticos convirjam corretamente para um estado de softmax online compartilhado.

Metodologia: TileMix

O TileMix introduz um kernel de roteamento de precisão centrado em tiles que trata a precisão numérica como uma decisão espacial sobre grupos de tiles de score dentro de uma computação de atenção densa fundida.

Mecanismos Principais

Roteamento de Precisão por Grupo de Tiles:
- A matriz de atenção ( $L_q \times L_k$ ) é particionada em tiles de computação alinhados ao hardware de tamanho $BLOCK_M \times BLOCK_N$ .
- Em vez de atribuir precisão por tile, o TileMix agrupa tiles de chaves adjacentes ao longo da dimensão da chave. Um único bit de roteamento governa um grupo de $g$ tiles de chave (onde $g$ é um fator de agrupamento).
- Isso permite que o kernel despache cada grupo de tiles para um caminho FP16 (multiplicação de matriz padrão) ou um caminho INT8 (usando Tensor Core MMA INT8 com acumulação INT32) dentro do mesmo loop de streaming.
Roteamento por Máscara de Bits Compacta:
- As decisões de roteamento são compactadas em inteiros de 64 bits (máscaras de bits), onde cada bit corresponde a um grupo de tiles de chave.
- Isso permite uma busca de tempo constante dentro do loop interno usando operações de deslocamento e máscara ( $R_{hk,m,g_j} = ((b_{hk,m} \gg g_j) \& 1)$ ).
- O overhead de metadados é $O(H_k T_m)$ , onde $H_k$ é o número de cabeças KV e $T_m$ é o número de tiles de query, permanecendo compacto mesmo conforme o comprimento do contexto cresce.
Execução Heterogênea de Estado Compartilhado:
- Ambos os caminhos FP16 e INT8 atualizam um estado de softmax online compartilhado (máximo por linha, normalizador e acumulador de saída não normalizado).
- O caminho INT8 computa $Q_8 K_8^\top$ , aplica escalas de bloco e reescala o resultado para entrar em um domínio de ponto flutuante comum antes de atualizar o estado compartilhado.
- Este design preserva o grafo de atenção densa (todas as interações legais de tokens são computadas) enquanto varia a precisão aritmética com base em políticas de roteamento espacial.
Políticas e Layouts:
- As políticas de roteamento são definidas por templates estruturados estáticos e sem dados (ex: Band, Global, Row-Random, Aligned Sparse, BigBird, Sparse Transformer).
- Esses templates determinam quais regiões espaciais da matriz de atenção são roteadas para FP16 (alta precisão) e quais para INT8 (baixa precisão) com base em uma razão de cobertura INT8 alvo (ex: 25%, 50%, 75%).
- O sistema suporta Grouped-Query Attention (GQA), batching de comprimento variável e caches de chave/valor (KV) em INT8.

Contribuições Principais

Roteamento de Precisão Centrado em Tiles: O artigo introduz a precisão de grupo de tiles como uma abstração de execução espacial para atenção densa fundida, permitindo a alocação fina de FP16/INT8 em todas as interações legais de tokens sem esparsificar o grafo.
Execução de Score Heterogênea de Estado Compartilhado: Um kernel fundido é projetado para alinhar os caminhos de score FP16 e INT8 a um domínio de score comum, integrando-os através de uma única recorrência de softmax online.
Roteamento de Kernel Nativo, Compacto e Escalável: Os autores desenvolvem roteamento por máscara de bits compactada com busca de tempo constante no loop interno e agrupamento de precisão escalável que mantém tiles de computação alinhados ao hardware enquanto suporta contextos longos.
Inferência Prática de Contexto Longo: A implementação suporta GQA, batching de comprimento variável e caches KV em INT8, validados em múltiplas famílias de modelos (LLaMA, Qwen, Vicina) e benchmarks.

Resultados Experimentais

A avaliação foi conduzida em GPUs NVIDIA A100 usando LLaMA 3.2 3B, Qwen 2/2.5 7B e Vicuna 7B.

Recuperação de Precisão: A quantização INT8 uniforme ("One") frequentemente degrada o desempenho de recuperação e de resposta a perguntas à medida que o contexto aumenta. O TileMix recupera essa qualidade perdida ao manter a precisão FP16 para grupos de tiles selecionados. Por exemplo, no benchmark LongEval, configurações de roteamento misto (como layouts SpTrans ou BigBird) superaram consistentemente o INT8 uniforme, muitas vezes aproximando-se ou igualando os baselines de FP16 enquanto mantinham alta cobertura INT8.
Eficiência de Throughput: O TileMix melhora o throughput de prefill em relação ao FlashAttention padrão em FP16. A 4k tokens, configurações como SpTrans75 (75% INT8) alcançaram ~31.8 K tokens/s comparado a ~14.3 K tokens/s para o FlashAttention. As configurações do TileMix geralmente superaram ou igualaram os baselines de INT8 uniforme em throughput de ponta a ponta, demonstrando que o overhead de roteamento e manipulação de precisão mista é insignificante comparado aos ganhos dos Tensor Cores INT8.
Estabilidade Numérica: A análise do desvio absoluto médio em relação a uma referência FP16 fixa mostrou que o desvio aumenta com a cobertura INT8, mas permanece controlável. A escolha do layout de roteamento (distribuição espacial de FP16) impacta significativamente a precisão, com certos layouts (como SpTrans) mostrando robustez em diferentes modelos e tarefas.
Proteção de Heavy-Hitters: Layouts de roteamento estáticos como SpTrans mostraram proteger naturalmente as interações de atenção de alta importância (heavy hitters) de serem roteadas para INT8, enviando apenas uma pequena fração das interações de alta massa para baixa precisão, apesar da alta cobertura nominal de INT8.

Significância e Alegações

O artigo afirma que o TileMix estabelece o roteamento de precisão por grupo de tiles de score como uma abstração nativa de kernel para atenção de contexto longo. Sua principal significância reside em:

Desacoplamento de Conectividade e Precisão: Permite a preservação da conectividade densa de tokens (ao contrário da atenção esparsa) enquanto possibilita a execução de precisão heterogênea (ao contrário da quantização uniforme).
Fronteira Controlável de Precisão-Eficiência: Fornece um mecanismo prático para trocar precisão por eficiência ao ajustar a razão de cobertura INT8 e o layout de roteamento, ofereando uma "fronteira controlável de precisão-eficiência" através de diversas famílias de modelos.
Implantação Sem Retreinamento: O método opera sem necessidade de retreinamento, utilizando templates espaciais estáticos para guiar o roteamento de precisão.
Otimização ao Nível de Kernel: Ao integrar o roteamento de precisão diretamente no kernel de atenção fundida com busca de tempo constante, evita os custos de latência associados à seleção de precisão em nível de invocação ou agendamento externo.

Os autores concluem que o controle de precisão espacial é uma dimensão prática e eficaz para otimizar kernels de atenção densa, recuperando a qualidade de contexto longo perdida sob métodos de baixa precisão uniforme, enquanto acelera significativamente a inferência.