Profiling in PyTorch: optimize Attention and GPU backends

Attention is a fascinating black box—until you profile it. In this technical installment we walk through real traces, GPU kernels, and layer-level decisions that change time and memory. What happens when you go from a hand-written NaiveCausalAttention to F.scaled_dot_product_attention and to each available backend? Let’s see it step by step.

Why profile attention?

Because attention looks simple in pseudocode, but in practice it moves a lot of memory. Matmuls, scaling, masking and softmax are familiar primitives, yet how they're executed on the GPU determines whether your model runs fast or spends time waiting. Is your model suffering a bottleneck? The profiler answers with kernel names and real timings.

Naive attention and the memcpy surprise

This is the clear, direct implementation:

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

When you profile it you see exactly what you expect: matmul, mul, masked_fill, softmax, matmul. But the GPU trace reveals an extra kernel: a Memcpy. Where does it come from? From out-of-place operations. masked_fill often creates an intermediate copy, and that copy shows up as a memcpy in the profile.

What's the fix? Use the in-place version when it's safe:

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

With that change the Memcpy disappears from the trace and you gain time and memory. Heads-up: in-place is safe only when you run forward without gradients or you know you won’t corrupt values needed for the backward pass.

The SDPA one-liner that hides several paths

PyTorch offers a single call that replaces all of the above:

F.scaled_dot_product_attention(q, k, v, is_causal=True)

Nice, right? But the profiler shows surprises depending on which backend that function uses. Look at the math backend: instead of simplifying, it launches about 20 kernels per step and ends up roughly 3.7x slower than our naive in-place version. Why? Because the math backend is the reference implementation: it upcasts to FP32 for numerical safety, reconstructs the mask on every call, and uses _safe_softmax to avoid NaNs in fully masked rows. It's correct and robust, but slow.

Lesson: a clean API doesn't guarantee the fastest path. Sometimes convenience moves extra work to less visible places.

Fused backends: efficient, flash and cuDNN

All of them aim to avoid writing the full scores matrix to global memory and to keep work on-chip in bf16 leveraging Tensor Cores. But each backend does it in its own way.

efficient (xformers)

It launches a single kernel with a name like fmha_cutlassF_bf16_aligned_64x64_rf_sm80. It's a fused multi-head attention kernel built on CUTLASS, runs in bf16 without upcasting, and keeps tiles in registers and shared memory. Result: one kernel per forward and high efficiency.

flash (FlashAttention-2)

Vendored as pytorch_flash, it implements the trick of walking K and V in tiles, maintaining an online softmax and never materializing the full [seq, seq] matrix. The profiler can show low occupancy (for example 13%) and that may scare you at first, but it's intentional. Flash uses many registers and shared memory per block to keep data on-chip. Low occupancy does not mean inefficient; it means each block is heavy and does a lot of work with few HBM accesses.

cuDNN

cuDNN generates a kernel per problem; its long name reflects a generated and retunable version. Two vital points:

• It doesn't need extra transpose metadata because its generator consumes the native layout.
• It launches with cuLaunchKernelEx, so the profiler can report 0% occupancy as a measurement limitation, not because the GPU is doing poorly.

The detail: cuDNN often wins on raw GPU execution, but shifts time to the CPU to select and prepare the strategy on every call. In traces that appears as a large CPU bar, unlike flash or efficient.

Quick numbers from the original article

And between the naive in-place and SDPA math there was a large difference in CUDA avg time: something like 1.955 ms versus 7.239 ms for math in that reference form. The moral: an "all-in-ATen" backend tries to be correct and safe, not necessarily the fastest.

Practical checklist for profiling attention

• Before looking at traces, guess what you should see. Which kernels do you expect? That expectation helps you spot anomalies.
• If you see a suspicious Memcpy, inspect out-of-place ops like masked_fill. Try the in-place version carefully.
• Try F.scaled_dot_product_attention but switch backends to compare math, efficient, flash and cudnn using torch.nn.attention.sdpa_kernel to pin them and profile each one.
• Watch dtype: if a backend upcasts to FP32, memory cost can double; bf16 + Tensor Cores is where gains live on A100.
• Don’t trust a clean-looking trace alone. Hidden work can move to the CPU or inside a library and show up as an opaque bar.

Final reflection

Is profiling a privilege of GPU experts? No. It's simply the habit of looking, guessing and asking about differences until everything makes sense. In this walkthrough we saw everything from a memcpy hidden by an out-of-place op to fused kernels that leverage Tensor Cores and strategies generated by cuDNN. Your task? Open a trace, make your hypothesis, and go find the explanation. That practice will give you latency and memory wins that truly matter in large models.

Original source

https://huggingface.co/blog/torch-attention-profile