Dois números em uma GPU
Você já sabe fazer c = a + b em Python. Já fez mil vezes. Aqui vamos falar sobre isso em tensores do PyTorch. Um tensor é apenas uma matriz de números. Colocá-lo em uma GPU significa que essa matriz reside na memória da GPU em vez da memória comum da CPU. Quando a e b são dois tensores em uma GPU, essa única linha termina rápido o suficiente para que você nunca pense nisso.
Agora reduza. Digamos que a e b sejam dois floats únicos, ambos na GPU. Mesma linha. c = a + b. O que realmente roda no chip?
A resposta é um kernel. Nesse contexto, um kernel é um pequeno programa que a GPU executa em algum pedaço de dados. Não é o kernel do sistema operacional que inicializa seu laptop. Não é o kernel matemático de um livro de álgebra linear. A palavra é reutilizada muitas vezes, e a culpa não é sua. No mundo das GPUs, um kernel significa apenas: uma pequena função que a GPU é instruída a executar, agora mesmo, em paralelo, nos dados que você passou para ela.
Ao final disso, você será capaz de olhar para um trecho de código PyTorch e contar quantos kernels a GPU executará. Isso parece um truque pequeno, e é, mas também é o primeiro passo para se livrar da sensação de "GPU como caixa-preta", aquela em que seu modelo está lento e você não sabe por quê. Toda pergunta que você pode fazer sobre desempenho de GPU acaba voltando para "quais kernels rodaram, e o que eles estavam fazendo". Então é por aqui que começamos.
Seu primeiro kernel
Vamos aumentar um pouco a e b. Desta vez, tensores de comprimento 8. Ainda uma única linha do PyTorch: c = a + b.
Quando você executa isso, sua CPU (a máquina que está rodando o Python) diz para a GPU: ei, vá executar este programa nestes dados. Essa instrução é chamada de lançamento (launch). O que é lançado é um kernel: um programa, pronto para rodar. Os lançamentos em si são baratos, microssegundos cada. O que está ao redor do lançamento (os dados indo para a GPU, os resultados voltando) é onde o custo real está, e é isso que vamos contar ao longo do caminho.
Dentro do kernel, o trabalho real é feito por pequenos trabalhadores chamados threads. Uma GPU tem milhares deles disponíveis. Para nossa soma de comprimento 8, 8 threads pegam o trabalho: thread 0 lida com o elemento 0, thread 1 lida com o elemento 1, e assim por diante até a thread 7. Cada thread executa o mesmo pequeno programa: ler um elemento de a, ler o elemento correspondente de b, somá-los, escrever o resultado em c.
(Na prática, a GPU lança threads em grupos de tamanho fixo chamados warps, sempre 32 threads em placas NVIDIA, e mascara os extras quando seu array não é divisível exatamente. Pode ignorar por enquanto.)
Então temos uma linha do PyTorch, um lançamento, um kernel, 8 threads fazendo 8 somas. Agora vamos contar o que realmente cruzou o chip. Para fazer a soma, cada thread precisa de seu elemento de a e seu elemento de b. Isso são 8 leituras de a e 8 leituras de b. Em seguida, cada thread escreve seu resultado em c. Isso são 8 escritas.
Essas leituras e escritas vão para a memória grande que fica ao lado do chip da GPU. Em placas de datacenter (A100, H100), essa memória é chamada de HBM (high-bandwidth memory). Em placas de consumo (RTX 4090) e nas T4s estilo Colab que as pessoas realmente tentam usar, é chamada de GDDR. De qualquer forma, é uma memória rápida ao lado do chip, e vamos chamá-la apenas de memória da GPU. É rápida, mas não é de graça, e cada viagem até ela tem um custo.
Um kernel = um lançamento = uma passada sobre os dados. O que o kernel faz dentro de seu corpo, as leituras e escritas em suas bordas (as viagens até a memória da GPU para buscar entradas, a viagem de volta para escrever saídas) são a parte que custa. Essa é a forma completa.
Nada disso muda quando os tensores ficam maiores. Mesma linha do PyTorch, mesmo kernel, apenas mais threads. Se a e b têm um milhão de elementos cada, a GPU lança o mesmo kernel com um enxame maior de threads. A matemática escala, os bytes movidos escalam, o modelo mental não. Uma linha, um kernel.
O que acontece entre duas operações
1c = (a + b).relu()
Você conhece Python o suficiente para saber que isso são duas operações, uma soma e depois uma relu, encadeadas. Em um interpretador, são duas chamadas de função. Em uma GPU, no PyTorch eager, são dois lançamentos de kernel: um para a soma, um para a relu. Até aqui, nada surpreendente.
O que é realmente interessante é o que acontece entre os dois kernels.
Quando a soma termina, ela precisa colocar seu resultado em algum lugar. Esse lugar é a memória da GPU. A soma escreve um array intermediário completo (vamos chamá-lo de tmp) na memória. Então, um momento depois, a relu é lançada, e seu primeiro trabalho é ler esse mesmo array tmp de volta da memória. Ela lê o array inteiro, aplica relu a cada elemento e escreve o resultado em c.
Conte o tráfego de memória para esses dois kernels:
- A soma: lê
a, lêb, escrevetmp. Três transferências do tamanho do array. - A relu: lê
tmp, escrevec. Mais duas.
Cinco transferências do tamanho do array no total. Compare com a soma de comprimento 8 sozinha da seção anterior, que foi de três. Adicionar .relu() à cadeia não custou apenas o processamento da relu. Custou uma viagem de ida e volta completa do array através da memória da GPU, porque tmp teve que ser escrito para que o próximo kernel pudesse lê-lo de volta.
Nada foi armazenado em cache. tmp não ficou em um registrador ou em um cache local rápido. Ele foi para a memória da GPU (o tipo lento e distante) e voltou imediatamente. Os dois kernels são estranhos um para o outro. Eles precisam passar dados através do único meio que ambos sabem usar: a memória da GPU.
Por que o PyTorch faz isso? Porque no modo eager, quando você escreve a + b, o PyTorch executa agora. Ele não sabe que você estava prestes a chamar .relu() em seguida. Cada operação é despachada no momento em que sua linha Python é executada. Não há plano, nem antecipação. Cada operação é independente, produz um array real e o entrega ao que vier depois através da memória.
Este é o padrão para se fixar. Cada valor intermediário em seu código PyTorch é fisicamente escrito na memória da GPU e lido de volta pela próxima operação. Cada um. É isso que a "contagem de kernels" realmente mede. Cada kernel extra é outra viagem de ida e volta que seus dados têm que fazer pela memória da GPU.
Fusão: duas operações, um kernel
Imagine um kernel que faz tudo de uma vez: lê seu elemento de a, lê seu elemento de b, soma-os, aplica relu ao resultado (tudo dentro do kernel, em um pequeno espaço de trabalho por thread que nunca sai do chip), e só então escreve o valor final em c. O intermediário (a + b) ainda existe, mas apenas dentro do kernel, no espaço privado de cada thread. Ele nunca é escrito na memória da GPU. tmp, como um array real, nem existe.
Conte as transferências agora. Leituras de a: 1 por elemento. Leituras de b: 1 por elemento. Escritas de c: 1 por elemento. Três transferências do tamanho do array. A mesma matemática da versão com dois kernels, mas com duas viagens de ida e volta a menos.
No comprimento 8, isso é um erro de arredondamento. Ninguém se importa. No comprimento de 1 milhão, ou 100 milhões, essas viagens extras se tornam uma grande parte do tempo de execução, e o relógio de parede reflete isso. Por que o tráfego de memória acaba dominando dessa forma é o assunto do Artigo 2 desta série, então vou deixar o "porquê" de lado aqui. O ponto agora é apenas: mesma matemática, menos viagens, mais rápido na prática.
Esse truque de combinar operações que seriam kernels separados em um único kernel para que o intermediário nunca precise visitar a memória da GPU tem um nome. Chama-se fusão. É isso. Essa é a ideia toda.
Agora a parte complicada. Escrever esse kernel combinado manualmente parece fácil para soma + relu. Duas operações. Uma linha de "computação" no meio. Mas o código PyTorch real tem dezenas de operações encadeadas, cada uma com suas próprias formas, tipos de dados e regras de broadcasting. Escrever um kernel fusionado que lide com tudo isso corretamente é um trabalho de engenharia de verdade. Geralmente você não escreveria esses kernels elementares de rotina manualmente.
Boas notícias: o PyTorch vem com uma ferramenta que faz a reescrita para você, exatamente para esse tipo de caso, automaticamente. Chama-se torch.compile.
Você quase certamente já viu torch.compile usado assim:
1model = torch.compile(model)
Uma linha. Alguém na internet disse que isso deixa as coisas mais rápidas. Aqui está o que realmente faz, em português claro: em vez de executar suas operações uma de cada vez como o modo eager faz, o torch.compile captura as operações de tensor que sua função realiza, procura por oportunidades de combiná-las e gera código otimizado. Chamadas posteriores que correspondam às mesmas premissas podem reutilizar esse código.
A fusão que fizemos no papel acima (soma e relu, compartilhando um kernel, tmp nunca tocando na memória) é exatamente o tipo de coisa que o torch.compile fará em seu código automaticamente, desde que as operações sejam simples o suficiente. Quando as pessoas dizem que torch.compile "deixa o PyTorch mais rápido", isso é uma grande parte do que elas querem dizer.
Para os casos em que o torch.compile não consegue fundir sozinho (operações personalizadas que ele não reconhece, reduções incomuns, layouts de memória estranhos), alguém ainda precisa escrever um kernel manualmente. É para isso que servem ferramentas como Triton e CUDA. Um artigo separado.
Veja com seus próprios olhos
Tudo acima foi contar kernels no papel. Hora de contá-los em uma GPU real. Se você tem qualquer máquina com uma GPU CUDA disponível (uma estação de trabalho, um notebook Colab, uma instância na nuvem), pode executar isso em alguns minutos.
A ferramenta é o torch.profiler. Está embutido no PyTorch. Tudo o que ele faz é registrar o que a GPU realmente fez enquanto seu código executava, e devolver uma tabela que você pode ler.
Passo 1: a versão eager.
Coloque a linha de duas operações em uma função para termos algo a chamar:
1import torch2from torch.profiler import profile, ProfilerActivity34def add_relu(a, b):5 return (a + b).relu()67a = torch.randn(1_000_000, device="cuda")8b = torch.randn(1_000_000, device="cuda")910with profile(activities=[ProfilerActivity.CUDA]) as prof:11 add_relu(a, b)12 torch.cuda.synchronize()1314print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
torch.cuda.synchronize() está aí apenas para garantir que a GPU terminou antes de lermos as medições. O trabalho da GPU é executado de forma assíncrona, e sem a sincronização você às vezes mediria a sobrecarga do lançamento em vez do trabalho real do kernel.
Passo 2: leia a saída.
A saída real do profiler terá mais linhas do que você pode esperar. Várias linhas de alocação de memória e contabilidade do PyTorch estarão misturadas. As linhas que nos interessam são os kernels CUDA, as funções reais que a GPU executou. Procure linhas com kernel no nome. As duas serão mais ou menos assim:
1vectorized_elementwise_kernel<...BinaryFunctor_add...> 1 12us2vectorized_elementwise_kernel<...threshold_kernel...> 1 9us
Os nomes exatos dos templates mudam entre versões do PyTorch (relu aparece frequentemente como threshold porque essa é a operação subjacente, e add às vezes como CUDAFunctor_add). Não tente analisar tudo. Apenas conte as linhas. Duas linhas. Dois kernels. Um para a soma, um para a relu. Exatamente o que dissemos que aconteceria algumas seções atrás.
Passo 3: a versão compilada.
Mude uma linha. Envolva a função em torch.compile:
1compiled = torch.compile(add_relu)
Agora, aqui vai uma pegadinha que vale a pena saber antes de executar. A primeira chamada a uma função com torch.compile é lenta, às vezes estranhamente lenta, porque é quando o torch.compile faz seu trabalho: analisar seu código, descobrir o que fundir, gerar o kernel fusionado. Se você perfilhar a primeira chamada, estará medindo a etapa de compilação, não o kernel. Então o padrão é: chame uma vez para aquecer, descarte o resultado, então perfilhe.
1compiled(a, b) # aquecimento, descarte2torch.cuda.synchronize()34with profile(activities=[ProfilerActivity.CUDA]) as prof:5 compiled(a, b)6 torch.cuda.synchronize()78print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
Passo 4: leia a saída novamente.
Desta vez:
1triton_poi_fused_add_relu_0 1 14us
Uma linha. Um kernel. O nome até diz o que fez: soma e relu fusionadas. Mesma matemática de antes, um lançamento em vez de dois.
Você acabou de fazer o que este artigo vem falando, em uma frase: pediu ao PyTorch para combinar duas operações em um kernel, observou o profiler e confirmou que a contagem foi de dois para um. Fusão, na prática, na sua máquina.
Se quiser ver de forma mais dramática, tente com alguns tamanhos de tensor diferentes. No comprimento 100, ambas as versões rodam tão rápido que a diferença se perde no ruído. No comprimento de 10 milhões ou 100 milhões, a versão compilada começa a ficar claramente à frente, porque a viagem de ida e volta que cortamos é uma parte real do trabalho nessa escala.
Contar kernels não é mais um conselho abstrato. Você tem uma maneira de verificar.
Vamos finalizar
Aqui está tudo em um resumo.
Seu código PyTorch, quando roda em uma GPU, se transforma em uma sequência de kernels. Cada kernel é um lançamento, uma passada sobre seus dados, uma viagem de ida e volta pela memória da GPU para buscar entradas e escrever saídas. Operações simples se tornam um kernel. Cadeias de operações se tornam um kernel por operação por padrão, com intermediários fazendo viagens de ida e volta pela memória entre eles. O torch.compile pode fundir cadeias simples para você, de modo que esses intermediários nunca toquem na memória. Menos kernels geralmente significam menos tráfego de memória. E menos tráfego de memória geralmente significa mais rapidez.
