Attention 的理论复杂度是二次方,但线上性能并不只由公式决定。Hugging Face 最新的 PyTorch Profiling 系列文章,从朴素 Attention 开始,对比原地因果遮罩、Scaled Dot Product Attention(SDPA)以及 math、efficient、Flash 和 cuDNN 等后端在 profiler 中的表现。

为什么一定要看 trace

同一段高层 PyTorch 代码,在不同输入形状、精度、GPU 和后端选择下,可能落到完全不同的 CUDA kernel。只看 Python 执行时间,通常无法判断耗时来自矩阵乘法、Softmax、掩码、内存复制还是 kernel 启动。

Profiler table 适合定位累计耗时最高的操作,timeline trace 则能看到 kernel 的顺序、空隙和并行关系。两者结合,才有机会回答“GPU 为什么没跑满”。

一个实用的分析顺序

面对 Attention 性能问题,可以按下面的顺序排查:

  1. 固定 batch、序列长度、head 数和数据类型,建立可重复基线。
  2. 同时记录 CPU 与 CUDA 活动,确认是否存在主机侧等待。
  3. 检查 SDPA 实际选择的后端,不要假设一定使用 Flash Attention。
  4. 观察显存分配与中间张量,判断是否受内存带宽限制。
  5. 改动一个变量后重新测量,避免把多个优化混在一起。

原地操作可能减少额外内存,但也可能影响 autograd 或编译优化;更快的 kernel 也可能对形状、精度和硬件有约束。因此优化必须以目标工作负载为准。

对工程实践的意义

模型推理成本往往不是“模型太大”这么简单。序列长度分布、动态 batching、KV Cache 和 Attention 后端都会影响吞吐与延迟。性能优化的第一步不是立即换框架或写自定义 kernel,而是先用 trace 建立证据。

这篇教程最值得带走的方法是:从最简单实现开始逐层替换,每次用 profiler 验证。只有能解释性能变化,优化结果才可能在生产环境里稳定复现。