Hugging Face 工程团队 7 月 10 日发布「Profiling in PyTorch」系列第 3 篇,在 A100 上用 torch.profiler 把 PyTorch SDPA 的 math、efficient(xformers)、flash(FlashAttention-2)、cuDNN 四个后端逐一跑过,挖出几条「直觉与 trace 对不上」的发现。 第一条:`masked_fill` 改成原地 `masked_fill_`,trace 里多出来的就是隐藏的 Memcpy。out-of-place 操作先复制再改写,原地版在 `no_grad` 下直接覆盖原 tensor 内存,对 logits 这种大张量既省 kernel 又省显存。改一行代码就砍掉整个 Memcpy,在多层 Transformer 里会叠加放大。 第二条:`F.scaled_dot_product_attention` 一行顶五行,但 math 后端一次 forward 启 20 个 kernel,比手写还慢 3.7 倍。原因是 math 后端为数值稳定把输入升到 FP32、绕开 Tensor Core,而且每次调用都重建 causal mask。它的角色是参考实现,不是性能路径。 第三条:四个 flash-style 后端都只跑一个 fused kernel,但 trace 上显示的「13% 占用率」很容易让人误判成性能差。flash 用 128 thread × 255 register/block,每个 SM 只能装下两个 block,理论占用率约 13%——这正是 FlashAttention 把 `[seq, seq]` score 矩阵永远钉在片上、不写回 HBM 的代价,不是「它慢」。 cuDNN 走另一条路:按输入 shape 现生成 kernel、用 `cuLaunchKernelEx` 驱动 API 直发、连 transpose 都省;代价是 plan 搜索全在 CPU,A100 这一组 shape 上 cuDNN(186.3µs)反而比 flash(146.8µs)慢。Profiler 报告「trace 变干净」并不等于「工作消失」,有时只是工作搬进了库函数看不到的地方。 整个系列的方法论只有一句:**先猜,再打开 trace。** 所有值得分享的洞见——隐藏的 Memcpy、20 个 kernel 的 math backend、flash 错位的占用率、cuDNN 的 CPU 大柱子——都来自「猜测和 trace 对不上」的那一瞬间。