Back to Articles Profiling in PyTorch (Part 3): Attention is all you profile
Run the in‑place masked_fill_ version of your attention modules and benchmark with torch.profiler to eliminate the Memcpy kernel and reduce per‑layer latency.
Patch your attention modules to use masked_fill_ and benchmark with torch.profiler to confirm kernel reduction.
Summary
Part 3 of the PyTorch profiling series dives into the attention mechanism, a core component of transformers. The author builds a naive causal attention module in PyTorch, profiling it on an NVIDIA A100‑SXM4‑80GB GPU using scripts 04_a_naive_attention.py, 04_b_inplace_ops_attention.py, 04_c_sdpa_attention.py, and 04_d_kernels_attention.py.
Profiling the naive implementation reveals five GPU kernels: matmul, mul, masked_fill, softmax, and a second matmul, plus an unexpected Memcpy kernel caused by the out‑of‑place masked_fill. Replacing masked_fill with its in‑place counterpart masked_fill_ removes the Memcpy, shrinking the trace and saving both time and memory for each attention forward. The author then demonstrates PyTorch’s built‑in scaled dot‑product attention (F.scaled_dot_product_attention) and shows how the framework dispatches to different backends via torch.nn.attention.SDPBackend. When the math backend is forced, the profiler reports 20 GPU kernels per forward and a 3.7× slowdown compared to the hand‑written version, contradicting the expectation that a single‑liner would be faster. The article also lists the available backends—math, flash, efficient, and cudnn—and explains how to pin a backend with a context manager for targeted profiling. Overall, the post illustrates how careful profiling and in‑place operations can uncover hidden overheads and guide performance‑critical optimizations in transformer models.
Key changes
- In‑place masked_fill_ eliminates a Memcpy kernel per attention forward, reducing GPU kernel count from 5 to 4.
- Naive attention implementation runs 3.7× faster than the math backend of SDPA on A100.
- SDPA math backend launches 20 GPU kernels per forward, compared to 5 in naive implementation.
- PyTorch’s F.scaled_dot_product_attention dispatches to one of four backends: math, flash, efficient, cudnn.
- The author demonstrates how to pin a specific backend using torch.nn.attention.sdpa_kernel context manager.
- Profiling shows that the naive in‑place attention removes the Memcpy kernel entirely, improving memory usage.
- The article provides scripts (04_a_naive_attention.py, 04_b_inplace_ops_attention.py, 04_c_sdpa_attention.py, 04_d_kernels_attention.py) for reproducible experiments.
- The A100‑SXM4‑80GB GPU is used for all profiling runs, highlighting hardware‑specific performance characteristics.