CUDA-L1: Melhorando a Otimização CUDA por meio de Aprendizado por Reforço Contrastivo

CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning

July 18, 2025

Autores: Xiaoya Li, Xiaofei Sun, Albert Wang, Jiwei Li, Chris Shum

cs.AI

Resumo

O crescimento exponencial na demanda por recursos de computação em GPU, impulsionado pelo rápido avanço dos Modelos de Linguagem de Grande Escala (LLMs), criou uma necessidade urgente por estratégias automatizadas de otimização CUDA. Embora os avanços recentes em LLMs mostrem potencial para geração de código, os modelos atuais de estado da arte (SOTA, por exemplo, R1, o1) alcançam taxas de sucesso baixas na melhoria da velocidade CUDA. Neste artigo, introduzimos o CUDA-L1, uma estrutura de aprendizado por reforço automatizada para otimização CUDA. O CUDA-L1 alcança melhorias de desempenho na tarefa de otimização CUDA: treinado em NVIDIA A100, ele proporciona um aumento médio de velocidade de 17,7x em todos os 250 kernels CUDA do KernelBench, com picos de aceleração chegando a 449x. Além disso, o modelo também demonstra excelente portabilidade entre arquiteturas de GPU, alcançando aumentos médios de velocidade de 17,8x em H100, 19,0x em RTX 3090, 16,5x em L40, 14,7x em H800 e 13,9x em H20, apesar de ter sido otimizado especificamente para A100. Além desses resultados de benchmark, o CUDA-L1 demonstra várias propriedades notáveis: 1) Descobre uma variedade de técnicas de otimização CUDA e aprende a combiná-las estrategicamente para alcançar desempenho ideal; 2) Revela princípios fundamentais de otimização CUDA; 3) Identifica gargalos de desempenho não óbvios e rejeita otimizações aparentemente benéficas que prejudicam o desempenho. As capacidades do CUDA-L1 demonstram que o aprendizado por reforço pode transformar um LLM inicialmente de baixo desempenho em um otimizador CUDA eficaz apenas com sinais de recompensa baseados em aceleração, sem expertise humana ou conhecimento de domínio. Mais importante, o modelo de RL treinado estende as habilidades de raciocínio adquiridas para novos kernels. Esse paradigma abre possibilidades para a otimização automatizada de operações CUDA e promete promover substancialmente a eficiência de GPUs e aliviar a crescente pressão sobre os recursos de computação em GPU.

English

The exponential growth in demand for GPU computing resources, driven by the rapid advancement of Large Language Models, has created an urgent need for automated CUDA optimization strategies. While recent advances in LLMs show promise for code generation, current SOTA models (e.g. R1, o1) achieve low success rates in improving CUDA speed. In this paper, we introduce CUDA-L1, an automated reinforcement learning framework for CUDA optimization. CUDA-L1 achieves performance improvements on the CUDA optimization task: trained on NVIDIA A100, it delivers an average speedup of x17.7 across all 250 CUDA kernels of KernelBench, with peak speedups reaching x449. Furthermore, the model also demonstrates excellent portability across GPU architectures, achieving average speedups of x17.8 on H100, x19.0 on RTX 3090, x16.5 on L40, x14.7 on H800, and x13.9 on H20 despite being optimized specifically for A100. Beyond these benchmark results, CUDA-L1 demonstrates several remarkable properties: 1) Discovers a variety of CUDA optimization techniques and learns to combine them strategically to achieve optimal performance; 2) Uncovers fundamental principles of CUDA optimization; 3) Identifies non-obvious performance bottlenecks and rejects seemingly beneficial optimizations that harm performance. The capabilities of CUDA-L1 demonstrate that reinforcement learning can transform an initially poor-performing LLM into an effective CUDA optimizer through speedup-based reward signals alone, without human expertise or domain knowledge. More importantly, the trained RL model extend the acquired reasoning abilities to new kernels. This paradigm opens possibilities for automated optimization of CUDA operations, and holds promise to substantially promote GPU efficiency and alleviate the rising pressure on GPU computing resources.