TorchInductor ganha suporte a GEMMs de alta performance com CuteDSL

O time do PyTorch anunciou a integração do CuteDSL como um novo backend de autotuning para multiplicações de matrizes (GEMMs) no TorchInductor. A novidade promete melhorar significativamente o desempenho de modelos de linguagem de grande porte em GPUs NVIDIA, especialmente em operações de atenção e camadas fully-connected. Segundo os desenvolvedores, o CuteDSL se destaca por oferecer controle de baixo nível sobre o hardware semelhante ao CUTLASS C++, mas com a vantagem de compilação mais rápida e manutenção simplificada. Enquanto backends como Triton já entregam bons resultados para operações pointwise e reduções, o CuteDSL vem para preencher uma lacuna crítica: otimizar GEMMs em arquiteturas modernas como Hopper e Blackwell.

A escolha do CuteDSL como novo backend não foi aleatória. O time do PyTorch estabeleceu três critérios rigorosos para integrar um novo sistema de autotuning: (1) não impor sobrecarga significativa de manutenção; (2) não degradar os tempos de compilação ou benchmark em relação aos backends existentes; e (3) entregar melhor desempenho em cargas de trabalho-alvo. O CuteDSL atendeu a todos esses requisitos, graças ao desenvolvimento ativo da NVIDIA, que fornece templates otimizados de kernels e mantém a complexidade sob controle. Além disso, a compilação via CuteDSL é tão rápida quanto a dos outros backends do TorchInductor, superando a abordagem CUTLASS C++ que exige compilações completas com nvcc para cada variante de kernel.

A integração do CuteDSL representa também um investimento estratégico de longo prazo. Como a NVIDIA continua a investir no desenvolvimento da ferramenta, espera-se que ela eventualmente substitua o backend CUTLASS C++ em gerações futuras de hardware. Isso simplificaria consideravelmente a base de código do TorchInductor, consolidando toda a lógica de GEMMs em um único sistema baseado em Python. A combinação de incentivos alinhados, adoção crescente no ecossistema open-source (como a biblioteca Quack de Tri Dao) e um modelo de programação de baixo nível que expõe totalmente a hierarquia de threads e memória torna o CuteDSL uma opção extremamente atraente para maximizar o desempenho de GEMMs em hardware NVIDIA atual e futuro.

É importante esclarecer que nem todas as operações se beneficiam igualmente do novo backend. Para tarefas como cálculos elementares, ativações e reduções — que são limitadas por memória — o Triton já gera código de alta qualidade. Seu modelo de programação em nível de bloco é ideal para esses casos, nos quais apenas acessos vetorizados à memória são necessários, e a diferença de desempenho entre kernels gerados automaticamente e kernels manuais é mínima. O CuteDSL também pode expressar operações pointwise e reduções, mas devido à sua natureza de baixo nível, a geração automática de kernels a partir do zero é complexa. Na prática, tanto o Triton quanto o CuteDSL produzem kernels que performam de forma comparável nesses cenários, o que torna essa complexidade adicional desnecessária. Experimentos internos do time confirmam essa teoria: ao testar kernels de softmax implementados tanto em Triton quanto em CuteDSL em diferentes tamanhos de entrada, ambos atingiram velocidades próximas à largura de banda máxima do GB200.

Já as multiplicações de matrizes (GEMMs) contam uma história diferente. Essas operações dominam o perfil computacional de modelos baseados em transformers: em uma passagem forward típica de um LLM, as GEMMs nas projeções de atenção, camadas FFN e cabeça de saída consomem a maioria dos ciclos de GPU. Para atingir uma utilização próxima ao pico nessas operações, é necessário controle preciso sobre recursos específicos de cada geração de GPU — como tamanhos de tile ajustados à pipeline dos Tensor Cores, gerenciamento explícito de estágios de memória compartilhada, escalonamento de warps e, em arquiteturas mais recentes como a B200, clusters de blocos de threads e memória compartilhada distribuída. São justamente esses os aspectos que linguagens de alto nível abstraem para facilitar o uso. Para simplificar a geração de código de baixo nível, o TorchInductor evita criar kernels do zero: em vez disso, parte de templates otimizados manualmente que expõem os parâmetros ajustáveis necessários para adaptar o desempenho a diferentes shapes de problema.

O backend CUTLASS C++ existente já aborda essa necessidade oferecendo controle de baixo nível, mas a sobrecarga de compilação em C++ impõe limitações práticas: cada variante de kernel exige uma invocação completa do nvcc, tornando custoso avaliar muitas candidatas durante o processo de autotuning e inviável benchmarkar decisões de fusão de epílogo no momento do agendamento. O CuteDSL resolve esse problema por meio de um compilador personalizado de Python para MLIR. A DSL é construída sobre as mesmas abstrações do CUTLASS C++ — mesma álgebra de tiles, mesmos primitivos de hierarquia de memória, mesmo modelo de fusão de epílogo — mas compila em velocidades comparáveis aos outros backends do TorchInductor. Essa combinação possibilita aplicar todo o pipeline de autotuning e benchmark de fusão que o TorchInductor utiliza para outros backends também às kernels GEMM com controle de hardware equivalente ao do CUTLASS. As propriedades específicas que viabilizam isso são: exposição total da hierarquia de threads e memória; melhorias significativas no tempo de compilação; e templates de GEMM otimizados pela NVIDIA.

A arquitetura moderna das GPUs tornou-se extremamente complexa com o avanço dos casos de uso em deep learning e IA. Como consequência, há inúmeras decisões a tomar ao projetar um kernel GEMM, como tamanhos de tile, especialização de warps, formatos de instrução e se usar transferências assíncronas de memória (TMA no Hopper e Blackwell). O Torch.compile está estrategicamente posicionado para resolver esse problema em tempo de execução porque, como compilador JIT, consegue identificar os shapes de problema de um modelo e selecionar a configuração de melhor desempenho usando essas informações. Essa técnica de ajuste automático de kernels a uma carga de trabalho específica é chamada de autotuning. O fluxo do sistema de autotuning de GEMMs do TorchInductor opera em várias etapas bem definidas que garantem otimização eficiente.

Quando o compilador encontra uma multiplicação de matriz durante o processo de lowering, ele primeiro consulta cada backend habilitado (Triton, CUTLASS, cuBLAS) para determinar se o backend suporta o shape, layout e tipo de dado do problema. Backends incapazes de lidar com a configuração são filtrados nesse estágio. Para cada backend elegível, o TorchInductor gera um conjunto de kernels candidatos a partir da biblioteca de templates do backend. Esses candidatos variam em tamanho de tile, configuração de warps e outros parâmetros específicos de cada backend. Todos os candidatos são então submetidos a benchmark na GPU alvo, e o kernel mais rápido é selecionado. O kernel vencedor e seu artefato compilado são armazenados no cache do TorchInductor, permitindo que compilações subsequentes com a mesma configuração de problema pulam completamente a fase de benchmark. Esse sistema de caching opera tanto no nível individual do kernel (código compilado) quanto no nível de seleção (qual candidato venceu para determinado shape e conjunto de backends).

Sobre esse pipeline base, o TorchInductor suporta fusão de epílogo para kernels GEMM. Durante o agendamento, o compilador avalia se é vantajoso fundir operações pointwise downstream no epílogo do GEMM. No caso do Triton, isso é implementado via o buffer MultiTemplate: os N melhores candidatos GEMM do lowering são levados adiante, e possíveis fusões são submetidas a benchmark durante o agendamento para determinar se uma variante fundida supera o GEMM não fundido seguido de um kernel pointwise separado. A seleção final do kernel é adiada até depois que as etapas de fusão sejam concluídas. Esse fluxo completo é especialmente relevante para modelos de linguagem, onde operações como adição de bias, ativações como ReLU ou normalização são frequentemente fundidas em GEMMs para melhorar a eficiência.

O backend CUTLASS C++ suporta fusão de epílogo através da Árvore de Visitantes de Epílogo (EVT), mas o custo de compilação do nvcc por variante limita o número de configurações que podem ser avaliadas na prática. Essa restrição de tempo de compilação é uma das principais motivações para introduzir o CuteDSL como alternativa. Vale notar que, atualmente, a fusão de epílogo não é suportada no backend CuteDSL, mas esse recurso está planejado para desenvolvimento futuro. O novo backend se integra ao pipeline de autotuning descrito acima de forma transparente. Quando o Inductor encontra uma multiplicação de matriz durante o lowering, o backend CuteDSL procede em três etapas: (1) consulta a biblioteca cutlass_api por todas as configurações de kernel compatíveis com o problema; (2) classifica essas configurações usando nvMatmulHeuristics para selecionar os principais candidatos; e (3) compila e faz benchmark desses candidatos no hardware alvo, em paralelo com ATen e Triton.

O CuteDSL se diferencia das abordagens do Triton e CUTLASS C++ de duas maneiras fundamentais. Primeiro, a seleção de kernels via cutlass_api: enquanto o backend Triton gera candidatos de kernel a partir de templates mantidos internamente no TorchInductor, o CuteDSL consulta a biblioteca cutlass_api — uma biblioteca Python mantida pela NVIDIA que contém todo o espaço de configurações de kernels GEMM do CuTeDSL, incluindo shapes de tiles, tamanhos de clusters e parâmetros de escalonamento. O Inductor descreve o problema (shape, tipo de dado, layout, modo de scaling e capacidade computacional da GPU) e a API retorna todos os kernels compatíveis. Quando a NVIDIA adiciona novas configurações de kernel ou suporte a hardware, elas ficam disponíveis na cutlass_api sem necessidade de alterações no Inductor. A API também é extensível: o TorchInductor pode registrar suas próprias classes de kernel na mesma biblioteca. Essa flexibilidade permitiu que a equipe adicionasse suporte a GEMMs FP4 (NVFP4, MXF4) antes mesmo de estarem disponíveis oficialmente na versão upstream — os kernels vendidos passam pelo mesmo filtro, classificação e pipeline de perfilamento dos kernels da NVIDIA.

Segundo, a redução do espaço de busca guiada por heurísticas: consultar a cutlass_api para um problema específico pode retornar centenas de configurações de kernel compatíveis. Avaliar todas elas seria proibitivamente caro. Para resolver isso, o backend CuteDSL integra o nvMatmulHeuristics, um modelo analítico de desempenho da NVIDIA que pontua cada configuração com base na estimativa de throughput de hardware — considerando eficiência de tile, largura de banda de memória e ocupação. Isso reduz centenas de candidatos a apenas algumas dezenas (5 por padrão, configurável via nvgemm_max_profiling_configs). Apenas essas configurações melhor ranqueadas são compiladas e submetidas a benchmark no hardware alvo. Nem o Triton nem o CUTLASS utilizam um modelo analítico desse tipo; eles dependem de benchmarking sobre um espaço de busca menor e definido por templates.

Uma vez que o autotuning seleciona um kernel vencedor, o artefato compilado é armazenado em cache na memória — chamadas subsequentes invocam a função compilada diretamente, sem overhead de compilação repetida. É importante destacar que o backend CuteDSL é puramente aditivo. Se um problema não for compatível com NVGEMM — devido a tipo de dado, layout ou hardware não suportado — nenhum candidato NVGEMM é gerado, e o processo de autotuning prossegue normalmente com ATen e Triton. Se candidatos NVGEMM forem gerados mas perderem no benchmark, o backend mais rápido é selecionado automaticamente. Habilitar o NVGEMM não pode causar regressão de desempenho em hipótese alguma.

Todos os benchmarks foram executados em uma única GPU NVIDIA B200 com 850W e clock dinâmico (sem paralelismo de tensor), utilizando PyTorch nightly e CUDA 13.1. Os resultados de nível de kernel medem a latência isolada de GEMMs via autotuning do Inductor, enquanto os resultados de ponta a ponta avaliam a latência de decodificação do vLLM V1 em modelos como Llama 3.1 8B, Qwen3 32B e Llama 3.3 70B, com prompt de entrada de 32 tokens e geração de 128 tokens, execução serial e cache limpo entre execuções. A equipe avaliou o Inductor NVGEMM contra os backends existentes do Inductor em três regimes de tipos de dado relevantes para LLMs, medindo o throughput das GEMMs em TFLOPS. Os resultados mostram melhorias significativas em shapes específicos.

No regime BF16 (brain floating-point de 16 bits), o NVGEMM apresentou melhorias notáveis em shapes típicos de decodificação (onde M varia de 8 a 64), alcançando speedups de até 1,73x, e em shapes “tall-skinny” como (4096, 256, 4096) com 1,54x de ganho. Shapes grandes típicos de preenchimento (prefill) mantiveram desempenho similar aos backends existentes. Para MXFP8 (matrix floating-point de 8 bits), o NVGEMM mostrou melhorias em shapes médios (até 1,78x) e desempenho equivalente em shapes grandes, enquanto shapes retangulares largas favoreceram o ATen. No regime NVFP4 (NVIDIA floating-point de 4 bits), o NVGEMM obteve melhor throughput em shapes de decodificação (M ≤ 256), com speedups de até 1,6x sobre o melhor backend existente. Em shapes maiores (M ≥ 512), o ATen demonstrou ser bem otimizado e os backends convergiram para desempenho semelhante.

A equipe também mediu a latência de inferência em três modelos com tamanhos de lote variando de 2 a 128 usando o vLLM V1. Como o vLLM utiliza shapes dinâmicos para a dimensão de lote, o TorchInductor não conhece o tamanho real do lote no momento da compilação. Por isso, é necessário usar a configuração autotune_batch_hint para especificar o tamanho de lote alvo, permitindo que o Inductor avalie candidatos de kernel exatamente no shape que será usado em runtime — isso é crucial porque as configurações ótimas de kernel são altamente dependentes do shape. Nos resultados com BF16, a adição do NVGEMM reduziu a latência em 90% das configurações testadas (19 de 21 pontos de dados). A maior melhora foi de 6,5% no Llama 3.3 70B com tamanho de lote 16. O Llama 3.1 8B apresentou ganhos consistentes de 2 a 4% em todos os tamanhos de lote, enquanto o Qwen3 32B mostrou melhorias mais modestas de 0,5 a 2,4%.

No regime NVFP4, o NVGEMM obteve uma taxa de vitória de 89% (16 de 18 pontos de dados). O Llama 3.1 8B melhorou em até 4,2%, o Qwen3 32B em até 3,5%, e o Llama 3.3 70B em até 3,3%. Os ganhos foram mais consistentes em tamanhos de lote entre 16 e 64. Esses resultados demonstram que, embora os ganhos em nível de kernel sejam modestos em alguns casos, a otimização de kernels específicos para determinados shapes e tipos de dado pode se traduzir em melhorias significativas em pipelines completos de inferência, especialmente em modelos grandes onde a latência de decodificação é crítica.

Para habilitar o backend CuteDSL na prática, é necessário instalar a biblioteca cutlass_api a partir de um branch específico do repositório CUTLASS. A expectativa é que essa biblioteca seja incorporada à branch principal do CUTLASS em uma versão futura, eliminando a necessidade dessa instalação separada. Além disso, é preciso ter PyTorch 2.11 ou superior (com suporte NVGEMM para mm, bmm, scaled_mm e grouped_mm), sendo que versões nightly do PyTorch são recomendadas para suporte a kernels FP4 e otimizações de desempenho. Atualmente, a cutlass_api requer a versão 4.3.5 ou anterior do CuTeDSL. Após a instalação, basta adicionar NVGEMM à lista de backends de autotuning do Inductor para ativar a nova funcionalidade. O time forneceu exemplos mínimos de código que demonstram como configurar e usar o backend em diferentes cenários.

A configuração pode ser feita tanto por linha de comando quanto via variáveis de ambiente. Para controlar quantas configurações de kernel são perfiladas por GEMM, basta definir a variável NVGEMM_MAX_PROFILING_CONFIGS. O time também disponibilizou documentação detalhada sobre como integrar o CuteDSL em pipelines de inferência existentes, incluindo considerações sobre compatibilidade com diferentes versões de hardware e software.

O roteiro de desenvolvimento planejado para o backend CuteDSL inclui várias frentes importantes. Primeiro, a implementação de benchmark e suporte a fusão de epílogo. Com os tempos de compilação do CuteDSL deixando de ser um gargalo, o TorchInductor poderá avaliar sistematicamente decisões de fusão de epílogo para kernels GEMM. Isso é crucial porque substituir o cublas por um GEMM individual nem sempre é vantajoso, mas a fusão de epílogo oferece uma oportunidade consistente de superar o cublas, que não consegue realizar fusões. Esse trabalho envolve adiar a seleção final do kernel até depois que as etapas de fusão sejam concluídas, avaliando variantes fundidas e não fundidas entre os backends e selecionando a configuração globalmente ótima. A biblioteca cutlass_api já fornece kernels com capacidade de fusão de epílogo (EFC) que suportam carregamentos/armazenamentos de tensores auxiliares, operações elementares (adição, multiplicação, subtração, divisão) e ativações (ReLU, sigmoid, tanh). O trabalho restante está no lado do Inductor: mapear as decisões de fusão do Inductor para a interface de kernels EFC e integrá-las ao pipeline de agendamento. Operações adicionais de epílogo, como reduções e transmissões de linhas/colunas, estão planejadas para futuras versões da cutlass_api.

Outra frente importante é o suporte a pré-compilação assíncrona e caching persistente. Atualmente, os candidatos a kernel são compilados sequencialmente inline via cute.compile(). A equipe está adicionando pré-compilação paralela entre subprocessos e caching persistente em disco de artefatos compilados, permitindo que execuções de autotuning “quentes” possam pular completamente a fase de compilação. Isso é especialmente valioso em ambientes de produção onde o overhead de autotuning precisa ser minimizado ao máximo.

O time também planeja introduzir caches de configuração exportáveis em formatos portáteis e legíveis por humanos (como JSON ou protobuf), com APIs para importação/exportação que permitam manipulação do cache. Isso possibilitaria a portabilidade de configurações de autotuning entre diferentes execuções e ambientes, facilitando o compartilhamento de otimizações entre equipes e implantações. Outra inovação interessante é uma API de multiplicação de matrizes no estilo FlexAttention, que permitiria aos usuários especificar preferências de backend, configurações de tile e epílogos diretamente no chamador da operação matmul. Isso forneceria controle explícito sobre o comportamento de autotuning e interoperaria com o sistema de cache de configurações exportáveis.

A integração com a biblioteca Quack de Tri Dao também está nos planos. O Quack possui implementações otimizadas de GEMMs para GPUs Blackwell, e a equipe investigará como o desempenho dessas implementações se compara aos templates atuais do CuteDSL, integrando-os caso sejam mais performantes. Além disso, está em estudo o suporte a compilação Ahead-of-Time (AOT) para implantações de inferência, que pré-compilaria kernels CuteDSL no momento da exportação do modelo, eliminando completamente o overhead de autotuning. Isso dependerá de uma API de pré-compilação planejada para a versão 4.4 do CuteDSL e exigirá investigação sobre a acessibilidade em C++ para integração com AOTI.

A longo prazo, espera-se que o CuteDSL atinja plena paridade de desempenho com o backend CUTLASS C++ em novas gerações de hardware. Nesse ponto, o CuteDSL poderia substituir completamente a integração CUTLASS, simplificando drasticamente a base de código do TorchInductor ao consolidar toda a lógica de GEMMs em um único caminho baseado em Python. Essa migração reduziria a complexidade de manutenção e aceleraria o desenvolvimento de otimizações para novas arquiteturas.

Neste artigo, apresentamos a arquitetura do backend CuteDSL do TorchInductor, como habilitá-lo hoje e os resultados de nossos benchmarks. Como mostrado no roteiro de desenvolvimento, este é apenas o início do trabalho, e há muito mais por vir. O time convida a comunidade a testar o novo backend, reportar problemas, fazer perguntas e sugerir novas ideias. Qualquer contribuição pode ser feita abrindo issues no GitHub com a tag apropriada. Para desenvolvedores interessados em se aprofundar, o PyTorch oferece documentação abrangente, tutoriais para iniciantes e avançados, além de recursos para desenvolvimento e suporte técnico. A evolução contínua do TorchInductor e de seus backends de autotuning é fundamental para manter o PyTorch na vanguarda do desenvolvimento de modelos de IA de ponta.