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

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

这篇博文使用 NVIDIA A100-SXM4-80GB GPU 对 PyTorch 中的注意力机制进行分析，并展示了每种变体在分析器（profiler）追踪结果中的表现。

Naive attention

使用原始操作（matmul, mul, masked_fill, softmax, matmul）构建的朴素注意力实现，每次前向传播会启动六个 GPU kernel，其中包括一个由非原地（out-of-place）masked_fill 引起的意外内存复制。

Naive attention with inplace causal masking

将 masked_fill 替换为原地（in-place）的 masked_fill_ 可以移除内存复制 kernel，将每次前向传播的 GPU kernel 数量从六个减少到五个。

Scaled Dot Product Attention

PyTorch 的 F.scaled_dot_product_attention 提供了一个单行接口，可分发到多个后端；每个后端都显示出截然不同的分析特征。

Math backend

当固定在 math 后端时，SDPA 每次前向传播会启动约二十个 GPU kernel，在 CUDA 核心上以 FP32 运行，每次调用都会重建因果掩码（causal mask），并使用 _safe_softmax 来避免 NaNs；它比朴素的原地版本慢大约 3.7 倍。

Efficient backend

高效后端 (xformers) 每次前向传播启动单个融合的 fmha_cutlassF_bf16_aligned_64x64_rf_sm80 kernel，在 Tensor cores 上保持 bfloat16 计算，并避免向 HBM 写入中间矩阵。

Flash backend

Flash 后端每次前向传播启动单个融合的 pytorch_flash kernel (FlashAttention-2)；尽管在分析器中显示的估计占用率较低（约 13%），但它是最快的后端，因为它使用寄存器和共享内存将注意力分块（tiles）保留在芯片上。

cuDNN backend

cuDNN 后端每次前向传播启动单个生成的 kernel（例如 cudnn_generated_fort_native_sdpa_sm80_flash_fprop_wmma_f16_knob_6_128x64x64_4x1x1_cga1x1x1_kernel0_0）；它避免了在 CPU 上进行转置操作，但由于运行时计划选择（runtime plan selection）而产生了较高的 CPU 时间（~214 µs），其 GPU 时间介于高效后端和 flash 后端之间。

Everything we covered, at a glance

Variant	What we changed	Kernels / forward	What the trace revealed
Naive attention	Attention built by hand from primitives (matmul, mul, mask, softmax, matmul)	6	A hidden Memcpy from the out-of-place masked_fill.
Naive in-place	masked_fill → masked_fill_	5	One line drops the Memcpy kernel entirely.
SDPA math	F.scaled_dot_product_attention pinned to the math backend	20	The reference: FP32 on CUDA cores, mask rebuilt every call, _safe_softmax. Correct but ~3.7× slower.
SDPA efficient	Efficient (xformers) backend	1	One fused fmha_cutlassF kernel, stays in bf16 on Tensor cores.
SDPA flash	Flash backend	1	One fused pytorch_flash kernel (FlashAttention-2). Fastest, despite "wrong-looking" 13% occupancy.
SDPA cuDNN	cuDNN backend	1	A per-problem generated kernel: no transposes, cuLaunchKernelEx, but the cost moved to a fat CPU bar.

Concluding the series

核心结论是：先对分析器追踪结果应该显示什么形成一个假设，然后检查追踪结果，并将任何不匹配之处视为最有价值的洞察；这种习惯揭示了隐藏的 Memcpy、math 后端的二十个 kernel、flash 的低占用率以及 cuDNN 的 CPU 端开销。

现在你可以应用这种“假设并验证”的方法来分析你自己的模型了。