Resumo Técnico: Sobreposição de Ativação em Nível de Tile para Inferência Eficiente de LLM

Definição do Problema

Os Modelos de Linguagem de Grande Escala (LLMs) modernos, incluindo Qwen-2.5, LLaMA, Mistral e Gemma, utilizam predominantemente a função de ativação SwiGLU. Esta estrutura de MLP com portão (gated) computa $Y = \text{SiLU}(A \times W_1) \odot (A \times W_2)$, exigindo duas multiplicações de matrizes (GeMMs) independentes seguidas por uma ativação de portão elemento a elemento.

Em implementações padrão, este processo materializa dois tensores intermediários completos ($Gate$e$Up$) na Memória de Alta Largura de Banda (HBM) entre as etapas de GeMM e ativação. O profiling em GPUs NVIDIA H100 revela que esta materialização intermediária e o processamento associado da ativação consomem de 9% a 37% do tempo total de execução do MLP. O overhead é mais severo para modelos menores (ex: Qwen-2.5 0.5B), onde o SwiGLU representa mais de 30% do tempo de MLP. À medida que a densidade aritmética dos Tensor Cores aumenta via quantização (FP8, INT4), o custo relativo destas operações limitadas por memória cresce, criando um gargalo significativo.

A infraestrutura de compiladores existente, especificamente o torch.compile do PyTorch, falha em abordar isso. Ele não consegue fundir operações através de dois GeMMs separados com matrizes de pesos diferentes devido a limitações fundamentais em passes de fusão ao nível de grafo. Experimentos mostram que o torch.compile alcança apenas 34–94% do desempenho do PyTorch eager para este padrão, com dicas de fusão explícitas fornecendo melhorias negligenciáveis.

Metodologia

Os autores apresentam dois kernels baseados em CUTLASS SM90 complementares e feitos à mão que fundem o SwiGLU no GeMM ao nível de tile, eliminando a materialização intermediária em HBM. Ambos os kernels utilizam o Tensor Memory Accelerator (TMA) e o escalonamento especializado em warps disponível no H100 (SM90).

1. Kernel-1: Sobreposição de Escalonamento Pingpong para Sincronização de SwiGLU

• Mecanismo: Este kernel utiliza o escalonamento especializado em warps Pingpong do CUTLASS, que particiona os warps em produtores (carregando dados via TMA) e consumidores (executando MMA).
• Estratégia: Ele cria uma grade de threadblocks $[M, N]$. Dentro da fase do consumidor, após o tile do Gate ser computado, mas antes que os dados do tile Up cheguem, há uma breve janela onde os warps consumidores ficam ociosos. O kernel explora esta janela para computar a ativação Swish ($\text{SiLU}(Gate)$) no acumulador do Gate residente nos registradores.
• Sobreposição: A computação Swish sobrepõe-se temporalmente com os carregamentos TMA dos tiles de peso Up do produtor. A multiplicação elemento a elemento final ($\text{SiLU}(Gate) \odot Up$) e o armazenamento ocorrem no epílogo.
• Otimização: Este design é otimizado para grandes tamanhos de lote ($M \ge 2048$), onde a grade $[M, N]$ satura suficientemente os Streaming Multiprocessors (SMs) da GPU.

2. Kernel-2: Intercalação de SwiGLU via Árvore de Visitante de Epílogo Customizada

• Mecanismo: Este kernel introduz um nó de Árvore de Visitante de Epílogo (EVT) customizado chamado PairMulStore.
• Estratégia: Ele concatena as matrizes de peso $W_1$ e $W_2$ em uma única matriz fundida $W_{fused} \in \mathbb{R}^{K \times 2N}$ com colunas intercaladas (colunas pares para Gate, ímpares para Up). O kernel lança um GeMM padrão sobre $[M, 2N, K]$, criando uma grade de threadblocks $[M, 2N]$.
• Fusão: O nó PairMulStore opera no acumulador $[M, 2N]$ dentro de cada threadblock. Ele lê pares de colunas adjacentes $(D[m, 2n], D[m, 2n+1])$, aplica $\text{SiLU}(D[m, 2n]) \odot D[m, 2n+1]$ e armazena o resultado.
• Sobreposição: Este design intercala a computação SwiGLU com os armazenamentos de tile. A aritmética do tile atual sobrepõe-se com os armazenamentos TMA assíncronos do tile anterior, proporcionando sobreposição entre computação e armazenamento (compute-store overlap).
• Otimização: O aumento de $2\times$ no número de threadblocks melhora significativamente a ocupação do SM para pequenos tamanhos de lote ($M \le 2048$), tornando-o o padrão preferencial para a maioria das configurações.

Principais Contribuições

1. Fusão ao Nível de Tile: A primeira implementação que funde o SwiGLU no GeMM ao nível de tile, eliminando 4 operações intermediárias de HBM (leitura/escrita de Gate e Up) por etapa de inferência.
2. Estratégia de Dual-Kernel: Uma abordagem complementar usando o Kernel-1 para grandes lotes (sequencial Gate $\to$ Swish $\to$ Up) e o Kernel-2 para pequenos lotes (pesos intercalados com EVT customizado), abordando restrições de ocupação e escalonamento de formas distintas.
3. Superioridade Numérica: Os kernels fundidos alcançam zero discrepâncias (0% de elementos com erro relativo >1%) em comparação com uma referência FP32, enquanto a implementação baseline do cuBLAS exibe 4,5–11% de discrepâncias devido à acumulação parcial agressiva em BF16 entre os tiles.
4. Validação da Limitação do Compilador: Evidência empírica demonstrando que o torch.compile não consegue replicar esta fusão, permanecendo de 3 a 7 vezes mais lento que os kernels hand-crafted.

Resultados

Avaliado em modelos Qwen-2.5 (0.5B a 72B) no NVIDIA H100:

• Aceleração (Speedup): Os kernels alcançam até 2.47× de aceleração sobre o PyTorch eager (cuBLAS + SwiGLU separado). O Kernel-2 vence em 13/20 configurações testadas, enquanto o Kernel-1 vence em 7/20 (principalmente para modelos grandes em altos tamanhos de lote).
• Mudança de Regime: A análise de Roofline confirma que a fusão desloca as cargas de trabalho de limitadas por memória para limitadas por computação. Para modelos pequenos, a intensidade aritmética aumenta em até 247%, movendo o ponto de operação através da crista do roofline.
• Utilização: A utilização de pico de BF16 atinge 79,5% no H100, quase dobrando a utilização do baseline para modelos pequenos.
• Eficiência: A eficiência de fusão (aceleração real vs. economia teórica de memória) é próxima de 100% para modelos pequenos. Para o Kernel-2, a eficiência excede 100% em várias configurações devido aos benefícios da sobreposição de computação-armazenamento.
• Escalabilidade: A aceleração diminui com o tamanho do modelo (2.47× para 0.5B vs. ~1.0× para 72B em altos tamanhos de lote) porque a computação GeMM domina o tempo total, e a fração relativa de tráfego de memória intermediária diminui.

Significância

O artigo argumenta que, para a inferência de LLMs modernos, particularmente em modelos de borda (edge) e de tamanho móvel (0.5B–14B), o gargalo mudou do processamento bruto para o tráfego de memória causado pela materialização de tensores intermediários. Os autores demonstram que o design hand-crafted e consciente do hardware é necessário para superar as limitações da infraestrutura de compiladores atual. Ao eliminar acessos intermediários à HBM e sobrepor a computação de ativação com a movimentação de dados, estes kernels reduzem significativamente a latência de inferência e melhoram a utilização de computação, oferecendo uma otimização crítica para sistemas de serviço de produção onde a latência e o throughput são limitados.