PyTorch逐算子性能基准测试：从profiler到独立benchmark的实践

做性能分析这几年我越来越觉得模型层面的benchmark是“看起来会实际很难做透”的一件事。用PyTorch profiler拉一轮forward看到的是每个op在总耗时里的占比但真要回答“这个op单独跑到底有多快”、“换成融合版能提升多少”、“同一块卡上不同shape下谁在退化”profiler给不了稳定可信的答案。于是我自己写了kern perf这个工具专门做全自动的模型forward逐op性能基准测试。这篇文章就是把我从思路到落地、从踩坑到修坑的完整过程整理出来给同样需要做算子级性能分析的朋友一份可以直接参考的实践记录。kern perf解决的核心问题很简单给定一个任意PyTorch模型自动把forward路径上每一个算子op捞出来单独构造输入、单独跑benchmark输出每个op的耗时分布、总耗时占比、带宽或FLOPs估算值最后汇总成一个可排序、可对比、可回归的报表。它不需要你给模型代码打任何埋点不需要手写每个op的bench脚本也不需要维护一个operator名单。你只需传一个模型实例和一份示例输入剩下的遍历、统计、报告都由工具完成。1. 为什么非要“每个op单独bench”不可1.1 模型慢往往不是大op慢而是小op多很多人做模型优化第一步就是看PyTorch profiler的表格按self CUDA time排序然后去优化排在前面的那几个大op。这个思路在大模型场景下有一定的合理性但如果你的模型是那种op数量多、单个op都很小的结构比如各种多分支attention变体、混合专家路由、多模态fusion模块profiler表上写“0.5%”的小op往往才是真凶。我用一个实际案例说明。之前调试一个多模态融合模型forward里有40多个不同op。profiler显示单次自注意力大算子稳定占30%以上按理说先优化它就行了。但当我们把大算子换成flash attention之后端到端延迟只下降了6%远低于预期的15%。后来逐个op去测发现模型在一层里调用了7次repeat、9次view、6次permute这些“小op”因为kernel launch次数多、访存模式差叠加起来占了接近25%的耗时。只看profiler总表根本无法把这个结论量化出来。所以核心痛点很清晰你需要知道的是“如果这个op被换掉、融合掉、或者改成另一种实现对整体forward的影响精确是多少”而不是“它在整体耗时中排第几”。这就需要把op之间的调度开销、显存分配噪声、L2 cache干扰全部剥离开单独量每一个op的性能特征。1.2 profiler为什么不能直接拿来当benchmark用PyTorch自带的profilertorch.profiler.profile是采集器不是benchmark工具。它做的是“记录一次或者多次forward里每个op的耗时”这个耗时是从整个执行流里截断出来的混入了很多干扰kernel launch的排队延迟CUDA上的op执行是异步的profiler记录到的CPU耗时里包含launch overheadGPU耗时里包含stream排队等待。如果你不是单独测这个op它的“耗时”其实取决于它在forward里的位置前面的op跑得慢它看起来也会慢一点。显存分配的抖动第一次执行某个op时torch的caching allocator可能需要切割新的block这个分配cost会被算进op耗时里。第二次、第三次执行同一op显存已经缓存好了耗时又会掉下来。同一个op在同一轮profiling里不同位置的耗时可能差好几倍。L2 cache和页表热度的偏差一个op在整体forward里跑的时候它的输入数据往往是刚从上一个op写出来的还在L2里能享受到比较高的cache命中率。如果你想知道它在“真normal”场景下的耗时这个cache状态反而是不真实的。kern perf的做法完全不同从模型里把每个op提取出来后每个op在受控条件下“单独”跑输入是预先构造好、固定在显存里的warmup足够轮次再连续跑上百次取样。这样才能拿到相对稳定、可比较、可复现的单op性能数据。1.3 什么场景下你会需要这个工具这个工具不是给你“看看”用的是给你“用”的。举几个我实际用到的场景模型上线前体检新模型从训练/微调环境搬到推理环境kernel库版本、CUDA版本都变了用工具跑一遍forward的op级报告看看有没有哪个op在目标环境上突然退化。算子替代验证你想把某个自定义op换成torch原生op、或者换成融合实现最直接的办法是用kern perf分别测两个版本的单op耗时再做端到端对比。硬件差异对比同型号模型在不同GPU比如A100、L40S、4090上的逐op耗时差异能帮你判断瓶颈是访存还是算力。精度优化回归你做量化、剪枝、算子融合之后用工具跑一遍op级基准确认没有哪个op因为shape变化或dtype变化产生性能回退。2. kern perf的架构设计与选型思路2.1 总的流程分成四步kern perf的全流程可以概括为拿到模型和示例输入 - 静态或动态提取forward中的算子序列 - 为每个算子抽象出独立的最小执行单元 - 逐个执行benchmark并汇总报告。我最早想用最简单的方案hook住nn.Module的forward或者用torch.jit.trace记录一份执行序列。但很快发现问题很大。nn.Module级hook只能拿到模块粒度一个nn.Linear内部其实还包含MatMul、Add、可能还有Reshape等好几个aten op而trace方式虽然能拿到aten op列表却会丢失很多动态信息的上下文比如某些op的输入是python原生list还是tensor list、某个维度是否需要保持整型、哪些tensor在trace之后不可复现。所以最终采用了更稳定的组合策略优先用torch.fx的symbolic_trace来抓取计算图拿到graph之后遍历node、按node粒度提取op针对fx无法处理的模型再回退到基于torch.compile的graph break信息或者手动注入hook来兜底。整套流程跑下来99%的普通模型都能自动处理。flow: 原始模型 示例输入 - torch.fx.symbolic_trace / 兜底hook - 得到有序的op节点列表 - 对每个节点记录op类型、参数、输入key、shape/dtype/stride信息 - 从缓存中分配独立的tensor - 每个op执行warmup N轮有效测试 - 统计mean/p50/p90/min/max估算FLOPs/带宽 - 输出文本报告 CSV JSON2.2 算子级别的粒度选择为什么是“node”而不是“module”实际跑起来的经验是按照fx图的node粒度来做bench信息量最大。一个node对应一个具体的计算可能是一个aten op也可能是一个调用内置函数的call_function。这个粒度恰好比“一行python代码”更细、比“单个kernel”更宏观。你不能直接用torch.profiler的kernel级数据因为一个op可能发射多个kernel而多次调用同名op也会有输入shape的差异。node粒度的好处是它和计算图一一对应你能明确知道这个op接在哪个前序op之后、输出哪个tensor、给哪个后续op用。这样在抽取bench单元时我们不需要重新推理数据流只需要按图拓扑顺序逐节点建立“输入tensor规范”和“输出tensor占位节点”就能保证整个forward的数据流是一致的、可复现的。2.3 有向无环图遍历与独立bench单元构建拿到fx.Graph之后遍历方式上有一点值得强调不要用递归用拓扑排序。因为fx.Graph本身自带节点顺序从placeholder到output你直接按graph.nodes的顺序遍历就行但要注意处理get_attr、placeholder这类不产生实际计算的node。这些node虽然不进bench单元但它们是op输入参数的来源。具体处理逻辑我放在下面的代码注释里def extract_op_bench_units(model, example_inputs): import torch.fx as fx gm fx.symbolic_trace(model) if not isinstance(model, fx.GraphModule) else model graph gm.graph # 先建立placeholder名到tensor规范的映射这部分节点对应模型输入 input_specs {} placeholders [] for node in graph.nodes: if node.op placeholder: # 这里要精确记录名字后面构造输入时一一映射 placeholders.append(node) input_specs[node.name] example_inputs[len(placeholders) - 1] elif node.op get_attr: # get_attr节点拿到的是模型参数/缓冲区先缓存实际tensor input_specs[node.name] gm.get_parameter(node.target) \ if hasattr(gm, get_parameter) else None elif node.op call_function: input_proto { node: node, target: node.target, args: [], kwargs: {}, inputs: [], # 存放实际tensor输入 output_dtype: None, } # 对每个args/kwargs如果是node引用则从input_specs里取tensor # 如果是常量就原样记录 for arg in node.args: if isinstance(arg, fx.Node): input_proto[args].append(input_specs[arg.name]) else: input_proto[args].append(arg) for key, arg in node.kwargs.items(): if isinstance(arg, fx.Node): input_proto[kwargs][key] input_specs[arg.name] else: input_proto[kwargs][key] arg # 注意这里只是一个记录真正的tensor要等bench时才从池里取 bench_units.append(BenchUnit(node.name, node.target, input_proto)) return bench_units这段代码看起来简单但在真实工程里要加很多防御逻辑。比如node.args里可能出现slice(None, None, None)这种常量对象可能kwargs里带设备字符串、dtype对象还可能出现fx的Proxy对象被包装进list/tuple的情况。我统一用了一个递归的normalize_arg函数遇到list/tuple/dict就递归展开遇到fx.Node就替换成实际的tensor规范遇到基本类型就原样保留。2.4 关于“模型融合”和torch.compile带来的新问题现在做算子bench绕不开torch.compile和模型融合的话题。很多人给我反馈说用torch.compile跑过的模型再去抽取op节点得到的图和原始模型对不上因为Inductor已经把多个op融合成一个大的fusion kernel了。kern perf对这种情况的处理策略是如果你要bench的是“原始模型的计算图”那就用torch.compile(model, disableTrue)或者直接用原始模型实例抽取节点如果你要bench的是“融合后模型”就需要在中后端接口上做hook按fusion的kernel粒度来分割。目前我的实现里默认抽取原始fx图因为从工程上看用户最容易理解、最容易做基准的永远是原始计算图。在报告里会额外标注“原始op”和“融合kernel”的映射关系方便你在调优时对照。3. 核心细节解析与实操要点3.1 tensor构造如何让每个op“单独”跑起来每个bench单元需要独立运行就必须有独立的输入。最简单粗暴的方法是直接复制原forward前向传播创建出来的tensor作为输入。但在自动抽取场景下很多op的输入是从前一个op的动态输出里拿到的单测时不可能真的串起全图再跑到这个op那就失去单独bench的意义了。所以kern perf的做法是按节点顺序静态解析输入规范。每个op的输入tensor记录它的shape、dtype、device、requires_grad、stride在需要时以及它作为“常量”还是“activations”。bench时根据这些规范重新构造一个相同shape、dtype、device的随机tensor。对于同一份tensor被多个op共享的情况会额外建立引用关系确保同一输入的多个op在bench时使用完全相同的tensor实例。构造tensor的时候有个坑有些op对输入tensor的实际数值范围敏感比如Softmax对极端大数、LayerNorm对极小方差。如果全部用标准正态分布随机tensor某些op可能落入数值不稳定区域导致bench时出现NaN或inf从而读取错误的时间。我对这类op做了白名单处理softmax/layernorm/transpose/sum这类op在构造输入数值时固定一个稳定区间比如U(-1,1)避免数值爆炸。3.2 warmup和重复策略统计口径怎么定才可信一个常见的问题是benchmark到底跑多少次数据才算数。我踩过坑之后现在固定采用“warmup 20次 正式跑100次”的默认配置同时支持命令行参数覆盖。为什么warmup这么重要因为CUDA kernel的第一次调用会有cuModuleLoad、cuFuncSetAttribute、context初始化等开销这些跟op本身性能毫无关系会严重拉高第一轮的耗时。20次warmup可以有效把这部分开销滤除。如果是CPU推理场景建议warmup次数至少加到50次因为CPU的JIT、MKL-DNN / oneDNN的primitive缓存是在第一次真正运行时才建立的。正式跑100次之后我取的不是均值而是P50和P90。这里的原因是GPU上的op耗时分不是正态分布均值会被极端值拉高。P50中位数代表最常见的情况P90反映长尾。最终的kern perf报告里同时输出mean / p50 / p90 / min / max并标出p90/p50的比值如果这个比值大于2说明op存在明显抖动需要重点排查。3.3 非tensor参数怎么处理这是很多自写bench工具漏掉的地方。一个op节点的输入不只是tensor还有大量标量、枚举、字符串、形状。举例torch.sum(input, dim1, keepdimTrue)dim是intkeepdim是bool。torch.narrow(input, dim, start, length)start/length是int。torch.permute(input, dims_list)dims_list是list。torch.cat(tensors, dim)tensors是tensor listdim是int。torch.nn.functional.interpolate(input, sizeNone, scale_factorNone, modenearest)mode是字符串。kern perf在抽取阶段把这些参数连同实际值一起记录下来保存成prototypebench阶段按prototype重组参数确保op调用的语义完全一致。这里面最麻烦的是list/tuple of tensors。比如torch.cat或者torch.stack你需要确保传入的一个tensor list中的所有tensorshape、device、dtype都按原计算图中的信息构造。抽取时就递归遍历list识别到list内元素是fx.Node时继续解析为tensor规范。这个逻辑我在前文代码注释里写过实际工程里是必须的。3.4 指标估算FLOPs和带宽不能漏单纯给一个op的耗时没有意义因为不同op的计算密集型程度完全不同。一个0.1ms的矩阵乘和一个0.1ms的transpose优化价值完全不同。所以kern perf报吿里加入了两个粗粒度指标FLOPs浮点运算量和HBM带宽显存带宽。FLOPs的估算对matmul、conv、linear这几种算子比较精确可以直接由输入shape计算对elementwise、transpose、view这类算子FLOPs意义不大反而需要看带宽。我实现了两个小函数做粗算def estimate_flops(node_target, args, kwargs): # 这里只做粗粒度估算 if node_target in (torch.mm, torch.bmm, matmul): # matmul: (m,k) (k,n) - 2*m*k*n ... elif node_target in (torch.nn.functional.linear, linear): # linear: input (b, in) weight (in, out) - 2*b*in*out ... elif node_target in (torch.nn.functional.conv2d, conv2d): # conv2d: b, c_in, h, w * c_out, c_in, kh, kw - 2*b*c_out*c_in*kh*kw*h_out*w_out ... return flops def estimate_bytes(node_target, args, kwargs): # 字节数直接算所有输入tensor和输出tensor的numel * dtype_size total 0 def add_tensor(t): nonlocal total total t.numel() * t.element_size() # 遍历args里的tensor加上对输出shape的估算 return total带宽的估算不能只算输入还要算输出。但输出tensor在bench前还不知道具体shape好在fx图里能通过meta信息拿到新版本fx会做shape propagation或者也可以先实际跑一遍op、拿到真实输出shape再回填。后者最简单我在工具里就是用这种“先试跑一次拿到真实输出shape”的方式因为每种op跑一次的成本很低。3.5 运行时环境控制锁频、清L2、隔离分配器噪声单op bench要想数据稳定有四个环节必须控制住锁GPU频率新一点的NVIDIA驱动里可以用nvidia-smi -lgc 1500把基准时钟锁在固定频率。如果不锁频GPU idle时降频、跑起来后又boost出来的数据前后能差20%。bench完记得nvidia-smi -rgc恢复。清L2 cache默认情况下每个op bench都是连续跑前一次的输出数据还留在L2里后面几次跑就能吃L2红利。想测真实访存性能应该每次跑之前刷一遍L2。kern perf用了一个小trick在每次op执行前对一个独立的、大小等于L2 cache的buffer做一次全量写操作把L2污染掉。这个开关默认关闭开启后测量的才是“冷缓存”数据。固定streamtorch默认用current stream但如果你程序里之前留下过异步操作可能计时不准。kern perf里每次bench前先torch.cuda.synchronize()用单独的stream执行op结束后再synchronize。显存分配去抖动反复创建tensor会让caching allocator频繁分配/释放属于严重噪声。做法是在bench前预分配一批相同shape的“假tensor”池每次op执行前从池里取tensor拷贝输入、执行、释放回池里。这样CUDA层面的所有内存块都在一开始就分配好了。4. 实操过程与核心环节实现4.1 命令行入口与参数kern perf用起来非常简单。最基本的用法是python -m kern_perf --model resnet50 --example-inputs ./assets/resnet_input.pt --warmup 20 --repeat 100 --output-dir ./reports--model传入模型名字工具内部会根据一个简单的注册表去构建模型。--example-inputs示例输入tensor的pt文件如果缺失则自动生成随机输入。--warmup/--repeatwarmup和正式跑次数。--lock-gpu-clock是否锁GPU频率。--output-dir结果输出目录会同时写出文本报告、CSV和JSON。自定义模型也支持直接传脚本文件python -m kern_perf --custom-model-file ./my_model.py --model-class MyNet --checkpoint ./weights.pt --example-inputs ./input.pt4.2 我实际跑一个ResNet-50样例的结果解读为了验证工具我拿ResNet-50torchvision标准版跑了一轮输入是(1, 3, 224, 224)GPU是L40Swarmup 20次、repeat 100次。抽取到的op节点共有163个实际进入bench池的是152个其余是get_attr和placeholder。报告节选op_nameshapemean(us)p50(us)p90(us)FLOPs(M)bytes(KB)占比conv2d_0(1,64,112,112)132.5130.2147.8118.3401.24.1%conv2d_1(1,64,112,112)332.7331.5367.9591.61109.810.3%batchnorm2d_0(1,64,112,112)28.428.230.10.9802.10.9%relu_0(1,64,112,112)6.26.17.80.9803.20.2%maxpool_0(1,64,112,112)72.571.984.230.1805.32.2%add_0(1,256,56,56)17.817.620.01.61608.40.5%avgpool_0(1,2048,7,7)42.141.649.7—802.11.3%linear_0(1,1000)14.214.117.01.42.40.4%数据透露出几个有意思的结论conv2d_1占比最高符合预期矩阵乘/卷积是这个模型的绝对主力。batchnorm和relu虽然单独看不高但整个模型里有几十次调用累计占比其实可观——这就是前面说的“小op多积少成多”问题。maxpool读4MB数据花了70多us明显是访存带宽受限不是算力受限。4.3 自定义模型的bench与算子融合对比再说一个自定义模型的案例。当时手头有个模型forward里有这么一段# 原始实现 h x.permute(0, 2, 1, 3) # (b, c, t, d) - (b, t, c, d) h h.reshape(b, t, -1) attn torch.matmul(h, h.transpose(-2, -1))直观上觉得这段permutereshapematmul可以融合成一次matmul但不确定收益有多大。用kern perf分别bench permute、reshape、matmul三个op结果permutereshape二者单op耗时合计约28usmatmul约12us。然后换成融合版实现后端到端只少了大约20us。这让我意识到真正的瓶颈反而不是这三个op反而是模型里一个看起来占比不高、每次都传入固定mask的torch.masked_fill。这说明单op bench的价值不只是看谁大而是能把问题拆到“到底这个op值不值得优化”的程度。4.4 动态shape模型的特殊处理很多新模型forward里有动态shape例如按输入长度生成attention mask、再做masked_fill。fx.symbolic_trace对动态shape支持差可能直接报错。我的处理方式是两级兜底第一级先尝试fx.symbolic_trace如果失败就建议用户加一行torch._dynamo.config.suppress_errors True改用torch.compile的dynamo来拿graph第二级还不行就退化到“手动标注”模式用一个--op-list参数让用户把要bench的op名字和输入规范用简单DSL写出来工具只负责bench。最常用的动态shape场景是NLP模型里输入长度可变。这时候bench的op序列其实一样只是shape不同。kern perf支持把示例输入扩展为多组shape--example-inputs {input_ids: torch.randint(0, 1000, (1, 128))} {input_ids: torch.randint(0, 1000, (1, 512))}。工具对每组shape分别做一次完整的抽取bench报告里会自动做group-by-shape对比能直观看到每个op的耗时随序列长度的扩展规律——这对定位LLM推理中的长序列退化特别有用。5. 常见问题与排查技巧实录5.1 CUDA caching allocator干扰读到“假优势”现象单独bench某个op跑50次前半段耗时高、后半段耗时低p50比mean低很多。原因torch的caching allocator在前几次分配内存时需要向CUDA申请新的block慢后面内存复用分配就被cache住了快。这会导致bench结果低估实际耗时。对策所有tensor预分配进内存池bench时只做“拷贝输入 执行op 拷贝输出回去”不新分配显存。kern perf实现了一个TensorPool类维护一个dict[(shape,dtype,device)] - list[tensor]需要时pop用完放回。5.2 使用fx symbolic_trace时遇到“Proxy tensor”不可迭代现象symbolic_trace跑到某个含python循环的模型报错Iterating over a tensor is not allowed。原因模型里有动态shape的循环fx无法展开循环体从而对tensor调用python迭代。对策如果是固定循环次数建议把循环写成torch.unbind或者用torch.arangegather等向量化写法如果实在要保留循环fx兜底失败后就退回到手动op-list的bench模式。这类问题本质上不是benchmark工具的锅是fx的图捕获能力边界。5.3 同一个op在模型里出现多次bench池怎么区分现象模型里conv2d被调用了10次每层shape不同但fx图里这些node名字不一样conv2d_0、conv2d_1……我们bench时需要每个实例都跑而不是合并成一个。对策kern perf按node.name区分实例不按op type聚合。同时报告里会额外加一列“op_type”方便按类型排序。这个设计对定位“同一类型在不同层的性能退化”非常有用。5.4 自定义op不支持fx现象模型里用了torch.autograd.Function自定义opfx.symbolic_trace可以追踪到但bench时如果op里有非tensor的静态变量比如提前编译好的CUDA kernel句柄直接传给bench单元会报错。对策kern perf在抽取到未知target时会优先检查该target是否实现了__kern_perf_bench__协议如果实现了就调用它的自定义bench逻辑没实现就把这个op标记为“需要手动bench”在工作里列出来提示用户。我不认为工具可以99%自动化解决所有自定义op的规格差异明确标记待办事项比强行跑出来的错误报告更有意义。5.5 多GPU环境的stream隔离现象多卡环境下bench默认在torch.device(cuda)上跑可能落在GPU0但模型参数在GPU1上也跑出结果但数据完全不可信。对策kern perf增加--device cuda:1选项bench单元构造tensor时全部显式指定device。同时默认在bench前做一次torch.cuda.set_device防止肉眼忽略卡号。5.6 锁频没生效现象设置了--lock-gpu-clock但报告里的mean耗时还是很跳。原因L40S这种卡如果进了持久化模式但没设--persistence-mode或者驱动版本不够nvidia-smi -lgc不能锁死。对策在你真正跑关键bench前先跑一次nvidia-smi -q -d CLOCK看当前频率是否稳定不稳定就加--persistence-mode启动持久化再lgc。锁频之后跑第一个op前最好先空跑1-2秒让GPU稳定在新频率再开始warmup。5.7 小op完全被overhead淹没怎么办统计发现模型里很多op本身计算量极小比如一个元素级relu、一个标量乘、一次加法单次耗时只有3-5us。这种情况下kernel launch overhead大概1-2us显得占比很大甚至导致耗时波动超过30%。kern perf里专门出了一个“batch execute”模式把同一个op在bench时连续构造多个不同的输入tensor塞进同一个stream里跑然后取“总耗时/实例数”作为单op单次执行耗时。这种模式下launch overhead被摊薄到接近0更能反映op本身的计算能力。当然这个模式测出来的数值会低于常规场景下的真实耗时报告里会特别标注是“无overhead的理论值”建议跟“有overhead的实际值”一起看。6. 结合工具使用后的几点体会做kern perf之前我花了很多时间在profiler输出的表格上反复对比效率很低。真正把每个op隔离出来bench之后很多原本看不清楚的性能问题变得非常直接到底是kernel慢还是调度慢到底是访存受限还是算力受限到底是优化收益大还是纯属白忙活几分钟就能给出数据支撑的结论。这个工具目前的适用范围是PyTorch原生模型和常见的torchvision、Hugging Face模型。对于含大量自定义C op、或者依赖外部推理引擎TensorRT、ONNX Runtime的模型需要额外扩展特定的bench适配层这也是后续可以继续做的方向之一。如果你平时也要做算子优化、模型压缩或者硬件的性能对比我建议不要只停留在用profiler拉一次总表而是把这个逐op bench的思路带进你的分析流程。先从自己的模型跑一遍看看你会发现有些结论跟你之前的直觉完全不一样。