Por que Perfilhar Atenção é Essencial
A atenção é o coração dos Transformers, mas sua complexidade quadrática a torna um alvo principal para otimização. Muitas técnicas existem, mas o verdadeiro desafio é entender como cada implementação se comporta por baixo dos panos. Neste tutorial, vamos perfilhar quatro variantes de atenção no PyTorch: uma versão naive escrita à mão, uma versão otimizada com in-place, e a função nativa scaled_dot_product_attention (SDPA) com diferentes backends. Você vai aprender a ler traces de profiler, identificar ineficiências e tomar decisões informadas para seus modelos.
Usaremos uma GPU NVIDIA A100 e a ferramenta trace-util do Hugging Face. Todo o código está disponível no repositório companheiro.
Passo a Passo: Atenção Naive e Otimização In-Place
Vamos começar com um módulo de atenção causal naive construído a partir de primitivas:
import torch
import torch.nn as nn
import math
class NaiveCausalAttention(nn.Module):
def __init__(self, head_dim):
super().__init__()
self.scale = 1.0 / math.sqrt(head_dim)
def forward(self, q, k, v, mask):
scores = torch.matmul(q, k.transpose(-2, -1))
scores = scores * self.scale
scores = scores.masked_fill(mask, float("-inf"))
attn = torch.softmax(scores, dim=-1)
out = torch.matmul(attn, v)
return out
Perfilhar isso revela um kernel Memcpy oculto causado pelo masked_fill out-of-place. Mudando para a versão in-place (masked_fill_) elimina essa cópia:
def forward(self, q, k, v, mask):
scores = torch.matmul(q, k.transpose(-2, -1))
scores = torch.mul(scores, self.scale)
scores.masked_fill_(mask, float("-inf")) # in-place
attn = torch.softmax(scores, dim=-1)
out = torch.matmul(attn, v)
return out
Dica de ouro: Operações in-place economizam memória e tempo, mas só são seguras sob torch.no_grad() (sem backward pass).
Comparando Backends SDPA: Math, Efficient, Flash, cuDNN
A função F.scaled_dot_product_attention do PyTorch despacha para múltiplos backends. Veja como eles se comparam:
| Backend | Kernels/forward | Característica principal | Nota de desempenho |
|---|---|---|---|
| Math | 20 | Implementação de referência: FP32 nos CUDA cores, reconstrói máscara a cada chamada, usa _safe_softmax | ~3.7x mais lento que o naive, mas numericamente seguro |
| Efficient | 1 | Kernel fmha_cutlassF fusionado (xformers), mantém bf16 nos Tensor Cores | Bom equilíbrio, sem upcast para FP32 |
| Flash | 1 | Kernel pytorch_flash fusionado (FlashAttention-2), baixa ocupação (~13%) mas o mais rápido | Melhor para sequências longas, evita tráfego HBM |
| cuDNN | 1 | Kernel gerado por problema (cudnn_generated_fort...), sem transposes, usa cuLaunchKernelEx | Custo movido para CPU (seleção de plano ~214µs) |
Dica de perfilamento: Sempre adivinhe o que o trace deve mostrar antes de abri-lo. As discrepâncias são onde os insights moram. Por exemplo, a baixa ocupação do Flash parece "errada" mas é intencional—troca ocupação por reuso de dados on-chip.
Conclusão & Próximos Passos
Perfilamento é um superpoder: formule uma hipótese, abra o trace e persiga a discrepância. Economizamos um kernel com um underscore, descobrimos o overhead de 20 kernels do backend math, e vimos por que a baixa ocupação do Flash é na verdade uma feature.
Sua vez: Perfilhe sua própria camada de atenção. Tente fixar backends com torch.nn.attention.sdpa_kernel e compare traces. Para um mergulho mais profundo em colaboração open-source em escala, confira como a Meta descontinuou seu fork interno do FFmpeg. E se você curte arquiteturas serverless, veja Serverless Genomics at Scale.
Bom perfilamento! 🤗