torch.compile e CUDA graphs

01 · Conceito

Conceito

Você perfila o decode de fluxo único no Qwen3.8-27B e o trace é quase todo buraco. A GPU faz uma rajada de trabalho, fica ociosa, faz outra rajada, fica ociosa de novo. A utilização fica baixa, o consumo de energia nunca chega perto do limite da placa, e trocar por um acelerador mais rápido quase não muda nada. Nada está quebrado. Você bateu num gargalo que nem sequer está na GPU.

A lição 7.3 estabeleceu a primeira metade da história: o decode tem intensidade aritmética péssima. Para emitir um token a máquina faz passar uma quantidade enorme de pesos pelas unidades aritméticas e realiza poucas operações sobre cada byte, então quem dita o ritmo é a banda, não a computação. Esta lição acrescenta a segunda metade. Toda operação em modo eager carrega custo do lado do host antes que qualquer aritmética aconteça: um frame Python, um dispatch pelo registro de operadores, checagem de argumentos e, por fim, um lançamento que entrega um kernel ao driver.

Conte-os. Uma única camada de decoder neste modelo realiza, grosso modo, duas normalizações, as projeções e as tripas do seu token mixer, três multiplicações matriciais da FFN mais uma multiplicação de gating, e um punhado de operações elementwise e reshapes — chame de algumas dezenas de kernels, e mais nas camadas DeltaNet, cuja convolução curta e atualização de estado recorrente são feitas de vários passos pequenos cada. Tome trinta como média ilustrativa ao longo de 64 camadas e você está disparando da ordem de dois mil lançamentos para produzir um token. A alguns microssegundos de tempo de host por lançamento, isso são vários milissegundos por token gastos puramente em dispatch. Transforme em teto: oito milissegundos de trabalho de host por token limitam você perto de 125 tokens por segundo, independentemente do que a GPU consiga fazer.

Agora compare os dois tetos com honestidade, porque a versão ingênua deste argumento está errada. Em bf16, fazer passar todo o conjunto de pesos de 27 B de parâmetros a cada token coloca o teto de banda nas dezenas de tokens por segundo em um acelerador de topo — abaixo do teto de lançamento. Ou seja, em bf16 com batch 1, o overhead de lançamento ainda não é a restrição que amarra; é um imposto de algumas dezenas de por cento montado sobre um muro de banda. Aqui está a virada: quantize para 4 bits (lição 7.9) e os bytes que você precisa mover por token caem por um fator de aproximadamente quatro, então o teto de banda sobe por um fator de aproximadamente quatro. O teto de lançamento não se move nem um pouco. O gargalo troca de mãos. É por isso que “compile isso” é conselho opcional para um modelo 27B em bf16 e conselho obrigatório para um quantizado, e por que modelos pequenos rodando em placas rápidas são quase sempre limitados por lançamento.

O primeiro remédio é a compilação. O torch.compile traça o bytecode Python do forward pass em um graph, protege esse graph com guards sobre os tensors que viu (dtypes, shapes, device e mais) e o rebaixa para kernels gerados. Duas coisas ficam mais baratas. Cadeias de trabalho elementwise — o escalonamento da normalização, a multiplicação do gate SiLU, as somas residuais — fundem em kernels únicos, o que corta tanto a contagem de lançamentos quanto as idas e vindas à memória que tornavam essas operações caras em primeiro lugar. E toda a região traçada é disparada como código compilado em vez de Python reinterpretado.

Os guards também são onde a compilação dá errado. Quando um guard falha, o graph é recompilado, e o decode muda o comprimento da sequência a cada passo. Um modelo que recompila uma vez por token é dramaticamente mais lento que o eager e vai parecer, num benchmark ingênuo, que a compilação é inútil. A abordagem corrigida é tornar dinâmica a dimensão que varia, para que um graph cubra muitos comprimentos, manter os shapes que precisam ficar estáticos num pequeno conjunto de buckets, fazer warm-up até o contador de recompilações parar de subir e só então medir. Medir a primeira chamada depois de compilar — a virada errada que quase todo mundo dá uma vez — cronometra o compilador, não o modelo.

O segundo remédio ataca o que sobra. Mesmo um graph bem fundido ainda dispara uma sequência de lançamentos, e essa sequência é idêntica a cada passo de decode. O CUDA graph capture a grava uma vez e a reproduz inteira como uma única submissão, de modo que o host emite um comando em vez de centenas. Os requisitos são estritos e explicam boa parte do design dos serving engines. Os shapes precisam ser estáticos, os endereços de memória precisam ser estáticos, e não pode haver sincronização com o host nem controle de fluxo dependente da CPU dentro da região capturada, porque uma reprodução não roda o seu Python de novo — ela reexecuta lançamentos gravados contra ponteiros gravados.

Guarde a contabilidade, não as flags. O eager gasta seu tempo dizendo à GPU o que fazer; a compilação reduz o quanto é preciso dizer, e o graph capture reduz a frequência com que você precisa dizer. Nenhuma das duas compra uma multiplicação matricial mais rápida. Elas compram de volta os buracos no trace — e os buracos são a razão inteira pela qual um modelo 27B consegue deixar uma placa caríssima quase ociosa.