Profiling in PyTorch (Half 3): Consideration is all you profile

The collection “Profiling in PyTorch” is supposed to make you comfy studying profiler traces and tables. In Half 1 we profiled primary math operations like addition and multiplication. We noticed how the profiler desk uncovers hotspots, and the way the profiler hint reveals the order through which an algorithm runs over time.

In Half 2 we wrapped that addition and multiplication right into a torch linear layer. We then stacked a number of linear layers on prime of one another (a multilayer perceptron) and profiled that. Alongside the way in which we additionally profiled fused and hand-tuned kernels.

From the attitude of the Transformer structure, the subsequent logical step for us to profile is one more basic algorithm, consideration. Whereas being notorious for its quadratic-time complexity, many intelligent methods exist to mitigate that challenge and make it quick. Our objective right here is to not cowl each trick intimately. As an alternative, we wish to see how each appears completely different beneath the profiler.

The scripts for this weblog publish stay right here: 04_a_naive_attention.py, 04_b_inplace_ops_attention.py, 04_c_sdpa_attention.py, and 04_d_kernels_attention.py. Like earlier than, it helps to open them in a separate tab and stroll by means of the code as you learn. We use an NVIDIA A100-SXM4-80GB GPU to run the scripts. It’s very easy to arrange a GPU on the Hugging Face infrastructure and experiment with the scripts utilizing Dev Mode with Areas. One may additionally run the scripts with the Hugging Face Jobs pipeline.

Naive consideration

Consideration works with Queries (q), Keys (ok), and Values (v). The interplay between them might be written as a brief sequence of steps:

Construct the eye scores scores: matmul(q, ok.T)
Scale the scores: scores * scale
Apply a causal masks to the scores: scores.masked_fill(masks, “-inf”)
Normalize the scores with softmax to get the eye weights attn: softmax(scores)
Reweight the values with these weights: matmul(attn, v)

So consideration can be a assortment of primitive operations. A few of them we already know (the matmuls), and the remainder are simple to identify. Let’s write a naive consideration module in PyTorch and profile it.

class NaiveCausalAttention(nn.Module):
def __init__(self, head_dim):
tremendous().__init__()
self.scale = 1.0 / math.sqrt(head_dim)

def ahead(self, q, ok, v, masks):
scores = torch.matmul(q, ok.transpose(-2, –1))
scores = scores * self.scale
scores = scores.masked_fill(masks, float(“-inf”))
attn = torch.softmax(scores, dim=-1)
out = torch.matmul(attn, v)
return out

Earlier than opening the hint, let’s do our traditional train and guess what we should always see. Tracing the ahead of this module, we anticipate:

a matmul kernel (q . ok.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 /traces

Determine 1 reveals the CPU lane of the profile (the GPU lane is folded so it doesn’t overwhelm us). Inside attn_fwd (our annotated ahead name) we will see precisely the operations we guessed. The matmul is an previous good friend by now, and the brand new operations are simple to identify:

mul: the scaling
masked_fill: the causal masking
softmax: the softmax kernel

Now let’s unfold the GPU lane and see which kernels have been truly launched.

Determine 2 reveals the GPU lane subsequent to the CPU lane. Let’s zoom right into a single attn_fwd block on the GPU lane to have a look at the kernels one after the other.

Determine 3 lets us learn off the person kernels for one profiler step:

matmul (question and key)
mul (scaling)
reminiscence copy 🤔
causal masking
softmax (produces the eye weights)
matmul (consideration weights and values)

5 of those are anticipated. The reminiscence copy is the odd one out, so the place does this come from? The clue is that PyTorch has in-place operations. While you function on a tensor the unusual (out-of-place) manner, PyTorch usually makes a replica, applies the requested operation to it, and returns the copy. Following the sequence of operations, the offender right here is our masked_fill.

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

Naive consideration with inplace causal masking

All we alter is masked_fill to masked_fill_ (word the trailing underscore, PyTorch’s conference for in-place operations), and we run the identical script.

def ahead(self, q, ok, v, masks):
# q, ok, v: [batch, heads, seq, head_dim]
scores = torch.matmul(q, ok.transpose(-2, -1)) # [batch, heads, seq, seq]
scores = torch.mul(scores, self.scale)
– scores = scores.masked_fill(masks, float(“-inf”))
+ scores.masked_fill_(masks, float(“-inf”))
attn = torch.softmax(scores, dim=-1)
out = torch.matmul(attn, v) # [batch, heads, seq, head_dim]
return out

Let us take a look at the hint and see if one thing modified.

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

The in-place model (Determine 5) wraps far fewer CPU ops contained in the masking step than the out-of-place model (Determine 4). That is an encouraging sign. Let’s unfold the GPU lane to verify what occurred there.

On the GPU lane the Memcpy kernel is gone for good (Figures 6 and seven). With a one line change we shaved a complete kernel off every ahead move. This may occasionally not seem like a lot by itself, however keep in mind it is a single consideration operation. Within the context of a transformer based mostly giant mannequin (LLMs, Diffusion fashions, and so forth.), it repeats as soon as per layer, and there are a lot of layers, so the saving provides up shortly (and if it earns you a increase, sharing at the least 10% with us feels solely truthful).

Out-of-place is PyTorch’s default for a cause. To compute gradients, autograd has to recollect the tensor values it noticed on the ahead move, as a result of many backward formulation reuse them. An in-place operation overwrites these values in reminiscence, so the backward move would learn the mistaken numbers. Because of the truth that we run ahead beneath torch.no_grad, in-place is secure for us, with no backward move and nothing to deprave. Additionally it is noteworthy that in-place operations don’t solely save time (like we see in our case) but in addition reminiscence (as a consequence of no further copy) which is nice for big tensors like logits!

Scaled Dot Product Consideration

We simply constructed consideration from primitives, and even shaved off a Memcpy. The excellent news is that the PyTorch group has finished all of this for us, and packaged the entire pipeline right into a single operate:

from torch.nn import useful as F

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

This one line replaces our hand written module, and is_causal=True even saves us from constructing the masks by hand. It’s price pausing to understand how a lot this one name hides. And it hides extra than simply code traces. Scaled Dot Product Consideration (SDPA) doesn’t have a single implementation. Below the hood it dispatches to one of many a number of backends and picks the quickest one which helps our inputs (dtype, head dimension, masks, {hardware}, and so forth.).

The official SDPA tutorial walks us by means of this choice, and the backends themselves are listed within the torch.nn.consideration.SDPBackend enum:

from torch.nn.consideration import SDPBackend

BACKENDS = {
“math”: SDPBackend.MATH,
“flash”: SDPBackend.FLASH_ATTENTION,
“environment friendly”: SDPBackend.EFFICIENT_ATTENTION,
“cudnn”: SDPBackend.CUDNN_ATTENTION,
}

Usually SDPA chooses for us, however we will pin a particular backend with the torch.nn.consideration.sdpa_kernel context supervisor. That is what we do in our scripts. This lets us profile every backend by itself and skim how in a different way they present up within the hint. Let’s go one by one.

Math backend

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

Earlier than we open something, let’s guess. We’ve changed hand written consideration (matmul, mul, masks, softmax, matmul) with a single one liner, so we should always anticipate the hint to get easier and sooner. Fewer kernels, much less CPU dispatch, perhaps even a fused kernel. Let’s verify the profiler desk first.

That is our first shock, the one liner is 3.7x slower.

Opening the hint (Determine 9) reveals why the alarm bells ring, the mathematics backend launches 20 GPU kernels per ahead as a substitute of the 5 launched with our naive consideration implementation (Determine 8). That is the other of what we guessed. Let’s work out why this occurs.

Tensor cores left vacant

In Half 2 we discovered to learn a kernel title like a fingerprint. Let’s use that behavior right here:

The A100s we used to seize these traces ship with Tensor Cores, specialised {hardware} for accelerated matmuls that’s identified to be far sooner than the unusual CUDA cores. To see why that issues right here, it helps to know what lives inside a GPU. A Streaming Multiprocessor (SM) is the compute unit of a GPU, and every SM has two sorts of arithmetic models, the CUDA cores and the Tensor Cores. CUDA cores are common objective and course of a handful of components at a time, whereas Tensor Cores multiply and accumulate a complete small matrix tile in a single instruction. So the query is easy, “Is every backend truly utilizing the quick path?”

The kernel names reply it. The s16816 within the naive kernel (Determine 10) is the signature of a bfloat16 Tensor Core matmul (the 16x8x16 Tensor Core instruction), so the naive model is on the quick path. sgemm (Determine 11) is the basic single precision (FP32) matmul that runs on the unusual CUDA cores. In different phrases, the mathematics backend by no means touches the Tensor Cores in any respect: to commerce pace for numerical accuracy it upcasts tensors to FP32 (doubling the info moved, even when the inputs are in bf16) and falls again to the slower CUDA cores.

Causal masks constructed

Within the naive model we constructed the causal masks as soon as and reused it. Right here we handed is_causal=True and the mathematics backend materialized one for us, on each single name. You may watch it occur on the CPU lane:

Here’s what we see in Determine 12

aten::ones -> aten::tril construct a [seq, seq] lower-triangular matrix
aten::scalar_tensor -> aten::fill_ make the -inf fill worth
aten::the place flip it into an additive bias (0 or -inf)

On the GPU this reveals up as a triu_tril_kernel, a number of the place kernels, and an add_. The comfort flag that permit us cease desirous about the masks didn’t take away the work, it simply moved it one layer down, the place the masks is rebuilt from scratch each ahead.

The secure softmax

Our hand written model known as plain aten::softmax. The mathematics backend calls aten::_safe_softmax, and the distinction is once more seen as further kernels (Determine 13):

A row that’s totally masked (each entry -inf) would make an unusual softmax compute exp(-inf)/sum(exp(-inf)) = 0/0 = NaN. _safe_softmax guards towards precisely that. Our naive kernel by no means bothered, and would have quietly produced NaNs in that nook case.

So what’s the math backend for?

Put collectively, the mathematics backend is the reference implementation. It’s a easy, dtype-safe, NaN-safe decomposition of consideration into primitive ATen ops. It’s basically the naive consideration we wrote by hand, however extra cautious. That carefulness is precisely what makes it extraordinarily sluggish.

Its job is to not be quick, however to at all times work. This makes it the right baseline. Each backend we profile subsequent (flash, environment friendly, cudnn) is making an attempt to break down the 20 GPU kernels into basically one fused kernel that stays in bf16 and by no means materializes the intermediate matrices in any respect.

Environment friendly backend

uv run 04_c_sdpa_attention.py –backend environment friendly
uvx trace-util -f traces -b /traces

The place the mathematics backend launched 20 kernels throughout one profiler step, the environment friendly backend launches just one fmha_cutlassF_bf16_aligned_64x64_rf_sm80 (as seen in Determine 14).

Let’s decode the title of the kernel:

fmha (fused multi-head consideration): All of the primitive ops in consideration is “fused” in a single op now.
cutlassF: constructed on CUTLASS (NVIDIA’s open-source templates for tensor-core GEMMs), F for ahead.
bf16_aligned: runs in bfloat16 (no FP32 upcast, not like math).
64×64: the tile measurement.
rf (register file): the working set is saved in registers, the quickest reminiscence on the chip.
sm80: compiled for Ampere (the A100’s compute functionality 8.0).

That is the reminiscence environment friendly consideration kernel that grew out of Meta’s xformers library and was upstreamed into PyTorch. When individuals say “the xformers backend,” this fmha_cutlassF kernel is what they imply.

Flash backend

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

The void pytorch_flash kernel (Determine 15) is FlashAttention-2 (Tri Dao’s implementation), vendored into PyTorch.

Earlier than we learn the hint any additional, it’s price answering the query you ought to be asking by now: why is there a complete backend named “flash”, and why does it matter a lot?

Why flash consideration exists?

Let’s return to the mathematics backend for a second. Its actual drawback was not the depend of 20 kernels, it was what these kernels handed to one another.

Step 1 builds the complete rating matrix attn = q . ok.T, which is [seq, seq] per head. For a sequence size of 4096 that’s 4096 x 4096 ≈ 16 million numbers for a single head. That matrix is written out to the HBM (the GPU’s major reminiscence), if there may be even sufficient area to take action. Then, it’s learn again to be scaled, written once more for the masks, learn once more for the softmax, and so forth. Consideration’s price is dominated by this forwards and backwards site visitors to HBM, not by the matmuls themselves.

FlashAttention assaults precisely this. As an alternative of computing the entire s matrix and solely then decreasing it, it walks over ok and v in tiles, retains a operating softmax because it goes (the “on-line softmax” trick), and accumulates the output one tile at a time. The complete [seq, seq] rating matrix is rarely written to HBM, it solely ever lives on-chip. That is the only concept that lets the whole consideration pipeline collapse into one fused kernel that stays in bf16 on the Tensor cores.

Why flash appears “mistaken” beneath the profiler

Right here is the place flash surprises individuals who learn profiler footprints. It’s the quickest backend, but the profiler studies it with very low occupancy (proven in Determine 16). To see why that’s high-quality, we’d like three fast definitions.

A GPU kernel is actually a collection of directions executed by many small execution models. These particular person execution models (threads) care for loading variables, including them collectively, storing them again, and so forth. For every kernel, we launch many, many threads, and to maintain observe of them, we group them by blocks.

Blocks are scheduled onto Streaming Multiprocessors (SMs), the principle compute models of a GPU. A block lives fully on one SM, and an SM can host a number of blocks without delay if it has sufficient sources. These sources embody registers, shared reminiscence, most resident threads, and most resident warps. So after we say a kernel has low occupancy, we imply every SM has fewer resident warps than it may theoretically assist.

If you wish to know extra about threads, blocks, grids, and so forth. right here is a superb useful resource.

In case you click on the flash kernel within the hint, its footprint tells the story (Determine 17).

Flash makes use of a variety of per-thread registers and a considerable amount of shared reminiscence per block. For instance, if a block has 128 threads and every thread makes use of 255 registers, that block wants 128 × 255 = 32,640 registers. On an Ampere SM with 65,536 registers, solely two such blocks match without delay. Every 128-thread block has 128 / 32 = 4 warps, so two blocks give solely 8 resident warps. Towards a most of 64 resident warps, that’s roughly 13% occupancy. Flash has low occupancy not as a result of it’s poorly optimized, however as a result of every block is intentionally very “heavy” in on-chip useful resource utilization.

And that’s the complete level. Excessive occupancy helps disguise latency by holding many warps able to run, however it doesn’t make the work itself environment friendly. Flash spends these registers and that shared reminiscence on objective, to maintain consideration tiles on-chip, reuse knowledge aggressively, and keep away from ever materializing the complete consideration matrix in world reminiscence.

cuDNN backend

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

By now the sample is acquainted. Like flash and environment friendly, cuDNN offers us one fused, flash-style kernel per ahead (Determine 18). So the pure query is: if flash already fuses consideration, why does PyTorch ship one more flash backend? The reply is who writes the kernel and the way it’s constructed, and that distinction is what makes the hint look completely different.

How is cuDNN kernel completely different

Flash and environment friendly are fastened, pre-compiled kernels vendored into PyTorch. You get the identical binary each time. cuDNN is NVIDIA’s personal deep studying library, and its consideration kernel is generated and tuned for the precise drawback at hand. It’s nearer in spirit to torch.compile’s codegen than to a hard and fast cuBLAS binary. You may learn that straight off the (very lengthy) kernel title:

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 consideration type ahead move. So the algorithm is similar household because the flash backend.
wmma_f16: it makes use of 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”). Completely different shapes choose completely different knobs, very like cuBLAS choosing a tile variant.
128x64x64: the tile dimensions it selected.

That one truth, generated per drawback, explains every part else that appears uncommon within the hint.

No transposes: The CPU lane goes from _cudnn_attention_forward straight to a few aten::empty allocations after which the kernel, with zero aten::transpose (Figures 19, 20 and 21). Flash and environment friendly every insert 4 (metadata) transposes to reshape the tensors whereas cuDNN consumes the native [B, H, S, D] format straight as a result of its generator emits a kernel for that format.

It launches by means of cuLaunchKernelEx, not cudaLaunchKernel: Each different kernel on this complete collection went by means of the runtime API cudaLaunchKernel. cuDNN makes use of the driver-level prolonged launch, which carries launch attributes (Determine 22).

The profiler studies 0% achieved occupancy: Don’t take that at face worth, it’s a measurement hole, not a stalled GPU. CUPTI (the profiling backend) can’t attribute occupancy to a driver-API (cuLaunchKernelEx) launch the way in which it does for cudaLaunchKernel, so the sphere reads 0. The footprint fills within the fact (Determine 23): 240 registers × 256 threads = 61,440 registers per block towards the SM’s 65,536, so just one block suits per SM (8 warps ≈ 12.5%), proper consistent with flash.

The price moved to the CPU

The “no transposes” story tempts us to anticipate cuDNN to be the leanest backend on the CPU. It’s the reverse.

Even with zero transpose ops, cuDNN spends about 214 µs per ahead on the CPU, greater than flash (138) or environment friendly (117). Nearly all of it sits in aten::scaled_dot_product_attention self time (26% of the entire run) and _cudnn_attention_forward. That’s cuDNN’s runtime engine choosing and making ready the plan (the “knob” search) on each name.

Fewer seen ATen ops didn’t imply much less CPU work, it moved the work into the library, the place the profiler can solely present it as one fats, opaque bar. When a hint immediately will get cleaner, the work has not at all times disappeared, generally it has simply moved someplace the profiler can’t break down.

On the GPU, cuDNN (186.3 µs) lands between environment friendly and flash. On this very flash-friendly form, hand-written FlashAttention-2 edges it out. cuDNN usually wins on different shapes (bigger head dimensions, completely different sequence lengths) exactly as a result of its generator retunes per drawback, however that retuning can also be what you simply paid for on the CPU.

All the things we lined, at a look

Earlier than we wrap up, here’s a single desk to evaluate each consideration variant we profiled and the one lesson every hint taught us.

Concluding the collection

In case you take away just one factor from the entire collection, let it’s the behavior we repeated earlier than each single hint which is to guess first, then look.

State out loud what you anticipate the hint to comprise, open it, and deal with any mismatch as probably the most attention-grabbing factor on the display. Each actual perception in these three posts, the hidden Memcpy, the addmm epilogue, the 20 kernel math backend, flash’s “wrong-looking” occupancy, cuDNN’s fats CPU bar, got here from a guess that didn’t match the hint.

Profiling shouldn’t be a separate, intimidating talent reserved for GPU consultants. It’s simply the self-discipline of trying carefully and asking “wait, why is that taking place?” till the reply clicks. You now have the vocabulary and the reflexes to try this by yourself fashions. Open a hint, type a guess, and go discover the mismatch.

Thanks for studying the Profiling in PyTorch collection. Now go profile one thing. 🤗

Due to Noe Flandre for his or her critiques on the early draft of the publish!

The weblog publish was polished utilizing an LLM. This by no means signifies that we have now let an agent run within the background and let it generate the weblog. A few of us within the group are non-english audio system and suppose LLMs (that are principally skilled within the English Language) can rectify foolish grammar errors or rephrase sentences that sound much less intimidating and cleaner. Hope this helps with the concept of “why ought to I learn, if this was LLM generated”. 🤗