Este artigo demonstra como perfilar mecanismos de atenção no PyTorch usando os rastros do profiler para identificar gargalos de desempenho e oportunidades de otimização. Ele compara implementações ingênuas de atenção contra a Atenção de Produto Ponto Escalado (SDPA) integrada do PyTorch para ilustrar o comportamento do kernel.
- A atenção causal ingênua envolve múltiplos kernels, incluindo matmul, escalonamento, mascaramento, softmax e uma cópia de memória inesperada causada por operações fora do lugar.
- Substituir masked_fill por masked_fill_ in-place remove o kernel desnecessário de cópia de memória do rastro da GPU.
- A função SDPA do PyTorch despacha automaticamente para o backend mais rápido disponível (como Flash Attention ou cuDNN) com base nos tipos de entrada e no hardware.
- O profiling de backends individuais do SDPA revela assinaturas de kernel distintas, permitindo que desenvolvedores verifiquem qual implementação está ativa.
Entender esses rastros do profiler ajuda os usuários a otimizar modelos de transformadores eliminando operações de memória redundantes e aproveitando os backends de atenção eficientes do PyTorch.