认知 Transformer 的正确姿势先看懂算子再谈优化做 AI 工程的人大概都有过这样的阶段模型结构背得烂熟多头注意力的公式能默写但一到模型变慢、显存爆掉、推理延迟压不下来的时候就陷入玄学调参——换 batch size、改精度、试各种 feature最后也说不清瓶颈到底在哪。问题的根源在于我们一直在“结构”层面理解 Transformer却很少站到“算子执行图”的视角去看它。同样一个 Attention在 PyTorch 里写着很简洁q k.transpose(-2, -1) * scale但到了 GPU 上它要被拆成若干 kernel 逐个调度每个 kernel 都要经历“数据搬运到显存→核心计算→结果写回显存”的完整过程。你看到的几行代码实际是几十次显存读写和 kernel 启动开销的叠加。这就是 Transformer 大量使用算子对硬件性能形成挑战的本质结构越复杂算子越碎计算效率就越依赖执行引擎的调度能力。这次分享我想用“第一性原理”的方式把 Transformer 重新拆解一遍。不背公式不抄框架而是回到最底层的问题每个算子到底在算什么、为什么要这么算、计算下界在哪、瓶颈在哪。搞懂了这一层你看模型的目光会完全不同。适合谁看如果你在写推理引擎、优化模型性能、处理长序列训练或者只是想知道自己 8 张卡跑出来的训练曲线为什么跟论文对不上这篇文章都很合适。它会帮你把“模型”拆成一组可以精确量化的“算子序列”让你在性能排查和架构改进时有凭有据。1. 为什么要把 Transformer 拆成“算子”来看1.1 从模型到算子你写的 PyTorch 代码在硬件上到底跑了什么先说一个经常被忽视的事实深度学习框架里的神经网络运行时并不是一个整体执行的“模型”而是一个由算子Operator / Kernel按依赖关系串起来的计算图。Transformer 只是对这个计算图的组织方式做了某种特定的描述真正在硬件上干活的是图中一个个独立的算子。拿编码器里最常见的一行代码举例比如hidden_states attn_output residual在程序视角它是一个残差连接一个非常“高层的语义”但在执行图视角它可能被拆成把 attn_output 和 residual 分别从显存读到寄存器执行 elementwise 加法再写回显存。如果这里还考虑了 dropout、LayerNorm 的融合算子边界又会变化。也就是说同一个模型结构在不同框架、不同推理后端里实际跑出来的算子图甚至可以被重排成完全不同的样子。这就要求我们建立一种“算子粒度”的思维不要问“Attention 为什么慢”而要问“Attention 被拆成哪几个算子哪个算子的访存量最大哪个算子的并行度最低”。这是从宏观模型走向微观性能分析的必经之路。明白了这一点后续所有关于执行图、融合、显存优化的讨论才有地基。1.2 第一性原理的三件事实计算量、访存量、并行度用第一性原理审视 Transformer不需要一开始就陷入复杂公式只需要抓住三个根本事实。第一任何算子都有计算量。所谓计算量就是这个算子需要执行多少次浮点运算。一般用 FLOPs 衡量Transformer 里矩阵乘法占绝对大头。第二任何算子都有访存量即从显存/HBM 读取和写回多少字节。这两个量的比值叫算术强度Arithmetic Intensity它决定了一个算子在硬件上到底是“计算密集”还是“访存密集”。第三任何算子都有并行度它可以同时被分解成多少条独立的流水线交给 GPU 的不同核心去算。Transformer 里矩阵乘法的并行度天然很高LayerNorm 这种归约操作并行度就会弱一些Softmax 中间还要穿插归一化并行模式更复杂。把所有算子都放到这三个维度的坐标轴上去审视你会发现很多概念的轮廓突然变得清晰。为什么说 Transformer 大量使用算子对硬件性能提出挑战因为它的执行图里既有 GEMM 这样典型的计算密集算子也有 LayerNorm、Softmax、残差、Dropout 这样高度访存密集的 elementwise/归约算子。这两类算子对硬件的需求是矛盾的前者要拼命堆算力后者要拼命提升带宽而现代 GPU 更擅长前者却也在持续改善后者。理解这种矛盾是进行任何性能优化的前提。2. Transformer 核心算子族拆解每个算子都在解决一个问题把一个标准 Encoder 结构打开抛开花哨的变体你会发现核心算子其实非常克制一只手数得过来查表、矩阵乘、缩放、Softmax、LayerNorm、残差加法、GELU、矩阵转置和相乘。先把这个清单列清楚后续所有性能分析才有着落。2.1 嵌入表示层与位置编码查表算子与常量算子输入 token 序列进入模型的第一步是把每个 token 映射成稠密向量。这一步在算子层面叫Embedding本质上是一次gather查表操作根据 token id 去权重矩阵里按索引取出一行向量。它没有浮点计算只有纯访存所以这是一个典型访存密集算子。如果词表很大比如 100KEmbedding 表本身会占用相当可观的显存这也解释了为什么大词汇表模型光 Embedding 参数就能到上亿规模。需要特别注意的是这种查表操作无法充分利用 GPU 的算力它的性能瓶颈在显存带宽和数据读取效率上。之后是位置编码。经典 Transformer 用的正弦余弦位置编码在执行图里实际上是一个常量生成算子。也就是说每个位置的向量是预先通过固定公式算好的常量在前向时直接相加即可。而可学习位置编码Learnable Positional Embedding则和 token embedding 一样是查表操作。但这里我建议你把位置编码视作“常量算子”来管理因为它在一次推理中保持不变可以在显存里预分配避免每次前向重复计算。到了 RoPE旋转位置编码时代情况发生改变它不再是在 embedding 上加一个位置向量而是把位置信息编码进 Q 和 K 的复数旋转里执行图里就多出角度计算、reshape、split、cat、逐元素乘法和加法等操作。这会让算子图变得稀疏细碎如果你在使用 RoPE 时发现前向变慢多半不是旋转计算本身慢而是这些切片拼接带来的 kernel 调度开销。要注意的是嵌入层还有个最常见的隐藏成本——能否把 Embedding 参数放到离计算单元更近的位置。在推理阶段如果只用 GPUEmbedding 表很小没什么问题但如果你做 CPU 推理或者混合异构推理Embedding 表可能成为内存访问的瓶颈。这些都属于“查表算子的执行图设计”范畴是从模型定义中看不见的。2.2 多头注意力MatMul、Scale 的三角关系与 Softmax 的特殊角色注意力机制是 Transformer 的灵魂执行图也是整个模型里最复杂的部分。拆开看主要有五个步骤Q、K、V 各自的线性变换GEMMQK^T 点积GEMM缩放elementwiseSoftmax 行归一化归约逐元素以及 Softmax 结果与 V 相乘GEMM。如果按执行顺序把它们写成算子序列最直观的写法是Q X W_q # GEMM, MxK KxN K X W_k # GEMM V X W_v # GEMM S Q K^T # GEMM, MxN NxM S S * scale # elementwise P softmax(S, dim-1) # reduce elementwise O P V # GEMM这串序列里有几个值得深剖的点。S 矩阵即 attention score形状是[batch, heads, seq_len, seq_len]。当序列长度增长时这部分显存按平方膨胀是长序列任务首先爆显存的位置。而且 S 的计算本身就包含一次大矩阵乘和一次转置转置操作在 GEMM kernel 中通常通过lda/ldb步长参数隐式实现不需要真的把数据搬一遍这一点是很多人在手写优化时容易踩的坑。Softmax 的执行并非简单“先指数再归一”它需要两次数据扫描第一次扫描寻找该行的最大值用于数值稳定第二次计算指数和总和。这两次扫描都发生在同一份 S 数据上如果 S 太大无法驻留 L2 缓存就会造成两次 HBM 访问让软乘法的访存开销翻倍。这也是为什么 FlashAttention 要把分块 Softmax 的“在线归一化”技术拿出来用重计算的方式减少一次 S 的读写。这里有个容易忽略的关键算子scale为什么会被单独提出来而不是直接乘到 Q 或者 K 里从数学上讲缩放可以提前作用到 Q 或 K但实操中为了保护 Softmax 的数值范围通常会在 QK^T 之后立刻乘1/sqrt(d_k)。第一性原理的问题是缩放会不会带来额外 kernel 开销如果框架的算子融合做得好scale 会并进前面的 GEMM epilogue 或 Softmax 的初始化过程几乎零开销如果融合不好它就是独立的一次 elementwise kernel意味着完整读写一遍 S 矩阵。同一个操作优化前后可以相差一个完整矩阵的访存量这是执行图视角给我们的第一个震动。2.3 LayerNorm 与残差连接被低估的访存瓶颈很多人在做模型加速时第一时间盯上 GEMM反而对 LayerNorm 和残差连接视而不见。但在算子执行图里这俩恰恰是制造访存压力的常客。原因很简单LayerNorm 是一个逐行归约算子它需要对输入矩阵的每一行求均值、方差然后逐元素做归一化和仿射变换。这个过程至少需要两次完整的显存读取写完中间结果后还需要一次写入而这一切几乎没有任何可复用的浮点计算量。用一个具体数据来看假设hidden_size 4096seq_len 2048batch 8那么 LayerNorm 处理的数据量约为8 * 2048 * 4096 * 2 bytes 134 MB以 fp16 计。如果总访存量按“读两次写一次”算大约是 400MB 的访问量。相比一个典型 GEMM计算量高达几十 TFLOPs但访存量通常在几百 MB 以内LayerNorm 的算术强度长期处于极低水平。如果 Transformer 块里叠加多个 LayerNorm它的执行时间可能占到整个 block 的 10%~20%一点都不夸张。残差连接的代价相对小一些毕竟只是两个张量相加读两个数写一个数但千万别忽略它的存在。如果执行图没有把残差融合到 LayerNorm 或者后续算子中你会在每一步都看到一串 elementwise 内核add、returnadd、return。每一个内核的启动都有开销小到几微秒但 Transformer 层数一多几十上百个小内核堆积起来绝对是你 profile 结果里那根刺眼的短柱。这也是我强烈建议在工程实现中启用算子融合比如把残差加进 LayerNorm 的参数重归一化流程或加进 GEMM 的 epilogue从执行图上把这些独立的“小算子”消灭掉。2.4 FFN 与 GELUTransformer 里真正的算力黑洞大部分人在讨论 Transformer 时把注意力放在 Attention 上但若计算整个模型的前向 FLOPsFFN前馈神经网络通常占了一半以上。标准 FFN 包含两个 GEMM 和一个激活函数x - Linear(4h) - GELU - Linear(h)。以h 4096为例FFN 中间维度是16384这个 GEMM 的规模远大于 Attention 里的任何一次 GEMM。GELU 的数学形式是GELU(x) 0.5 * x * (1 erf(x / sqrt(2)))这里涉及erf即误差函数。如果严格按数学定义去实现这是一个计算开销比较大的“元素级函数”而且不同类型硬件对erf的支持程度不一样。所以工程上有各种近似比如GELU(x) ≈ 0.5 * x * (1 tanh(sqrt(2/pi) * (x 0.044715 * x^3)))这就是常见的 tanh 近似。这两种实现方式在数据集上精度差异非常小但执行图上的计算量差异却很明显。如果你用 CUDA 实现纯 tanh 近似涉及多次乘法和一次 tanh 计算而erf近似在 CUDA 里有erff指令硬件直接支持实际也不慢。关键要看你的目标平台如果跑 GPU直接用erf数学实现或许更好如果跑某些加速卡或 CPUtanh 近似可能更划算。这就是“从第一性原理选择算子实现”的典型例子。GELU 在计算图里是 elementwise 算子意味着它要按元素读一遍、计算一遍、写一遍。如果把它融合到后面那个 GEMM 的 epilogue尾处理里这个读写开销就可以被消除。PyTorch 老版本里nn.GELU()是一个独立 kernel新版在推理模式中有时也会自动融合。但如果你自己写推理引擎这里是最容易产生一个不必要算子节点的位置。3. 执行图上的真实战场算子融合、显存规划与硬件映射3.1 为什么执行图是“图”而不是“列表”一个朴素的执行器拿到一串算子可能会挨个执行先算 Q、K、V再算 QK^T再 Softmax…… 但这未必最优。因为算子之间存在依赖关系有些可以并行有些可以合并有些可以重排。这就是“执行图”比“执行列表”更能表达并行机会的原因。举个例子Embedding 输出后通常紧接着会同时计算 Q、K、V 和残差分支在图上它们是并行的分支理论上有多个 kernel 可以在不同流上并发执行前提是硬件资源允许。再比如某些独立分支的计算顺序可以交换先算 FFN 再算 Attention或者反过来只要依赖关系不被破坏即可。可惜实际优化中很少人在计算图层面做这种全局重排因为收益不如算子融合直接。但作为工程人员理解执行图为有向无环图DAG依然有实际价值。你在用 PyTorch Profiler 或 TensorRT 的 graph dump 时看到的算子并不是线性列表而是带有依赖关系的图。借助编译器的调度器比如 ONNX Runtime 的 execution plan 或者 XLA 在编译期对算子块做融合、内存复用它们本质上都是在“图”上做文章。算子融合的直接收益是减少 kernel 启动次数和访存量显存复用则解决的是图生命周期分析问题——某个算子的中间结果如果后续不再被使用它的显存可以立即被后续张量回收。这也是为什么一个部署模型的性能好坏可能在你导出 ONNX 那一刻就已经注定了。我用过 TensorRT 转换一些模型得到的明显加速很大程度并非来自硬件优化而是来自图优化它会自动做层融合、权重合并、常量折叠把 PyTorch 里那些零零碎碎的小算子收敛成一整块大 kernel。3.2 三个必须掌握的融合策略QKV 融合、FlashAttention 与 LayerNorm 融合第一个策略把计算 Q、K、V 的三个权重矩阵合并成一个大的W_qkv矩阵。这样原本三次 GEMM 变成一次执行图节点数和 kernel 启动次数直接变为原来的三分之一。这既是模型定义层面的改动也是执行图优化最经典的例子。PyTorch 实现里可以写成self.qkv nn.Linear(d_model, 3 * d_model, biasFalse)然后把输出的 chunk 分成 Q、K、V。效果不只是在代码层面更简洁在运行时是一次更大的矩阵乘能更好发挥 Tensor Core 的并行能力。第二个策略FlashAttention。它针对的核心问题是 Attention 算子链中 S 矩阵的反复读写。FlashAttention 做了两件大事按块计算让 S 的中间分块驻留在 SRAM 中引入在线 Softmaxonline softmax通过维护行最大值、行指数和来避免对整个 S 做全局归约。这样 Attention 的前向过程中原始 S 矩阵绝大部分时间不用落回 HBMsoftmax 的两次扫描也只剩一次。实测中当 seq_len 达到 2048 以上FlashAttention 相对原始 PyTorch 实现的加速可以达到 2~3 倍甚至更多同时显存占用大幅下降。第三个策略LayerNorm 残差的融合以及 LayerNorm 和它后面的 GEMM 融合。在x x Attn(LN(x))的结构中LN 的输出同时被残差加法和后续 GEMM 使用如果把 LN 的计算结果直接通过 GEMM 的 epilogue 输出到残差缓冲途中至少省掉一次中间张量的显存写读。对于训练场景算子融合需要谨慎因为反向传播时需要重新利用部分中间结果但对于推理场景这种融合几乎只有好处。3.3 从执行图视角估算一次前向传播的时间这一节我们亲自动手动笔感受一下“从第一性原理估算执行时间”的爽快感。假设模型配置为d_model 4096ff_dim 16384num_heads 32d_k 128seq_len 2048batch 4num_layers 24fp16 推理。先估算 FLOPs。一点共识一次矩阵乘法的 FLOPs 约为2 * M * N * K。各主要部分每层 QKV 线性输入尺寸[8192, 4096]权重[4096, 12288]FLOPs 2 * 8192 * 4096 * 12288 ≈ 0.82 TFLOPs。Attention score 计算Q 尺寸[8192, 32, 128]K 转置后与之相乘等效 FLOPs 2 * 8192 * 8192 * 128 ≈ 17.2 GFLOPs。这里和 QKV 相比小了一个数量级。Attention output 投影O[8192, 4096]权重[4096, 4096]FLOPs 2 * 8192 * 4096 * 4096 ≈ 0.27 TFLOPs。FFN 两个线性层第一个2 * 8192 * 4096 * 16384 ≈ 1.1 TFLOPs第二个2 * 8192 * 16384 * 4096 ≈ 1.1 TFLOPs。每层总计约 3.3 TFLOPs24 层约 80 TFLOPs。在一张 A100 80Gfp16 理论上限约 312 TFLOPs实际利用率取 45%的算力水平下理论最优耗时约为80 / (312 * 0.45) ≈ 0.57 s。但实际跑一次前向你用 PyTorch eager 模式大约要 1.5~2 秒。多出来的时间去哪里了答案就在执行图里。这 24 层里还包含 48 个 LayerNorm、48 个残差 add、24 个 Softmax、大量 dropout/bias 相加。它们加起来本身计算量不到 1 TFLOPs但访存量却可能超过 200 GB。以 A100 的 HBM 带宽约 2 TB/s 来算这些算子访存时间就要约 0.1 s。更要命的是PyTorch eager 模式在 GPU 上每启动一个算子 kernel 有几十微秒的开销当整个网络由上千个算子构成时这个调度开销会变成 0.3~0.5 秒的显性成本。因此真实耗时与理论峰值之间的差距大多由低算术强度算子和 kernel 调度开销造成。这个估算方法你自己可以套到任何模型上。4. 踩坑实录从算子粒度出发的调试心得4.1 显存空洞与峰值计算为什么同样的模型别人显存少一半有一次我在调一个视频理解模型序列长度拉长后一个 batch 显存直接爆掉。把执行图输出 dump 下来发现问题不在权重而在大量中间激活值的一个峰值Attention score 矩阵在 seq_len 很长时按平方增长。结论很朴素——用 sublinear attention 或 FlashAttention 能直接把这个峰值压掉。但这里我真正想强调的是显存空洞问题频繁分配、释放的中间张量在显存中会产生碎片峰值时比实际需求高不少。处理技巧有三个。第一在 PyTorch 里尽量复用已有缓冲区避免反复new_empty或torch.empty。第二跑推理时固定形状可用torch.jit或 TensorRT 图优化把中间张量的生命周期尽量缩短。第三注意权重和激活分开管理显存训练时用梯度检查点技术激活重计算以减少存储激活的空间这一招能在几乎不损失性能的情况下让 batch 增大一两倍。执行图优化的目的不是让单个算子变快而是让所有张量的生命周期尽量短显存尽量复用。4.2 数值稳定性Softmax 的溢出、LayerNorm 的 eps、fp16 的表示范围写算子不是数学题你还得处理浮点数本身的脾气。Softmax 的经典数值稳定方法是在每行减去最大值即考察exp(x - max)。这个技巧在算子实现中不可省略。如果是 fp16 训练或推理exp对绝对值很大的负数会直接下溢为 0一整个行全变成 0 的话梯度回传也全部消失。所以实现 Softmax kernel 时我的建议是不管 FP16、BF16、FP32统一先减最大值再算哪怕因此多做一次读取也比溢出导致的不稳定要划算得多。LayerNorm 里的eps参数也别乱动。它出现在分母的方差中(x - mean) / sqrt(var eps)。如果 eps 设得太小fp16 下方差较小的时候除出来异常大或出现除零误报设大了归一化精度会受影响。常见做法是eps1e-5到1e-6之间但在 fp16 下试过不少模型我建议用1e-5起步不要因为 QuickStart 代码里是 1e-12 就直接抄那是给 fp32 用的。4.3 用等价算子替换法定位性能瓶颈最后分享一个非常实用的调试技巧等价算子替换法。当我觉得某个算子开销不合理的时候先把模型里的算子用手写等价版本替换再重新跑 profile对比前后时间变化。例如把 GELU 换成恒等映射x - x看整体耗时降多少就能反推出 GELU 及其相关访存的开销把 Attention 里的 Softmax 换成恒等映射同理。这是做算子级性能归因最简单直接的方法比分析理论计算量来得更快。做算子替换的原则是“一次只动一个变量”保证其它层都保持不变否则无法判断瓶颈。实际使用中我发现 PyTorch Profiler 的图表视图已经很直观配合torch.profiler的 kernel 级输出可以很清楚地看到每个算子占用的总时间、GPU 时间、调用次数。如果某个算子调用次数不多但 GPU 时间占比很高优先优化它如果调用次数非常多但每个都很短优先想融合。方向对了性能优化的大头就抓到了。一点个人心得算子级性能分析切忌只盯着 FLOPs。真正的瓶颈通常藏在访存开销和 kernel 启动开销里也就是“低算术强度的算子”和“碎小算子”这两个地方。Transformer 的结构决定了这两类问题必然存在——它的成功来自高度的结构表达力而它的性能压力也恰恰来自结构带来的算子多寡与访存开销。理解这一层才会理解“算子融合”、“内存复用”、“图优化”为什么是部署必选项而不是锦上添花。最后再分享一个落地建议新项目里别等模型写好后再做性能优化。在模型结构定稿前先把执行图草稿画出来——哪些算子可以合并哪些中间张量可以省掉哪一步会形成显存峰值。你在这个阶段能解决的问题比事后跑 profiler 再去修要省一半力气。Transformer 十年前第一次出现时我们还在为精度欢呼现在它遍地开花性能优化的“第一性原理”也该回到算子执行图这个最本质的层面了。
Transformer性能优化关键：从算子执行图看透访存与融合本质
认知 Transformer 的正确姿势先看懂算子再谈优化做 AI 工程的人大概都有过这样的阶段模型结构背得烂熟多头注意力的公式能默写但一到模型变慢、显存爆掉、推理延迟压不下来的时候就陷入玄学调参——换 batch size、改精度、试各种 feature最后也说不清瓶颈到底在哪。问题的根源在于我们一直在“结构”层面理解 Transformer却很少站到“算子执行图”的视角去看它。同样一个 Attention在 PyTorch 里写着很简洁q k.transpose(-2, -1) * scale但到了 GPU 上它要被拆成若干 kernel 逐个调度每个 kernel 都要经历“数据搬运到显存→核心计算→结果写回显存”的完整过程。你看到的几行代码实际是几十次显存读写和 kernel 启动开销的叠加。这就是 Transformer 大量使用算子对硬件性能形成挑战的本质结构越复杂算子越碎计算效率就越依赖执行引擎的调度能力。这次分享我想用“第一性原理”的方式把 Transformer 重新拆解一遍。不背公式不抄框架而是回到最底层的问题每个算子到底在算什么、为什么要这么算、计算下界在哪、瓶颈在哪。搞懂了这一层你看模型的目光会完全不同。适合谁看如果你在写推理引擎、优化模型性能、处理长序列训练或者只是想知道自己 8 张卡跑出来的训练曲线为什么跟论文对不上这篇文章都很合适。它会帮你把“模型”拆成一组可以精确量化的“算子序列”让你在性能排查和架构改进时有凭有据。1. 为什么要把 Transformer 拆成“算子”来看1.1 从模型到算子你写的 PyTorch 代码在硬件上到底跑了什么先说一个经常被忽视的事实深度学习框架里的神经网络运行时并不是一个整体执行的“模型”而是一个由算子Operator / Kernel按依赖关系串起来的计算图。Transformer 只是对这个计算图的组织方式做了某种特定的描述真正在硬件上干活的是图中一个个独立的算子。拿编码器里最常见的一行代码举例比如hidden_states attn_output residual在程序视角它是一个残差连接一个非常“高层的语义”但在执行图视角它可能被拆成把 attn_output 和 residual 分别从显存读到寄存器执行 elementwise 加法再写回显存。如果这里还考虑了 dropout、LayerNorm 的融合算子边界又会变化。也就是说同一个模型结构在不同框架、不同推理后端里实际跑出来的算子图甚至可以被重排成完全不同的样子。这就要求我们建立一种“算子粒度”的思维不要问“Attention 为什么慢”而要问“Attention 被拆成哪几个算子哪个算子的访存量最大哪个算子的并行度最低”。这是从宏观模型走向微观性能分析的必经之路。明白了这一点后续所有关于执行图、融合、显存优化的讨论才有地基。1.2 第一性原理的三件事实计算量、访存量、并行度用第一性原理审视 Transformer不需要一开始就陷入复杂公式只需要抓住三个根本事实。第一任何算子都有计算量。所谓计算量就是这个算子需要执行多少次浮点运算。一般用 FLOPs 衡量Transformer 里矩阵乘法占绝对大头。第二任何算子都有访存量即从显存/HBM 读取和写回多少字节。这两个量的比值叫算术强度Arithmetic Intensity它决定了一个算子在硬件上到底是“计算密集”还是“访存密集”。第三任何算子都有并行度它可以同时被分解成多少条独立的流水线交给 GPU 的不同核心去算。Transformer 里矩阵乘法的并行度天然很高LayerNorm 这种归约操作并行度就会弱一些Softmax 中间还要穿插归一化并行模式更复杂。把所有算子都放到这三个维度的坐标轴上去审视你会发现很多概念的轮廓突然变得清晰。为什么说 Transformer 大量使用算子对硬件性能提出挑战因为它的执行图里既有 GEMM 这样典型的计算密集算子也有 LayerNorm、Softmax、残差、Dropout 这样高度访存密集的 elementwise/归约算子。这两类算子对硬件的需求是矛盾的前者要拼命堆算力后者要拼命提升带宽而现代 GPU 更擅长前者却也在持续改善后者。理解这种矛盾是进行任何性能优化的前提。2. Transformer 核心算子族拆解每个算子都在解决一个问题把一个标准 Encoder 结构打开抛开花哨的变体你会发现核心算子其实非常克制一只手数得过来查表、矩阵乘、缩放、Softmax、LayerNorm、残差加法、GELU、矩阵转置和相乘。先把这个清单列清楚后续所有性能分析才有着落。2.1 嵌入表示层与位置编码查表算子与常量算子输入 token 序列进入模型的第一步是把每个 token 映射成稠密向量。这一步在算子层面叫Embedding本质上是一次gather查表操作根据 token id 去权重矩阵里按索引取出一行向量。它没有浮点计算只有纯访存所以这是一个典型访存密集算子。如果词表很大比如 100KEmbedding 表本身会占用相当可观的显存这也解释了为什么大词汇表模型光 Embedding 参数就能到上亿规模。需要特别注意的是这种查表操作无法充分利用 GPU 的算力它的性能瓶颈在显存带宽和数据读取效率上。之后是位置编码。经典 Transformer 用的正弦余弦位置编码在执行图里实际上是一个常量生成算子。也就是说每个位置的向量是预先通过固定公式算好的常量在前向时直接相加即可。而可学习位置编码Learnable Positional Embedding则和 token embedding 一样是查表操作。但这里我建议你把位置编码视作“常量算子”来管理因为它在一次推理中保持不变可以在显存里预分配避免每次前向重复计算。到了 RoPE旋转位置编码时代情况发生改变它不再是在 embedding 上加一个位置向量而是把位置信息编码进 Q 和 K 的复数旋转里执行图里就多出角度计算、reshape、split、cat、逐元素乘法和加法等操作。这会让算子图变得稀疏细碎如果你在使用 RoPE 时发现前向变慢多半不是旋转计算本身慢而是这些切片拼接带来的 kernel 调度开销。要注意的是嵌入层还有个最常见的隐藏成本——能否把 Embedding 参数放到离计算单元更近的位置。在推理阶段如果只用 GPUEmbedding 表很小没什么问题但如果你做 CPU 推理或者混合异构推理Embedding 表可能成为内存访问的瓶颈。这些都属于“查表算子的执行图设计”范畴是从模型定义中看不见的。2.2 多头注意力MatMul、Scale 的三角关系与 Softmax 的特殊角色注意力机制是 Transformer 的灵魂执行图也是整个模型里最复杂的部分。拆开看主要有五个步骤Q、K、V 各自的线性变换GEMMQK^T 点积GEMM缩放elementwiseSoftmax 行归一化归约逐元素以及 Softmax 结果与 V 相乘GEMM。如果按执行顺序把它们写成算子序列最直观的写法是Q X W_q # GEMM, MxK KxN K X W_k # GEMM V X W_v # GEMM S Q K^T # GEMM, MxN NxM S S * scale # elementwise P softmax(S, dim-1) # reduce elementwise O P V # GEMM这串序列里有几个值得深剖的点。S 矩阵即 attention score形状是[batch, heads, seq_len, seq_len]。当序列长度增长时这部分显存按平方膨胀是长序列任务首先爆显存的位置。而且 S 的计算本身就包含一次大矩阵乘和一次转置转置操作在 GEMM kernel 中通常通过lda/ldb步长参数隐式实现不需要真的把数据搬一遍这一点是很多人在手写优化时容易踩的坑。Softmax 的执行并非简单“先指数再归一”它需要两次数据扫描第一次扫描寻找该行的最大值用于数值稳定第二次计算指数和总和。这两次扫描都发生在同一份 S 数据上如果 S 太大无法驻留 L2 缓存就会造成两次 HBM 访问让软乘法的访存开销翻倍。这也是为什么 FlashAttention 要把分块 Softmax 的“在线归一化”技术拿出来用重计算的方式减少一次 S 的读写。这里有个容易忽略的关键算子scale为什么会被单独提出来而不是直接乘到 Q 或者 K 里从数学上讲缩放可以提前作用到 Q 或 K但实操中为了保护 Softmax 的数值范围通常会在 QK^T 之后立刻乘1/sqrt(d_k)。第一性原理的问题是缩放会不会带来额外 kernel 开销如果框架的算子融合做得好scale 会并进前面的 GEMM epilogue 或 Softmax 的初始化过程几乎零开销如果融合不好它就是独立的一次 elementwise kernel意味着完整读写一遍 S 矩阵。同一个操作优化前后可以相差一个完整矩阵的访存量这是执行图视角给我们的第一个震动。2.3 LayerNorm 与残差连接被低估的访存瓶颈很多人在做模型加速时第一时间盯上 GEMM反而对 LayerNorm 和残差连接视而不见。但在算子执行图里这俩恰恰是制造访存压力的常客。原因很简单LayerNorm 是一个逐行归约算子它需要对输入矩阵的每一行求均值、方差然后逐元素做归一化和仿射变换。这个过程至少需要两次完整的显存读取写完中间结果后还需要一次写入而这一切几乎没有任何可复用的浮点计算量。用一个具体数据来看假设hidden_size 4096seq_len 2048batch 8那么 LayerNorm 处理的数据量约为8 * 2048 * 4096 * 2 bytes 134 MB以 fp16 计。如果总访存量按“读两次写一次”算大约是 400MB 的访问量。相比一个典型 GEMM计算量高达几十 TFLOPs但访存量通常在几百 MB 以内LayerNorm 的算术强度长期处于极低水平。如果 Transformer 块里叠加多个 LayerNorm它的执行时间可能占到整个 block 的 10%~20%一点都不夸张。残差连接的代价相对小一些毕竟只是两个张量相加读两个数写一个数但千万别忽略它的存在。如果执行图没有把残差融合到 LayerNorm 或者后续算子中你会在每一步都看到一串 elementwise 内核add、returnadd、return。每一个内核的启动都有开销小到几微秒但 Transformer 层数一多几十上百个小内核堆积起来绝对是你 profile 结果里那根刺眼的短柱。这也是我强烈建议在工程实现中启用算子融合比如把残差加进 LayerNorm 的参数重归一化流程或加进 GEMM 的 epilogue从执行图上把这些独立的“小算子”消灭掉。2.4 FFN 与 GELUTransformer 里真正的算力黑洞大部分人在讨论 Transformer 时把注意力放在 Attention 上但若计算整个模型的前向 FLOPsFFN前馈神经网络通常占了一半以上。标准 FFN 包含两个 GEMM 和一个激活函数x - Linear(4h) - GELU - Linear(h)。以h 4096为例FFN 中间维度是16384这个 GEMM 的规模远大于 Attention 里的任何一次 GEMM。GELU 的数学形式是GELU(x) 0.5 * x * (1 erf(x / sqrt(2)))这里涉及erf即误差函数。如果严格按数学定义去实现这是一个计算开销比较大的“元素级函数”而且不同类型硬件对erf的支持程度不一样。所以工程上有各种近似比如GELU(x) ≈ 0.5 * x * (1 tanh(sqrt(2/pi) * (x 0.044715 * x^3)))这就是常见的 tanh 近似。这两种实现方式在数据集上精度差异非常小但执行图上的计算量差异却很明显。如果你用 CUDA 实现纯 tanh 近似涉及多次乘法和一次 tanh 计算而erf近似在 CUDA 里有erff指令硬件直接支持实际也不慢。关键要看你的目标平台如果跑 GPU直接用erf数学实现或许更好如果跑某些加速卡或 CPUtanh 近似可能更划算。这就是“从第一性原理选择算子实现”的典型例子。GELU 在计算图里是 elementwise 算子意味着它要按元素读一遍、计算一遍、写一遍。如果把它融合到后面那个 GEMM 的 epilogue尾处理里这个读写开销就可以被消除。PyTorch 老版本里nn.GELU()是一个独立 kernel新版在推理模式中有时也会自动融合。但如果你自己写推理引擎这里是最容易产生一个不必要算子节点的位置。3. 执行图上的真实战场算子融合、显存规划与硬件映射3.1 为什么执行图是“图”而不是“列表”一个朴素的执行器拿到一串算子可能会挨个执行先算 Q、K、V再算 QK^T再 Softmax…… 但这未必最优。因为算子之间存在依赖关系有些可以并行有些可以合并有些可以重排。这就是“执行图”比“执行列表”更能表达并行机会的原因。举个例子Embedding 输出后通常紧接着会同时计算 Q、K、V 和残差分支在图上它们是并行的分支理论上有多个 kernel 可以在不同流上并发执行前提是硬件资源允许。再比如某些独立分支的计算顺序可以交换先算 FFN 再算 Attention或者反过来只要依赖关系不被破坏即可。可惜实际优化中很少人在计算图层面做这种全局重排因为收益不如算子融合直接。但作为工程人员理解执行图为有向无环图DAG依然有实际价值。你在用 PyTorch Profiler 或 TensorRT 的 graph dump 时看到的算子并不是线性列表而是带有依赖关系的图。借助编译器的调度器比如 ONNX Runtime 的 execution plan 或者 XLA 在编译期对算子块做融合、内存复用它们本质上都是在“图”上做文章。算子融合的直接收益是减少 kernel 启动次数和访存量显存复用则解决的是图生命周期分析问题——某个算子的中间结果如果后续不再被使用它的显存可以立即被后续张量回收。这也是为什么一个部署模型的性能好坏可能在你导出 ONNX 那一刻就已经注定了。我用过 TensorRT 转换一些模型得到的明显加速很大程度并非来自硬件优化而是来自图优化它会自动做层融合、权重合并、常量折叠把 PyTorch 里那些零零碎碎的小算子收敛成一整块大 kernel。3.2 三个必须掌握的融合策略QKV 融合、FlashAttention 与 LayerNorm 融合第一个策略把计算 Q、K、V 的三个权重矩阵合并成一个大的W_qkv矩阵。这样原本三次 GEMM 变成一次执行图节点数和 kernel 启动次数直接变为原来的三分之一。这既是模型定义层面的改动也是执行图优化最经典的例子。PyTorch 实现里可以写成self.qkv nn.Linear(d_model, 3 * d_model, biasFalse)然后把输出的 chunk 分成 Q、K、V。效果不只是在代码层面更简洁在运行时是一次更大的矩阵乘能更好发挥 Tensor Core 的并行能力。第二个策略FlashAttention。它针对的核心问题是 Attention 算子链中 S 矩阵的反复读写。FlashAttention 做了两件大事按块计算让 S 的中间分块驻留在 SRAM 中引入在线 Softmaxonline softmax通过维护行最大值、行指数和来避免对整个 S 做全局归约。这样 Attention 的前向过程中原始 S 矩阵绝大部分时间不用落回 HBMsoftmax 的两次扫描也只剩一次。实测中当 seq_len 达到 2048 以上FlashAttention 相对原始 PyTorch 实现的加速可以达到 2~3 倍甚至更多同时显存占用大幅下降。第三个策略LayerNorm 残差的融合以及 LayerNorm 和它后面的 GEMM 融合。在x x Attn(LN(x))的结构中LN 的输出同时被残差加法和后续 GEMM 使用如果把 LN 的计算结果直接通过 GEMM 的 epilogue 输出到残差缓冲途中至少省掉一次中间张量的显存写读。对于训练场景算子融合需要谨慎因为反向传播时需要重新利用部分中间结果但对于推理场景这种融合几乎只有好处。3.3 从执行图视角估算一次前向传播的时间这一节我们亲自动手动笔感受一下“从第一性原理估算执行时间”的爽快感。假设模型配置为d_model 4096ff_dim 16384num_heads 32d_k 128seq_len 2048batch 4num_layers 24fp16 推理。先估算 FLOPs。一点共识一次矩阵乘法的 FLOPs 约为2 * M * N * K。各主要部分每层 QKV 线性输入尺寸[8192, 4096]权重[4096, 12288]FLOPs 2 * 8192 * 4096 * 12288 ≈ 0.82 TFLOPs。Attention score 计算Q 尺寸[8192, 32, 128]K 转置后与之相乘等效 FLOPs 2 * 8192 * 8192 * 128 ≈ 17.2 GFLOPs。这里和 QKV 相比小了一个数量级。Attention output 投影O[8192, 4096]权重[4096, 4096]FLOPs 2 * 8192 * 4096 * 4096 ≈ 0.27 TFLOPs。FFN 两个线性层第一个2 * 8192 * 4096 * 16384 ≈ 1.1 TFLOPs第二个2 * 8192 * 16384 * 4096 ≈ 1.1 TFLOPs。每层总计约 3.3 TFLOPs24 层约 80 TFLOPs。在一张 A100 80Gfp16 理论上限约 312 TFLOPs实际利用率取 45%的算力水平下理论最优耗时约为80 / (312 * 0.45) ≈ 0.57 s。但实际跑一次前向你用 PyTorch eager 模式大约要 1.5~2 秒。多出来的时间去哪里了答案就在执行图里。这 24 层里还包含 48 个 LayerNorm、48 个残差 add、24 个 Softmax、大量 dropout/bias 相加。它们加起来本身计算量不到 1 TFLOPs但访存量却可能超过 200 GB。以 A100 的 HBM 带宽约 2 TB/s 来算这些算子访存时间就要约 0.1 s。更要命的是PyTorch eager 模式在 GPU 上每启动一个算子 kernel 有几十微秒的开销当整个网络由上千个算子构成时这个调度开销会变成 0.3~0.5 秒的显性成本。因此真实耗时与理论峰值之间的差距大多由低算术强度算子和 kernel 调度开销造成。这个估算方法你自己可以套到任何模型上。4. 踩坑实录从算子粒度出发的调试心得4.1 显存空洞与峰值计算为什么同样的模型别人显存少一半有一次我在调一个视频理解模型序列长度拉长后一个 batch 显存直接爆掉。把执行图输出 dump 下来发现问题不在权重而在大量中间激活值的一个峰值Attention score 矩阵在 seq_len 很长时按平方增长。结论很朴素——用 sublinear attention 或 FlashAttention 能直接把这个峰值压掉。但这里我真正想强调的是显存空洞问题频繁分配、释放的中间张量在显存中会产生碎片峰值时比实际需求高不少。处理技巧有三个。第一在 PyTorch 里尽量复用已有缓冲区避免反复new_empty或torch.empty。第二跑推理时固定形状可用torch.jit或 TensorRT 图优化把中间张量的生命周期尽量缩短。第三注意权重和激活分开管理显存训练时用梯度检查点技术激活重计算以减少存储激活的空间这一招能在几乎不损失性能的情况下让 batch 增大一两倍。执行图优化的目的不是让单个算子变快而是让所有张量的生命周期尽量短显存尽量复用。4.2 数值稳定性Softmax 的溢出、LayerNorm 的 eps、fp16 的表示范围写算子不是数学题你还得处理浮点数本身的脾气。Softmax 的经典数值稳定方法是在每行减去最大值即考察exp(x - max)。这个技巧在算子实现中不可省略。如果是 fp16 训练或推理exp对绝对值很大的负数会直接下溢为 0一整个行全变成 0 的话梯度回传也全部消失。所以实现 Softmax kernel 时我的建议是不管 FP16、BF16、FP32统一先减最大值再算哪怕因此多做一次读取也比溢出导致的不稳定要划算得多。LayerNorm 里的eps参数也别乱动。它出现在分母的方差中(x - mean) / sqrt(var eps)。如果 eps 设得太小fp16 下方差较小的时候除出来异常大或出现除零误报设大了归一化精度会受影响。常见做法是eps1e-5到1e-6之间但在 fp16 下试过不少模型我建议用1e-5起步不要因为 QuickStart 代码里是 1e-12 就直接抄那是给 fp32 用的。4.3 用等价算子替换法定位性能瓶颈最后分享一个非常实用的调试技巧等价算子替换法。当我觉得某个算子开销不合理的时候先把模型里的算子用手写等价版本替换再重新跑 profile对比前后时间变化。例如把 GELU 换成恒等映射x - x看整体耗时降多少就能反推出 GELU 及其相关访存的开销把 Attention 里的 Softmax 换成恒等映射同理。这是做算子级性能归因最简单直接的方法比分析理论计算量来得更快。做算子替换的原则是“一次只动一个变量”保证其它层都保持不变否则无法判断瓶颈。实际使用中我发现 PyTorch Profiler 的图表视图已经很直观配合torch.profiler的 kernel 级输出可以很清楚地看到每个算子占用的总时间、GPU 时间、调用次数。如果某个算子调用次数不多但 GPU 时间占比很高优先优化它如果调用次数非常多但每个都很短优先想融合。方向对了性能优化的大头就抓到了。一点个人心得算子级性能分析切忌只盯着 FLOPs。真正的瓶颈通常藏在访存开销和 kernel 启动开销里也就是“低算术强度的算子”和“碎小算子”这两个地方。Transformer 的结构决定了这两类问题必然存在——它的成功来自高度的结构表达力而它的性能压力也恰恰来自结构带来的算子多寡与访存开销。理解这一层才会理解“算子融合”、“内存复用”、“图优化”为什么是部署必选项而不是锦上添花。最后再分享一个落地建议新项目里别等模型写好后再做性能优化。在模型结构定稿前先把执行图草稿画出来——哪些算子可以合并哪些中间张量可以省掉哪一步会形成显存峰值。你在这个阶段能解决的问题比事后跑 profiler 再去修要省一半力气。Transformer 十年前第一次出现时我们还在为精度欢呼现在它遍地开花性能优化的“第一性原理”也该回到算子执行图这个最本质的层面了。
