Profiling Attention in PyTorch: kernels, backends and where time goes

This article is part three of the "Profiling in PyTorch" series, which teaches how to read profiler traces and tables and use them to guide optimizations. Previous installments profiled basic tensor math and linear layers; here we focus on attention, a core Transformer primitive that is quadratic in sequence length but has multiple implementation tricks to make it fast.

The example scripts used are: 04_a_naive_attention.py, 04_b_inplace_ops_attention.py, 04_c_sdpa_attention.py and 04_d_kernels_attention.py. Measurements were captured on an NVIDIA A100‑SXM4‑80GB GPU. We inspect CPU and GPU lanes in the profiler traces and compare how different implementations show up.

Naive causal attention

Attention uses Queries (q), Keys (k) and Values (v). The algorithm is a short sequence of primitives: compute scores = matmul(q, k.T), scale scores, apply causal mask (masked_fill), softmax to get weights, and then matmul(attn, v). A straightforward PyTorch implementation looks like this:

class NaiveCausalAttention(nn.Module): def init(self, head_dim): super().init() self.scale = 1.0 / math.sqrt(head_dim) def forward(self, q, k, v, mask): scores = torch.matmul(q, k.transpose(-2, -1)) scores = scores * self.scale scores = scores.masked_fill(mask, float("-inf")) attn = torch.softmax(scores, dim=-1) out = torch.matmul(attn, v) return out

We expected to see two matmuls, a mul, a masked_fill, and a softmax in the trace. The CPU lane shows those ops. Unfolding the GPU lane, however, reveals an extra Memcpy kernel in addition to the expected kernels. The cause: out‑of‑place masked_fill triggers a copy—PyTorch often makes a copy, performs the operation, and returns it.

In‑place masking

Replacing masked_fill with the in‑place masked_fill_ (PyTorch convention: trailing underscore) removes that GPU memcpy kernel entirely:

scores.masked_fill_(mask, float("-inf"))

With this single change, one kernel is shaved off each forward. While one kernel might seem minor, attention runs many times per layer and models have many layers, so the savings accumulate. Note that in‑place ops are unsafe if autograd needs the original forward values; in these tests forward runs under torch.no_grad, so in‑place is safe and also saves memory by avoiding the extra copy.

Scaled Dot Product Attention (SDPA) and backends

PyTorch provides F.scaled_dot_product_attention(q, k, v, is_causal=True), a single API that dispatches to multiple backends based on dtype, shapes, mask, and hardware support. The relevant backends include MATH, FLASH_ATTENTION, EFFICIENT_ATTENTION and CUDNN_ATTENTION. We pin each backend to profile them individually.

Math backend (reference)

The math backend is a careful, reference implementation that decomposes attention into primitive ATen ops, uses FP32 for some paths, and calls a NaN‑safe softmax. Results observed:

• On the profiler table the Naive in‑place forward has an average CUDA time of ~1.955 ms, while the SDPA math backend reports ~7.239 ms for the *_fwd op — about 3.7× slower for this example. Self CUDA time totals were ~7.194 ms vs ~27.279 ms.
• The math backend launches around 20 GPU kernels per forward, whereas the hand‑written naive attention launches about 5.

Why so slow? Several reasons combine:

• The math backend upcasts to FP32 (even when inputs are bf16), so it uses slower sgemm kernels on CUDA cores instead of bf16 Tensor‑Core kernels, increasing data movement and slowing matmuls.
• With is_causal=True the backend materializes the causal mask on every call (aten::ones, aten::tril, aten::where, etc.), so mask construction happens repeatedly.
• It calls a _safe_softmax implementation, which inserts extra kernels to guard against all‑masked rows that would produce NaNs.

In short, math is correct, robust and dtype/NaN safe, but not optimized for raw speed. It is a useful baseline.

Efficient backend (xformers origin)

The efficient backend collapses attention into a single fused kernel: fmha_cutlassF_bf16_aligned_64x64_rf_sm80. Its characteristics:

• A fused multi‑head attention (fmha) kernel built on CUTLASS, running in bf16 with register‑file working sets, minimizing global memory traffic.
• It keeps the computation on Tensor Cores and avoids materializing large [seq, seq] matrices in HBM.

Flash backend (FlashAttention‑2)

The flash backend uses a single fused pytorch_flash kernel (FlashAttention‑2). FlashAttention avoids ever writing the full score matrix to global memory by processing k and v in tiles and maintaining an online softmax, accumulating outputs tile by tile. This enables a single fused bf16 kernel on Tensor Cores.

Profiler readers are often surprised that flash reports low occupancy (≈13%). That is expected: the kernel uses many registers and a lot of shared memory per block, so few blocks can be resident per SM. Low occupancy here reflects heavy on‑chip resource usage by design, which is how the kernel avoids HBM traffic and achieves high throughput.

cuDNN backend

cuDNN also provides a single fused attention kernel, but with important differences:

• cuDNN generates and tunes kernels per problem shape, so it can consume the native [B, H, S, D] layout directly and avoid the transposes other backends perform. The kernel name indicates it is generated and tuned (e.g. cudnn_generated_..._knob_6_128x64x64 ...).
• cuDNN launches via the driver API (cuLaunchKernelEx) rather than the runtime launch; CUPTI may then report 0% achieved occupancy (a measurement gap), but the kernel footprint shows heavy register/shared memory usage similar to flash.
• cuDNN moves some cost onto the CPU: on this tested shape the measured average CPU time per forward was ~214 µs for cuDNN, compared to ~138 µs for flash and ~117 µs for efficient. That CPU time comes from cuDNN selecting and preparing the tuned plan (knob) on every call.

On this specific shape flash was fastest on the GPU (GPU avg times: efficient ~277.9 µs, flash ~146.8 µs, cuDNN ~186.3 µs), but cuDNN can win on other shapes because it retunes per problem at the expense of CPU prep time.

At a glance: what each trace revealed

• Naive attention (hand‑written): 6 kernels; hidden Memcpy from out‑of‑place masked_fill.
• Naive in‑place: 5 kernels; a single line removes the Memcpy kernel.
• SDPA math: 20 kernels; FP32 upcast, mask rebuilt every call, _safe_softmax — correct but ~3.7× slower.
• SDPA efficient: 1 fused fmha_cutlassF kernel, bf16 Tensor Core path.
• SDPA flash: 1 fused pytorch_flash kernel (FlashAttention‑2), fastest on this shape despite low reported occupancy.
• SDPA cuDNN: 1 generated kernel per problem, no transposes, driver API launch, but higher CPU cost for per‑call tuning.

Final takeaway

Before opening a profiler trace, make a hypothesis about what you expect to see, then inspect the trace and treat any mismatch as the most interesting clue. The series’ key insights (hidden Memcpy, addmm epilogue, math backend’s 20 kernels, flash’s apparent low occupancy, cuDNN’s big CPU bar) all came from asking "why" when the trace contradicted our guess.

Profiling is an approachable discipline: look closely and ask "wait, why is that happening?" until the answer becomes clear. You now have the vocabulary and reflex to open traces, form hypotheses, and find actionable mismatches.

Thanks to Noe Flandre for reviewing early drafts.

(Note: an LLM was used to polish the post’s prose; this does not mean the content was generated autonomously.)