Profiling in PyTorch (Part 3): Attention is all you profile

Profiling the attention mechanism in PyTorch reveals how its fundamental operations, like matrix multiplication and softmax, execute on hardware. A naive implementation shows up in the profiler as a sequence of distinct steps, with each one launching separate kernels on the GPU. This creates noticeable gaps and overhead between operations as data is written to and read from memory multiple times. By contrast, more advanced implementations like Scaled Dot Product Attention (SDPA) fuse these steps into a single, highly optimized kernel. Analyzing these different profiler traces provides a clear visual guide on how to optimize code by reducing kernel launch overhead and improving memory access patterns.