TileFuse: A Fused Mixed-Precision Kernel Library for Efficient Quantized LLM Inference on AMD NPUs

Resumo Técnico: TileFuse

Definição do Problema

A implantação de Grandes Modelos de Linguagem (LLMs) em dispositivos de borda depende cada vez mais de Unidades de Processamento Neural (NPUs) para atender a rigorosos orçamentos de potência e térmicos. No entanto, a inferência prática de LLMs em NPUs de clientes, especificamente a arquitetura AMD XDNA2, enfrenta barreções significativas. O principal obstáculo é a fragmentação das pilhas de software e a falta de programabilidade de baixo nível. A maioria dos frameworks de NPU (por exemplo, Apple Core ML, Intel OpenVINO, Qualcomm QNN) expõe apenas interfaces de alto nível ou caminhos de implantação proprietários que não suportam nativamente formatos de quantização amplamente utilizados e prontos para uso, como o AWQ (Activation-Weight Quantization).

Consequentemente, os desenvolvedores são frequentemente forçados a remodelar modelos para se ajustarem aos esquemas de quantização específicos da NPU ou a depender de uma execução de precisão total menos eficiente. Embora trabalhos anteriores tenham explorado a quantização consciente do hardware para plataformas específicas (por exemplo, ScalingLLM da Qualcomm), eles frequentemente falham em suportar formatos padrão como $W4A16$ ou não conseguem lidar com as grandes dimensões de matriz (por exemplo, >8K) comuns nos LLMs modernos. Além disso, as otimizações de NPU existentes muitas vezes carecem da capacidade de fundir o desempacotamento (unpacking) e a desquantização diretamente no kernel de computação, levando a um tráfego de memória desnecessário e movimentação de dados ineficiente.

Metodologia

O TileFuse é uma biblioteca de kernels de precisão mista próxima ao metal (close-to-metal) projetada para NPUs AMD XDNA2, visando as camadas lineares (GEMM e GEMV) que dominam a inferência de LLMs. O sistema utiliza o fluxo de programação aberto MLIR-AIE (IRON) para o co-design de layouts de pesos, posicionamento de metadados, microkernels e fluxos de dados em nível de array.

1. Microkernels de Precisão Mista Fundidos

O TileFuse funde três operações distintas em um único fluxo de kernel:

• Desempacotamento (Unpacking): Extração de pesos INT4/INT8 compactados.
• Desquantização (Dequantization): Conversão de inteiros de baixa bitagem para BF16 usando escalas e pontos zero por grupo ou canal.
• Execução: Realização da multiplicação de matrizes (GEMM) ou multiplicação de matriz-vetor (GEMV).

Ao manter os pesos em formato comprimido até que alcancem o núcleo de computação, o TileFuse evita a ida e volta (round-trip) na memória global necessária para materializar os pesos de precisão total (BF16). Os microkernels são projetados para desempacotar e desquantizar cada tile de peso uma única vez, armazenar os valores BF16 resultantes na memória local do núcleo de computação e reutilizá-los para operações de matriz subsequentes, preservando assim os padrões de reutilização essenciais para o desempenho de GEMM em tiles.

2. Pré-Tiling Intercalado

Para abordar as limitações das restrições de streaming de DMA em tempo de execução (que tipicamente limitam as dimensões suportadas abaixo de 8K), o TileFuse emprega uma estratégia de pré-tiling intercalado offline.

• Empacotamento de Metadados: Metadados de quantização (escalas e pontos zero) são empacotados diretamente nos tiles de peso na memória, garantindo que todos os dados necessários para a desquantização estejam disponíveis de forma contígua.
• Reorganização de Layout: A matriz de pesos é reorganizada em um layout intercalado de coluna principal (column-major). Isso garante que os tiles de peso atribuídos à mesma coluna AIE sejam contíguos na memória, eliminando grandes saltos de memória que, de outra forma, forçariam uma única operação de GEMM a ser dividida entre múltiplas invocações de kernel. Esta abordagem suporta dimensões de GEMM de até 32K.

3. Co-Design de Fluxo de Dados GEMV

A geração de tokens em LLMs de borda é tipicamente uma operação GEMV de tamanho de lote (batch size) 1, que costuma ser limitada pela memória em precisão total, mas torna-se limitada pelo processamento quando a quantização e a desquantização são introduzidas.

• Limitação de Base (Baseline): Mapeamentos de GEMV padrão costumam utilizar apenas uma linha do array de computação AIE de 4×8, deixando a maior parte do hardware ociosa.
• Otimização: O TileFuse redesenha o fluxo de dados para distribuir os pesos por todo o array 4×8. Ele utiliza núcleos de memória como uma camada de distribuição intermediária, transmitindo pacotes de tiles de peso de núcleos shim para núcleos de memória, que então distribuem para as quatro linhas de computação. Isso maximiza a utilização espacial e o rendimento aritmético para a geração de tokens.

4. Otimização Consciente de Reconfigurabilidade

O sistema aproveita explicitamente a reconfigurabilidade da arquitetura XDNA2 para personalizar os microkernels dos núcleos de computação e os padrões de acesso DMA para cargas de trabalho específicas. No entanto, os autores reconhecem que essa flexibilidade introduz overhead de despacho em tempo de execução e de reconfiguração de fabric. O TileFuse é otimizado para amortizar esse overhead durante os GEMMs de preenchimento (prefilling) de grande escala, mas reconhece suas limitações para operações GEMV de pequena escala e sensíveis à latência.

Principais Contribuições

1. Kernels Fundidos para Quantização Prática: Implementação de kernels AIE fundidos que consomem diretamente pesos no estilo AWQ ($W4A16$ e $W8A16$), eliminando a necessidade de remodelagem de modelos ou esquemas de quantização proprietários.
2. Pré-Tiling Intercalado: Uma estratégia de layout que suporta dimensões de GEMM de até 32K e permite o streaming eficiente em tempo de execução ao alinhar o acesso à memória com as restrições de descritores de buffer de DMA.
3. Redesenho de Fluxo de Dados GEMV: Um novo fluxo de dados que distribui os pesos por todo o array AIE 4×8, melhorando significativamente o rendimento da geração de tokens.
4. Eficiência de Borda Ponta a Ponta: Integração desses kernels em um runtime híbrido (NPU para camadas lineares, iGPU para atenção/camadas não lineares) que demonstra ganhos substanciais em latência e eficiência energética.

Resultados

As avaliações foram conduzidas em laptops Ryzen AI (Ryzen AI 7 350 e Ryzen AI 9 HX 370) usando modelos como Llama3-8B, Gemma-2B e Qwen2.5-3B.

• Desempenho de Kernel:
  • GEMM: O TileFuse alcança até 121,6% de melhoria de desempenho sobre as bases de precisão total.
  • GEMV: O TileFuse alcança até 281% de melhoria de desempenho sobre as bases de precisão total.
  • Comparação: Em GEMM, o TileFuse entrega mais de 2× ganhos de desempenho e eficiência energética sobre bases fortes de iGPU (kernels fundidos do llama.cpp).
• Inferência de LLM Ponta a Ponta:
  • Prefilling (Preenchimento): No Ryzen AI 7 350, o TileFuse alcança até 2,0× menor latência de prefilling (ex: 2,00× de aceleração para Llama3-8B com 2048 tokens) comparado à base de iGPU.
  • Energia: O caminho da NPU reduz o consumo de energia em mais de 64,6% no Ryzen AI 7 350 e 52,8% no Ryzen AI 9 HX 370.
  • Geração de Tokens: A NPU apresenta desempenho inferior à iGPU para geração de tokens devido ao overhead de milissegundos de despacho de kernel e reconfiguração em relação à execução de microssegundos dos operadores GEMV.
• Comparação com Trabalhos Anteriores: O TileFuse supera frameworks recentes baseados em NPU, como FastFlowLM e Scaling NPU, em throughput de prefilling (ex: acelerações de 1,45× e 1,61× respectivamente em Llama3-8B e Qwen2.5-3B).

Significância e Alegações

O artigo afirma que o TileFuse demonstra que as NPUs AMD XDNA2 são um alvo prático para a inferência eficiente de LLMs quantizados prontos para uso. Ao trazer formatos padrão como o AWQ diretamente para o hardware através de otimização próxima ao metal, este trabalho preenche a lacuna entre o potencial bruto da NPU e a implantação real em clientes.

Os autores enfatizam que o sucesso do TileFuse depende de uma estratégia de implantação híbrida: utilizando o array espacial reconfigurável da NPU para o prefilling de contexto grande (onde o overhead é amortizado) enquanto confia na iGPU para a geração de tokens sensível à latência. Esta abordagem valida que as NPUs podem melhorar substancialmente a usabilidade de LLMs na borda quando o hardware é adaptado aos formatos de quantização dependentes do usuário, em vez de forçar os modelos a se conformarem a restrições proprietárias. O trabalho estabelece que, com o devido co-design de kernel e fluxo de dados, as NPUs podem entregar benefícios significativos de desempenho e energia para IA de borda, mesmo na presença de fragmentação de pilha de software.