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

Source: https://huggingface.co/blog/torch-attention-profile Back to Articles

• Naive attention
• Naive attention with inplace causal masking
• Scaled Dot Product Attention- Math backend - Efficient backend - Flash backend - cuDNN backend
• Everything we covered, at a glance
• Concluding the series

The series “Profiling in PyTorch” is meant to make you comfortable reading profiler traces and tables. InPart 1we profiled basic math operations like addition and multiplication. We saw how the profiler table uncovers hotspots, and how the profiler trace shows the order in which an algorithm runs over time.

InPart 2we wrapped that addition and multiplication into a torch linear layer. We then stacked several linear layers on top of each other (a multilayer perceptron) and profiled that. Along the way we also profiled fused and hand-tuned kernels.

From the perspective of the Transformer architecture, the next logical step for us to profile is yet another fundamental algorithm, attention. While being infamous for its quadratic-time complexity, many clever tricks exist to mitigate that issue and make it fast. Our goal here is not to cover every trick in detail. Instead, we want to see how each one looks different under the profiler.

The scripts for this blog post live here:04\_a\_naive\_attention\.py,04\_b\_inplace\_ops\_attention\.py,04\_c\_sdpa\_attention\.py, and04\_d\_kernels\_attention\.py. Like before, it helps to open them in a separate tab and walk through the code as you read. We use anNVIDIA A100\-SXM4\-80GBGPU to run the scripts. It is really easy to set up a GPU on the Hugging Face infrastructure and experiment with the scripts usingDev Mode with Spaces. One could also run the scripts with theHugging Face Jobs pipeline.

https://huggingface.co/blog/torch-attention-profile#naive-attentionNaive attention

Attention works with Queries (q), Keys (k), and Values (v). The interaction between them can be written as a short sequence of steps:

1. Build the attention scoresscores:matmul\(q, k\.T\)
2. Scale the scores:scores \* scale
3. Apply a causal mask to the scores:scores\.masked\_fill\(mask, "\-inf"\)
4. Normalize the scores with softmax to get the attention weightsattn:softmax\(scores\)
5. Reweight the values with those weights:matmul\(attn, v\)

So attention is really a collection of primitive operations. Some of them we already know (the matmuls), and the rest are easy to spot. Let’s write a naive attention module in PyTorch and profile it.

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

Before opening the trace, let’s do our usual exercise and guess what we should see. Tracing theforwardof this module, we expect:

• a matmul kernel (q \. k\.T)
• a mul kernel (the scaling)
• an operation for the masking
• a softmax kernel
• a matmul kernel (atten \. v)

uv run 04_a_naive_attention.py
uvx trace-util -f traces/ -b <hf_uname>/traces

Figure 1: The CPU lane of the profile trace for naive attention highlighting the discrete operations Figure 1 shows the CPU lane of the profile (the GPU lane is folded so it does not overwhelm us). Insideattn\_fwd(our annotated forward call) we can see exactly the operations we guessed. The matmul is an old friend by now, and the new operations are easy to spot:

• mul: the scaling
• masked\_fill: the causal masking
• softmax: the softmax kernel

Now let’s unfold the GPU lane and see which kernels were actually launched.

Figure 2: GPU and CPU lanes of the profile trace for naive attention highlighting a collection of kernels corresponding to one profiler step. Figure 2 shows the GPU lane next to the CPU lane. Let’s zoom into a singleattn\_fwdblock on the GPU lane to look at the kernels one by one.

Figure 3: Zoomed in GPU lane of the profiler trace for naive attention implementation. Figure 3 lets us read off the individual kernels for one profiler step:

1. matmul (query and key)
2. mul (scaling)
3. memory copy 🤔
4. causal masking
5. softmax (produces the attention weights)
6. matmul (attention weights and values)

Five of these are expected. The memory copy is the odd one out, so where does this come from? The clue is that PyTorch has in-place operations. When you operate on a tensor the ordinary (out-of-place) way, PyTorch often makes a copy, applies the requested operation to it, and returns the copy. Following the sequence of operations, the culprit here is ourmasked\_fill.

What if we replaced this with an in-place operation?

https://huggingface.co/blog/torch-attention-profile#naive-attention-with-inplace-causal-maskingNaive attention with inplace causal masking

All we change ismasked\_filltomasked\_fill\_(note the trailing underscore, PyTorch’s convention for in-place operations), and we run the same script.

def forward(self, q, k, v, mask):
    # q, k, v: [batch, heads, seq, head_dim]
    scores = torch.matmul(q, k.transpose(-2, -1))  # [batch, heads, seq, seq]
    scores = torch.mul(scores, self.scale)
-    scores = scores.masked_fill(mask, float("-inf"))
+    scores.masked_fill_(mask, float("-inf"))
    attn = torch.softmax(scores, dim=-1)
    out = torch.matmul(attn, v)  # [batch, heads, seq, head_dim]
    return out

Let’s look at the trace and see if something changed.

uv run 04_b_inplace_ops_attention.py
uvx trace-util -f traces/ -b <hf_uname>/traces

The in-place version (Figure 5) wraps far fewer CPU ops inside the masking step than the out-of-place version (Figure 4). This is an encouraging signal. Let’s unfold the GPU lane to confirm what happened there.

On the GPU lane theMemcpykernel is gone for good (Figures 6 and 7). With a one line change we shaved a whole kernel off each forward pass. This may not look like much on its own, but remember this is a single attention operation. In the context of a transformer based large model (LLMs, Diffusion models, etc.), it repeats once per layer, and there are many layers, so the saving adds up quickly (and if it earns you a raise, sharing at least 10% with us feels only fair).

Out-of-place is PyTorch’s default for a reason. To compute gradients, autograd has to remember the tensor values it saw on the forward pass, because many backward formulas reuse them. An in-place operation overwrites those values in memory, so the backward pass would read the wrong numbers. Due to the fact that we runforwardundertorch\.no\_grad, in-place is safe for us, with no backward pass and nothing to corrupt. It is also noteworthy that in-place operations do not only save time (like we see in our case) but also memory (due to no extra copy) which is great for large tensors like logits!

https://huggingface.co/blog/torch-attention-profile#scaled-dot-product-attentionScaled Dot Product Attention

We just built attention from primitives, and even shaved off aMemcpy. The good news is that the PyTorch team has done all of this for us, and packaged the whole pipeline into a single function:

from torch.nn import functional as F

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

This one line replaces our hand written module, andis\_causal=Trueeven saves us from building the mask by hand. It is worth pausing to appreciate how much this one call hides. And it hides more than just code lines. Scaled Dot Product Attention (SDPA) does not have a single implementation. Under the hood itdispatchesto one of the several backends and picks the fastest one that supports our inputs (dtype, head dimension, mask, hardware, etc.).

Theofficial SDPA tutorialwalks us through this selection, and the backends themselves are listed in thetorch\.nn\.attention\.SDPBackendenum:

from torch.nn.attention import SDPBackend

BACKENDS = {
    "math": SDPBackend.MATH,
    "flash": SDPBackend.FLASH_ATTENTION,
    "efficient": SDPBackend.EFFICIENT_ATTENTION,
    "cudnn": SDPBackend.CUDNN_ATTENTION,
}

Normally SDPA chooses for us, but we can pin a specific backend with thetorch\.nn\.attention\.sdpa\_kernelcontext manager. This is what we do in our scripts. This lets us profile each backend on its own and read how differently they show up in the trace. Let’s go one at a time.

https://huggingface.co/blog/torch-attention-profile#math-backendMath backend

uv run 04_c_sdpa_attention.py --backend math
uvx trace-util -f traces/ -b <hf_uname>/traces

Before we open anything, let’s guess. We have replaced hand written attention (matmul, mul, mask, softmax, matmul) with a single one liner, so we should expect the trace to getsimpler and faster. Fewer kernels, less CPU dispatch, maybe even a fused kernel. Let’s check the profiler table first.

MetricWhere to look?Naive in-placeSDPA math\*\_fwdCUDA time avgThe “CUDA time avg” column for the\*\_fwdop1.955 ms7.239 msSelf CUDA time totalAt the bottom of the profiler table7.194 ms27.279 ms This is our first surprise, the one liner is3\.7xslower.

Opening the trace (Figure 9) shows why the alarm bells ring, the math backend launches20GPU kernels per forward instead of the5launched with our naive attention implementation (Figure 8). This is the opposite of what we guessed. Let’s figure out why this happens.

https://huggingface.co/blog/torch-attention-profile#tensor-cores-left-vacantTensor cores left vacant

InPart 2we learned to read a kernel name like a fingerprint. Let’s use that habit here:

The A100s we used to capture these traces ship withTensor Cores, specialised hardware for accelerated matmuls that is known to be far faster than the ordinary CUDA cores. To see why that matters here, it helps to know what lives inside a GPU. A Streaming Multiprocessor (SM) is the compute unit of a GPU, and each SM has two kinds of arithmetic units, the CUDA cores and the Tensor Cores. CUDA cores are general purpose and process a handful of elements at a time, while Tensor Cores multiply and accumulate a whole small matrix tile in a single instruction. So the question is simple, “Is each backend actually using the fast path?”

The kernel names answer it. Thes16816in the naive kernel (Figure 10) is the signature of abfloat16Tensor Core matmul (the16x8x16Tensor Core instruction), so the naive version is on the fast path.sgemm(Figure 11) is the classic single precision (FP32) matmul that runs on the ordinary CUDA cores. In other words, the math backend never touches the Tensor Cores at all: to trade speed for numerical accuracy it upcasts tensors toFP32(doubling the data moved, even when the inputs are inbf16) and falls back to the slower CUDA cores.

https://huggingface.co/blog/torch-attention-profile#causal-masks-builtCausal masks built

In the naive version we built the causal mask once and reused it. Here we passedis\_causal=Trueand the math backend materialized one for us, oneverysingle call. You can watch it happen on the CPU lane:

Figure 12: CPU lane showing the ops for masking Here is what we see in Figure 12

aten::ones -> aten::tril            build a [seq, seq] lower-triangular matrix
aten::scalar_tensor -> aten::fill_  make the -inf fill value
aten::where                         turn it into an additive bias (0 or -inf)

On the GPU this shows up as atriu\_tril\_kernel, severalwherekernels, and anadd\_. The convenience flag that let us stop thinking about the mask did not remove the work, it just moved it one layer down, where the mask is rebuilt from scratch every forward.

https://huggingface.co/blog/torch-attention-profile#the-safe-softmaxThe safe softmax

Our hand written version called plainaten::softmax. The math backend callsaten::\_safe\_softmax, and the difference is again visible as extra kernels (Figure 13):

Figure 13: Safe softmax highlighting the extra kernels compared to generic softmax A row that is fully masked (every entry\-inf) would make an ordinary softmax computeexp\(\-inf\)/sum\(exp\(\-inf\)\) = 0/0 = NaN.\_safe\_softmaxguards against exactly that. Our naive kernel never bothered, and would have quietly producedNaNs in that corner case.

https://huggingface.co/blog/torch-attention-profile#so-what-is-the-math-backend-forSo what is the math backend for?

Put together, the math backend is the reference implementation. It is a straightforward, dtype-safe, NaN-safe decomposition of attention into primitive ATen ops. It is essentially the naive attention we wrote by hand, but more careful. That carefulness is exactly what makes it extremely slow.

Its job is not to be fast, but toalwayswork. This makes it the perfect baseline. Every backend we profile next (flash, efficient, cudnn) is trying to collapse the20GPU kernels into essentially one fused kernel that stays in bf16 and never materializes the intermediate matrices at all.

https://huggingface.co/blog/torch-attention-profile#efficient-backendEfficient backend

uv run 04_c_sdpa_attention.py --backend efficient
uvx trace-util -f traces -b <hf_uname>/traces

Figure 14: The profiler trace for sdpa with efficient backend Where the math backend launched 20 kernels across one profiler step, the efficient backend launches only onefmha\_cutlassF\_bf16\_aligned\_64x64\_rf\_sm80(as seen in Figure 14).

Let’s decode the name of the kernel:

• fmha(fused multi-head attention): All the primitive ops in attention is “fused” in one op now.
• cutlassF: built on CUTLASS (NVIDIA’s open-source templates for tensor-core GEMMs),Ffor forward.
• bf16\_aligned: runs in bfloat16 (no FP32 upcast, unlike math).
• 64x64: the tile size.
• rf(register file): the working set is kept in registers, the fastest memory on the chip.
• sm80: compiled for Ampere (the A100’s compute capability 8.0).

This is the memory efficient attention kernel that grew out of Meta’sxformerslibrary and was upstreamed into PyTorch. When people say “the xformers backend,” thisfmha\_cutlassFkernel is what they mean.

https://huggingface.co/blog/torch-attention-profile#flash-backendFlash backend

uv run 04_c_sdpa_attention.py --backend flash
uvx trace-util -f traces -b <hf_uname>/traces

Figure 15: The flash backend trace, one fusedpytorch\_flashkernel per forward Thevoid pytorch\_flashkernel (Figure 15) isFlashAttention-2(Tri Dao’s implementation), vendored into PyTorch.

Before we read the trace any further, it is worth answering the question you should be asking by now:why is there a whole backend named “flash”, and why does it matter so much?

https://huggingface.co/blog/torch-attention-profile#why-flash-attention-existsWhy flash attention exists?

Let’s go back to the math backend for a moment. Its real problem was not the count of 20 kernels, it was what those kernels handed to each other.

Step 1 builds the full score matrixattn = q \. k\.T, which is\[seq, seq\]per head. For a sequence length of 4096 that is4096 x 4096 ≈ 16 millionnumbers for a single head. That matrix is written out to the HBM (the GPU’s main memory), if there is even enough space to do so. Then, it is read back to be scaled, written again for the mask, read again for the softmax, and so on. Attention’s cost is dominated by thisback and forth traffic to HBM, not by the matmuls themselves.

FlashAttention attacks exactly this. Instead of computing the wholesmatrix and only then reducing it, it walks overkandvintiles, keeps a running softmax as it goes (the “online softmax” trick), and accumulates the output one tile at a time. The full\[seq, seq\]score matrix isnever written to HBM, it only ever lives on-chip. This is the single idea that lets the entire attention pipeline collapse into one fused kernel that stays in bf16 on the Tensor cores.

https://huggingface.co/blog/torch-attention-profile#why-flash-looks-wrong-under-the-profilerWhy flash looks “wrong” under the profiler

Figure 16: Estimated occupancy of flash kernel is seen to be 13% Here is where flash surprises people who read profiler footprints. It is the fastest backend, yet the profiler reports it with verylow occupancy(shown in Figure 16). To see why that is fine, we need three quick definitions.

A GPU kernel is essentially a series of instructions executed by many small execution units. These individual execution units (threads) take care of loading variables, adding them together, storing them back, etc. For each kernel, we launch many, many threads, and to keep track of them, we group them by blocks.

Blocks are scheduled onto Streaming Multiprocessors (SMs), the main compute units of a GPU. A block lives entirely on one SM, and an SM can host multiple blocks at onceif it has enough resources. Those resources include registers, shared memory, maximum resident threads, and maximum resident warps. So when we say a kernel has lowoccupancy, we mean each SM has fewer resident warps than it could theoretically support.

If you want to know more about threads, blocks, grids, etc. here is agreat resource.

If you click the flash kernel in the trace, its footprint tells the story (Figure 17).

Figure 17: The flash kernel footprint, heavy on registers and shared memory per block. Flash uses a lot of per-thread registers and a large amount of shared memory per block. For example, if a block has 128 threads and each thread uses 255 registers, that block needs128 × 255 = 32,640registers. On an Ampere SM with 65,536 registers, only two such blocks fit at once. Each 128-thread block has128 / 32 = 4warps, so two blocks give only 8 resident warps. Against a maximum of 64 resident warps, that is roughly 13% occupancy. Flash has low occupancy not because it is poorly optimized, but because each block is deliberately very “heavy” in on-chip resource usage.

And that is the whole point. High occupancy helpshide latencyby keeping many warps ready to run, but it does not make the work itself efficient. Flash spends those registers and that shared memory on purpose, to keep attention tiles on-chip, reuse data aggressively, and avoid ever materializing the full attention matrix in global memory.

https://huggingface.co/blog/torch-attention-profile#cudnn-backendcuDNN backend

uv run 04_c_sdpa_attention.py --backend cudnn
uvx trace-util -f traces -b <hf_uname>/traces

Figure 18: The cuDNN backend trace, a single generated attention kernel per forward. By now the pattern is familiar. Like flash and efficient, cuDNN gives us one fused, flash-style kernel per forward (Figure 18). So the natural question is:**if flash already fuses attention, why does PyTorch ship yet another flash backend?**The answer iswho writes the kernel and how it is built, and that difference is what makes the trace look different.

https://huggingface.co/blog/torch-attention-profile#how-is-cudnn-kernel-differentHow is cuDNN kernel different

Flash and efficient arefixed, pre-compiled kernelsvendored into PyTorch. You get the same binary every time. cuDNN is NVIDIA’s own deep learning library, and its attention kernel isgenerated and tuned for the specific problemat hand. It is closer in spirit totorch\.compile’s codegen than to a fixed cuBLAS binary. You can read that straight off the (very long) kernel name:

cudnn_generated_fort_native_sdpa_sm80_flash_fprop_wmma_f16_knob_6_128x64x64_4x1x1_cga1x1x1_kernel0_0

• cudnn\_generated: not a pre-shipped binary, it was generated by cuDNN.
• flash\_fprop: a flash attention style forward pass. So the algorithm is the same family as the flash backend.
• wmma\_f16: it uses the warp-level matrix multiply-accumulate (WMMA) API, the Tensor-core path on the 16-bit float pipeline.
• knob\_6: cuDNN picks from a set of pre-tuned configurations (“knobs”). Different shapes select different knobs, much like cuBLAS picking a tile variant.
• 128x64x64: the tile dimensions it chose.

That one fact,generated per problem, explains everything else that looks unusual in the trace.

1. No transposes: The CPU lane goes from\_cudnn\_attention\_forwardstraight to a couple ofaten::emptyallocations and then the kernel, with zeroaten::transpose(Figures 19, 20 and 21). Flash and efficient each insert four (metadata) transposes to reshape the tensors while cuDNN consumes the native\[B, H, S, D\]layout directly because its generator emits a kernel for that layout.
2. It launches throughcuLaunchKernelEx, notcudaLaunchKernel: Every other kernel in this whole series went through the runtime APIcudaLaunchKernel. cuDNN uses the driver-levelextendedlaunch, which carries launch attributes (Figure 22). Figure 22: CPU lane of the cuDNN backend showing the cuLaunchKernelEx driver-level launch instead of cudaLaunchKernel
3. The profiler reports 0% achieved occupancy: Do not take that at face value, it is a measurement gap, not a stalled GPU. CUPTI (the profiling backend) cannot attribute occupancy to a driver-API (cuLaunchKernelEx) launch the way it does forcudaLaunchKernel, so the field reads 0. The footprint fills in the truth (Figure 23):240 registers × 256 threads = 61,440registers per block against the SM’s 65,536, so onlyone blockfits per SM (8 warps ≈ 12.5%), right in line with flash. Figure 23: cuDNN kernel reporting 0% achieved occupancy, with 240 registers per thread and 256 threads per block

https://huggingface.co/blog/torch-attention-profile#the-cost-moved-to-the-cpuThe cost moved to the CPU

The “no transposes” story tempts us to expect cuDNN to be theleanestbackend on the CPU. It is the opposite.

backendCUDA avg timeCPU avg timeefficient277.9 µs117 µsflash146.8 µs138 µscudnn186.3 µs214 µs Even with zero transpose ops, cuDNN spends about214 µs per forward on the CPU, more than flash (138) or efficient (117). Almost all of it sits inaten::scaled\_dot\_product\_attentionself time (26% of the whole run) and\_cudnn\_attention\_forward. That is cuDNN’s runtime engine selecting and preparing the plan (the “knob” search) on every call.

Fewer visible ATen ops did not mean less CPU work, itmoved the work into the library, where the profiler can only show it as one fat, opaque bar. When a trace suddenly getscleaner, the work has not always disappeared, sometimes it has just moved somewhere the profiler cannot break down.

On the GPU, cuDNN (186.3 µs) lands between efficient and flash. On this very flash-friendly shape, hand-written FlashAttention-2 edges it out. cuDNN often wins onothershapes (larger head dimensions, different sequence lengths) precisely because its generator retunes per problem, but that retuning is also what you just paid for on the CPU.

https://huggingface.co/blog/torch-attention-profile#everything-we-covered-at-a-glanceEverything we covered, at a glance

Before we wrap up, here is a single table to review every attention variant we profiled and the one lesson each trace taught us.

VariantWhat we changedKernels / forwardWhat the trace revealedNaive attentionAttention built by hand from primitives (matmul, mul, mask, softmax, matmul)6A hiddenMemcpyfrom the out-of-placemasked\_fill.Naive in-placemasked\_fill→masked\_fill\_5One line drops theMemcpykernel entirely.SDPA mathF\.scaled\_dot\_product\_attentionpinned to the math backend20The reference: FP32 on CUDA cores, mask rebuilt every call,\_safe\_softmax. Correct but ~3.7x slower.SDPA efficientEfficient (xformers) backend1One fusedfmha\_cutlassFkernel, stays in bf16 on Tensor cores.SDPA flashFlash backend1One fusedpytorch\_flashkernel (FlashAttention-2). Fastest, despite “wrong-looking” 13% occupancy.SDPA cuDNNcuDNN backend1A per-problem generated kernel: no transposes,cuLaunchKernelEx, but the cost moved to a fat CPU bar.

https://huggingface.co/blog/torch-attention-profile#concluding-the-seriesConcluding the series

If you take away only one thing from the whole series, let it be the habit we repeated before every single trace which is toguess first, then look.

State out loud what you expect the trace to contain, open it, and treat any mismatch as the most interesting thing on the screen. Every real insight in these three posts, the hiddenMemcpy, theaddmmepilogue, the 20 kernel math backend, flash’s “wrong-looking” occupancy, cuDNN’s fat CPU bar, came from a guess that did not match the trace.

Profiling is not a separate, intimidating skill reserved for GPU experts. It is just the discipline of looking closely and asking “wait, why isthathappening?” until the answer clicks. You now have the vocabulary and the reflexes to do that on your own models. Open a trace, form a guess, and go find the mismatch.

Thanks for reading theProfiling in PyTorchseries. Now go profile something. 🤗

Thanks toNoe Flandrefor their reviews on the early draft of the post!

The blog post was polished using an LLM. This in no way means that we have let an agent run in the background and let it generate the blog. Some of us in the team are non-english speakers and think LLMs (which are mostly trained in the English Language) can rectify silly grammar mistakes or rephrase sentences that sound less intimidating and cleaner. Hope this helps with the idea of “why should I read, if this was LLM generated”. 🤗