PyTorch profiling shows in-place masking removes memory copy and SDPA dispatches to optimized backends

This article demonstrates how to profile attention mechanisms in PyTorch using the profiler traces to identify performance bottlenecks and optimization opportunities. It compares naive attention implementations against PyTorch's built-in Scaled Dot Product Attention (SDPA) to illustrate kernel behavior.

• Naive causal attention involves multiple kernels including matmul, scaling, masking, softmax, and an unexpected memory copy caused by out-of-place operations.
• Replacing masked_fill with the in-place masked_fill_ removes the unnecessary memory copy kernel from the GPU trace.
• PyTorch's SDPA function automatically dispatches to the fastest available backend (such as Flash Attention or cuDNN) based on input types and hardware.
• Profiling individual SDPA backends reveals distinct kernel signatures, allowing developers to verify which implementation is active.

Understanding these profiler traces helps users optimize transformer models by eliminating redundant memory operations and leveraging PyTorch's efficient attention backends.