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 性能问题,可以按下面的顺序排查:
- 固定 batch、序列长度、head 数和数据类型,建立可重复基线。
- 同时记录 CPU 与 CUDA 活动,确认是否存在主机侧等待。
- 检查 SDPA 实际选择的后端,不要假设一定使用 Flash Attention。
- 观察显存分配与中间张量,判断是否受内存带宽限制。
- 改动一个变量后重新测量,避免把多个优化混在一起。
原地操作可能减少额外内存,但也可能影响 autograd 或编译优化;更快的 kernel 也可能对形状、精度和硬件有约束。因此优化必须以目标工作负载为准。
对工程实践的意义
模型推理成本往往不是“模型太大”这么简单。序列长度分布、动态 batching、KV Cache 和 Attention 后端都会影响吞吐与延迟。性能优化的第一步不是立即换框架或写自定义 kernel,而是先用 trace 建立证据。
这篇教程最值得带走的方法是:从最简单实现开始逐层替换,每次用 profiler 验证。只有能解释性能变化,优化结果才可能在生产环境里稳定复现。