From primitives to production attention
Hugging Face just published the third installment of their "Profiling in PyTorch" series, and it's exactly the kind of deep technical dive the AI infrastructure community needs right now. Attention is all you profile walks through how attention mechanisms—the beating heart of transformers—look under PyTorch's profiler, kernel by kernel.
The series follows a pedagogical arc that actually works: Part 1 covered basic ops like addition and multiplication, Part 2 tackled nn.Linear and fused MLPs, and now Part 3 arrives at attention itself. The progression mirrors how you'd actually build understanding of transformer internals, bottom-up.
What makes this post valuable isn't just explaining what attention does—we all know the queries, keys, and values dance by now. It's showing you how to read what your GPU is actually doing when you run attention, and why seemingly small code changes produce radically different performance profiles.
Naive attention: counting kernels
The post starts with a hand-rolled attention implementation that matches the textbook formula: matmul queries and keys, scale the scores, apply a causal mask, softmax to get attention weights, then matmul those weights with values.
Before profiling, the authors ask you to predict what you'll see. This is smart pedagogy—it forces you to build intuition. You'd expect five discrete steps: two matmuls, one multiply for scaling, masking, and a softmax.
The CPU lane of the profiler trace confirms exactly that. Each operation shows up as a distinct block: matmul, mul, masked_fill, softmax, matmul again.
But then they unfold the GPU lane, and something odd appears: a memory copy kernel that shouldn't be there. The culprit? PyTorch's masked_fill makes a copy of the tensor before masking it, then returns the copy. That's an extra kernel launch and an extra round-trip to global memory.
The underscore that saved a kernel
Here's where the post gets practical. PyTorch has a convention: operations ending in an underscore (like masked_fill_) are in-place. They modify the tensor directly instead of copying it first.
Changing scores.masked_fill(mask, float("-inf")) to scores.masked_fill_(mask, float("-inf")) eliminates the memory copy entirely. The profiler trace confirms it—the Memcpy kernel just disappears from the GPU lane.
One character, one fewer kernel per forward pass. In a 32-layer transformer running thousands of times per second, that adds up fast. The post notes this is safe here because they're running under torch.no_grad, so there's no backward pass to corrupt.
This is the kind of micro-optimization that separates hobby code from production inference. The authors have a sense of humor about it: "if it earns you a raise, sharing at least 10% with us feels only fair."
SDPA: the one-liner that hides everything
PyTorch ships a built-in function for scaled dot-product attention: F.scaled_dot_product_attention(q, k, v, is_causal=True). It replaces all the hand-written primitives with a single call and even builds the causal mask for you.
But here's the twist: SDPA doesn't have one implementation. It dispatches to different backends depending on your hardware, dtype, and inputs. The backends—math, flash attention, efficient attention, and cuDNN—show up in wildly different ways under the profiler.
The post pins each backend using torch.nn.attention.sdpa_kernel context manager and profiles them individually. This is where things get interesting.
Math backend: the slow safe path
The math backend is 3.7× slower than the naive in-place version. The profiler table shows it: 7.2ms versus 1.9ms per forward pass.
Opening the trace reveals why: the math backend launches 20 GPU kernels per attention operation instead of 5. It also never touches the A100's Tensor Cores—the specialized hardware units that make matmuls fast.
You can read this in the kernel names. The naive version shows s16816, the signature of a bfloat16 Tensor Core matmul using the 16×8×16 tile size. The math backend shows sgemm, the classic single-precision FP32 matmul that runs on ordinary CUDA cores.
The math backend upcasts everything to FP32 for numerical accuracy, doubling the data moved even when inputs are bf16. It also materializes a new causal mask on every call instead of reusing one. The profiler's CPU lane shows aten::ones followed by aten::tril building that mask over and over.
This backend exists for correctness and reproducibility, not speed. But seeing exactly how much it costs—kernel by kernel—is the point of profiling.
Flash, efficient, and cuDNN: production backends
The post walks through the other three backends with the same kernel-level scrutiny. Each one shows distinct tradeoffs.
Flash Attention uses a custom fused kernel that avoids materializing the full attention matrix in memory. The efficient backend applies similar memory-saving tricks. cuDNN uses NVIDIA's heavily optimized library kernels.
The profiler traces show these backends launching far fewer kernels than the math backend. They also hit the Tensor Cores consistently, keeping the fast arithmetic units busy instead of falling back to CUDA cores.
What's valuable here isn't just knowing "Flash is faster"—everyone knows that by now. It's being able to see why it's faster by reading the kernel launches, memory operations, and which hardware units are actually doing the work.
Building the skill that matters
The real contribution of this series is teaching you to read profiler output like a primary source. Not trusting benchmarks or vibes, but opening the trace and counting kernels yourself.
The post is opinionated about methodology: always predict what you'll see before opening the trace, learn to read kernel names as fingerprints of what hardware is active, and use the profiler table alongside the visual trace to cross-check your understanding.
This matters because transformer performance engineering is increasingly about choosing between pre-built attention implementations rather than writing your own. But "just use Flash Attention" isn't helpful if you don't know when it's actually being used, or what to check when performance doesn't match expectations.
The scripts live in a public repo, captured on an A100-SXM4-80GB. The post encourages you to run them yourself using Hugging Face Spaces with Dev Mode or the Jobs pipeline. Profiling is a hands-on skill—reading about it helps, but you need to generate and read your own traces to build intuition.
Where this fits in the bigger picture
This series arrives at a good time. We're in an era where model architectures are converging but inference optimization is diverging wildly. Kernel-level profiling is how you close the gap between theoretical FLOPs and actual throughput.
The attention mechanism is where most of the compute goes in transformers, and it's where most of the optimization leverage lives. Understanding how different implementations show up in the profiler is table-stakes for anyone doing serious inference work.
The Hugging Face team is clearly building toward something—this is Part 3, and the progression from basic ops to linear layers to attention suggests more complex patterns are coming. Multi-head attention, grouped-query attention, or full end-to-end transformer profiling would be logical next steps.
For now, this post gives you a practical framework: write naive attention, profile it, apply in-place ops, profile again, switch to SDPA, profile each backend, and learn to read what actually happened at the kernel level. That's the tight loop that builds real understanding.
If you're working on transformer inference, this is required reading. And if you're not profiling your attention kernels yet, you probably should be.