Resumo Técnico: CUDA-L2

Declaração do Problema

A multiplicação de matrizes (matmul) é uma operação fundamental em Grandes Modelos de Linguagem (LLMs), representando uma parte significativa do tempo de computação tanto no treinamento quanto na inferência. Embora existam kernels CUDA otimizados manualmente, alcançar o desempenho ideal em diversas dimensões de matriz ($M, N, K$) e arquiteturas de GPU continua sendo um desafio significativo. Diferentes tripletos de dimensões frequentemente exigem estratégias de otimização distintas que não se transferem bem entre hardwares. Além disso, mesmo para dimensões idênticas, as estratégias ideais variam com base em detalhes de configuração, como a precisão do acumulador (FP16 vs. FP32).

Abordagens automatizadas existentes, como aquelas que utilizam Grandes Modelos de Linguagem (LLMs) para geração de kernels, focaram primariamente em benchmarks com configurações únicas e fixas (ex: KernelBench). Permanece incerto como esses métodos se traduzem para ambientes de produção, onde devem competir contra bibliotecas altamente otimizadas e de código fechado, como o cuBLAS e o cuBLASLt da NVIDIA. Há uma falta de sistemas automatizados capazes de explorar sistematicamente o vasto espaço de configuração de Multiplicação de Matrizes Generalizada de Meia Precisão (HGEMM) para superar baselines otimizados por especialistas humanos.

Metodologia: CUDA-L2

Os autores propõem o CUDA-L2, um sistema que combina Grandes Modelos de Linguagem (LLMs) e Aprendizado por Reforço (RL) para otimizar automaticamente kernels CUDA de HGEMM. O sistema foi projetado para lidar com 1.000 configurações distintas, representando todas as combinações de valores de $M, N, K$ de $\{64, 128, 256, 512, 1024, 2048, 4096, 8192, 12288, 16384\}$.

Arquitetura do Sistema e Treinamento

O CUDA-L2 estende o framework anterior, CUDA-L1, com quatro principais melhorias técnicas:

1. Pré-treinamento Continuado: O modelo é pré-treinado em um corpus diversificado de código CUDA além dos benchmarks padrão, incluindo fontes do PyTorch, ATen, CUTLASS e tutoriais da NVIDIA. Descrições de instrução são geradas para fragmentos de código bruto, e o contexto de recuperação aumentada é adicionado para incorporar novos conhecimentos e características arquiteturais.
2. Treinamento de RL em Múltiplos Estágios: O treinamento progride de tipos de kernels gerais para domínios especializados (especificamente matmul).
   • RL de Kernel Geral: Treina o modelo em uma ampla gama de operações (álgebra linear, convolução, redução, etc.) usando a velocidade de execução como recompensa.
   • RL de HGEMM: Foca especificamente em kernels HGEMM com variações de configurações $M, N, K$.
3. Contexto Aprimorado com Métricas de Profiling: O contexto de RL inclui métricas abrangentes de profiling do NVIDIA Nsight Compute (NCU) (ex: throughput de memória, ocupação de SM, eficiência de cache) para permitir melhores decisões de otimização além do simples tempo de execução de ponta a ponta.
4. Design da Função de Recompensa: A recompensa para um kernel gerado é baseada no score médio de aceleração (speedup) através de iterações de teste. Para garantir robustez, a recompensa é penalizada por:
   • Desvio numérico em relação a um ground truth de CPU em FP32.
   • Comprimento excessivo de código (para incentivar a concisão).
   • O sistema impõe critérios rigorosos de correção, incluindo correspondências exatas para entradas binárias e desvio limitado contra baselines confiáveis (cuBLAS, cuBLASLt).

Protocolo de Avaliação

O sistema é avaliado sob dois cenários para refletir diferentes realidades de implantação:

• Modo Offline: Os kernels são executados consecutivamente sem intervalos de tempo, simulando throughput máximo.
• Modo Servidor: Os kernels são executados em intervalos aleatórios para simular inferência em tempo real, contabilizando a dinâmica térmica da GPU (períodos de resfriamento e clock boosting).

A avaliação compara o CUDA-L2 contra quatro baselines:

1. torch.matmul (dispatch padrão do PyTorch).
2. cuBLAS (usando seleção de layout ótima entre NN e TN).
3. cuBLASLt-heuristic (selecionando o algoritmo de ranking mais alto da API de heurística da NVIDIA).
4. cuBLASLt-AutoTuning (testando exaustivamente até 100 candidatos para encontrar o mais rápido).

Principais Contribuições e Técnicas Descobertas

O CUDA-L2 demonstra que o RL guiado por LLM pode descobrir implementações superiores ao explorar sistematicamente espaços de configuração impraticáveis para ajuste manual. O sistema descobre e aplica autonomamente várias técnicas de otimização avançadas:

• Seleção de Abstração: O modelo seleciona automaticamente entre intrínsecos WMMA brutos para matrizes pequenas e abstrações CuTe para operações de tiles mais complexas e grandes.
• Padding de Entrada: O sistema descobre estratégias onde as matrizes de entrada são preenchidas com zeros para dimensões não divisíveis (ex: padding de $M=8192$ para $8320$) para permitir tamanhos de tile ($B_M$) mais eficientes que seriam inválidos de outra forma, equilibrando o overhead de computação extra contra ganhos de performance.
• Sobreposição Avançada de Memória e Computação:
  • Fragmentos de Registrador com Double-Buffering: Alternar entre dois conjuntos de fragmentos de registrador para sobrepor o movimento de dados com a computação (execução ping-pong).
  • Prefetching Agressivo: Ir além do prefetching de etapa única para prefetching de múltiplas etapas (carregando várias iterações à frente) quando a pressão de registradores permite.
  • Escalonamento de Prefetch Escalonado: Intercalar o prefetching das matrizes A e B em torno da operação de multiply-accumulate de matriz (MMA) para maximizar o paralelismo ao nível de instrução.
  • Cópia Direta de Registrador para Shared-Memory: Eliminar tensores temporários intermediários no epílogo quando os layouts de registrador coincidem com os layouts de memória compartilhada, usando tipos de dados largos (ex: uint128_t) para reduzir operações de cópia.
  • Swizzling de Bloco: Reordenar a execução do bloco de threads para melhorar as taxas de acerto do cache L2, particularmente para grandes tamanhos de problema.

Resultados

O CUDA-L2 supera sistematicamente todos os principais baselines em 1.000 configurações em GPUs A100.

Performance em Modo Offline (Aceleração Média sobre Baselines):

• +22.0% sobre torch.matmul.
• +19.2% sobre cuBLAS (usando layout ótimo).
• +16.8% sobre cuBLASLt-heuristic.
• +11.4% sobre cuBLASLt-AutoTuning (o baseline mais forte, que testa exaustivamente 100 candidatos).

Performance em Modo Servidor:
As acelerações aumentam no cenário de servidor, provavelmente devido à capacidade do sistema de lidar melhor com a dinâmica térmica e efeitos de resfriamento de cache:

• +28.7% sobre torch.matmul.
• +26.0% sobre cuBLAS.
• +22.4% sobre cuBLASLt-heuristic.
• +15.9% sobre cuBLASLt-AutoTuning.

As taxas de vitória variam de 79.3% a 95.7% em todos os baselines, indicando que as melhorias são sistemáticas e não impulsionadas por outliers. O sistema também mostra que as acelerações são geralmente maiores para matrizes menores, onde a utilização da GPU é menor, enquanto matrizes maiores saturam o throughput da GPU, estreitando a lacuna com os baselines.

Significância

O artigo afirma que o CUDA-L2 demonstra que mesmo os kernels mais críticos para performance e altamente otimizados, como o HGEMM, podem ser melhorados através da automação de RL guiada por LLM. Ao escalar a otimização através de espaços de configuração que são difíceis de ajustar manualmente, o sistema alcança ganhos de performance que superam bibliotecas de código fechado de última geração.

Embora a versão atual seja limitada às arquiteturas A100 (Ampere), o framework é projetado para ampla aplicabilidade. Os autores observam que o trabalho contínuo visa estender o sistema para outras arquiteturas de GPU, incluindo Ada Lovelace, Hopper e Blackwell. O trabalho sugere que métodos automatizados podem explorar sistematicamente vastos espaços de configuração para encontrar otimizações que o ajuste manual ou bibliotecas baseadas em heurísticas podem perder.