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

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

The blog post profiles attention mechanisms in PyTorch using the NVIDIA A100-SXM4-80GB GPU and demonstrates how each variant appears in profiler traces.

Naive attention

A naive attention implementation built from primitive operations (matmul, mul, masked_fill, softmax, matmul) launches six GPU kernels per forward pass, including an unexpected memory copy caused by the out‑of‑place masked_fill.

Naive attention with inplace causal masking

Replacing masked_fill with the in‑place masked_fill_ removes the memory copy kernel, reducing the GPU kernel count from six to five per forward pass.

Scaled Dot Product Attention

PyTorch’s F.scaled_dot_product_attention provides a single‑line interface that dispatches to multiple backends; each backend shows distinct profiling characteristics.

Math backend

When pinned to the math backend, SDPA launches about twenty GPU kernels per forward pass, runs in FP32 on CUDA cores, rebuilds the causal mask on every call, and uses _safe_softmax to avoid NaNs; it is roughly 3.7× slower than the naive in‑place version.

Efficient backend

The efficient backend (xformers) launches a single fused fmha_cutlassF_bf16_aligned_64x64_rf_sm80 kernel per forward pass, keeping the computation in bfloat16 on Tensor cores and avoiding intermediate matrix writes to HBM.

Flash backend

The flash backend launches a single fused pytorch_flash kernel (FlashAttention‑2) per forward pass; despite showing low estimated occupancy (around 13%) in the profiler, it is the fastest backend because it uses registers and shared memory to keep attention tiles on‑chip.

cuDNN backend

The cuDNN backend launches a single generated kernel per forward pass (e.g., cudnn_generated_fort_native_sdpa_sm80_flash_fprop_wmma_f16_knob_6_128x64x64_4x1x1_cga1x1x1_kernel0_0); it avoids transpose operations on the CPU but incurs higher CPU time (~214 µs) due to runtime plan selection, and its GPU time lands between the efficient and flash backends.

Everything we covered, at a glance

Variant	What we changed	Kernels / forward	What the trace revealed
Naive attention	Attention built by hand from primitives (matmul, mul, mask, softmax, matmul)	6	A hidden Memcpy from the out‑of‑place masked_fill.
Naive in-place	masked_fill → masked_fill_	5	One line drops the Memcpy kernel entirely.
SDPA math	F.scaled_dot_product_attention pinned to the math backend	20	The reference: FP32 on CUDA cores, mask rebuilt every call, _safe_softmax. Correct but ~3.7× slower.
SDPA efficient	Efficient (xformers) backend	1	One fused fmha_cutlassF kernel, stays in bf16 on Tensor cores.
SDPA flash	Flash backend	1	One fused pytorch_flash kernel (FlashAttention‑2). Fastest, despite "wrong‑looking" 13% occupancy.
SDPA cuDNN	cuDNN backend	1	A per‑problem generated kernel: no transposes, cuLaunchKernelEx, but the cost moved to a fat CPU bar.

Concluding the series

The main takeaway is to form a guess about what a profiler trace should show, then examine the trace and treat any mismatch as the most interesting insight; this habit revealed the hidden Memcpy, the math backend’s twenty kernels, flash’s low occupancy, and cuDNN’s CPU‑side cost.

Now you can apply this guessing‑and‑checking approach to profile your own models.