📚 课程目录第3章推理性能基准测试与剖析工具实战

3. 推理性能基准测试与剖析工具实战

物理代理模型（如神经算子、图神经网络代理、PINN 推理器）在训练完成后，真正决定其能否落地的往往是推理阶段的性能。与常规视觉模型不同，物理代理模型的输入常包含高维网格、边界条件张量、多物理场通道，单次推理可能同时触发大量小算子、非规则访存和跨设备数据搬运。因此，基准测试不能只测一个"端到端延迟"，而要能回答三个问题：时间花在哪些算子上、内存访问是否成为瓶颈、并行度是否被充分利用。本节围绕这三个问题搭建可复现的基准环境，并用 Nsight Systems 与 PyTorch Profiler 采集对应数据。

一、基准测试环境的三层结构

一个可靠的推理基准由三层构成：负载层负责生成与真实工况一致的输入；执行层固定模型、精度、批大小与设备；采集层决定用哪种工具、采哪些指标。三者必须解耦，否则一次改动就无法判断性能变化来自哪里。

负载层的关键是"形状真实"。以二维稳态热传导代理模型为例，训练时常用 128×128 网格，但工程推理可能面对 256×256 或 512×512。建议准备一组形状阶梯，例如 (1,3,128,128)、(1,3,256,256)、(1,3,512,512)，并额外准备一个变批量序列 [1,2,4,8]，用于观察吞吐随批量的变化。输入张量用固定随机种子生成并缓存到显存，避免每轮都做主机到设备拷贝，把 H2D 传输时间混入算子耗时。

执行层要显式固定以下变量：torch.inference_mode() 关闭自动求导；model.eval() 关闭 Dropout 与 BatchNorm 的训练行为；torch.backends.cudnn.benchmark 的取值要记录（开启会让首次运行做算法搜索，导致首轮异常慢）；精度统一为 FP32 或 AMP 的 FP16，不可混用。预热轮数建议不少于 10 轮，正式计时不少于 50 轮，取中位数与 P95 两个统计量，因为物理代理模型常因显存分配抖动出现长尾。

二、用 PyTorch Profiler 定位算子与访存

PyTorch Profiler 适合回答"哪些算子最贵、访存效率如何"。核心是开启 record_shapes、profile_memory 与 with_stack，并把 schedule 设为"跳过预热、只采若干步"，避免把预热阶段的大量一次性分配计入统计。

import torch
from torch.profiler import profile, ProfilerActivity, schedule

model = model.eval().cuda()
x = cached_input.cuda()

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=2, warmup=5, active=10, repeat=1),
    record_shapes=True,
    profile_memory=True,
    with_stack=False,
) as prof:
    with torch.inference_mode():
        for _ in range(20):
            model(x)
            torch.cuda.synchronize()

print(prof.key_averages().table(
    sort_by="cuda_time_total", row_limit=20))

读表时有三个判断要点。第一，看 cuda_time_total 的头部算子是否与模型结构吻合：若卷积类算子占比高，说明计算密集；若 elementwise、copy、cat 类算子排名靠前，往往意味着存在大量碎片化小算子，是融合优化的目标。第二，看 Self CUDA Mem 与 CUDA Mem 的差值，若某算子自身计算时间短但内存列很高，说明它是访存受限而非算力受限。第三，看 Input Shapes 列，确认实际执行形状与预期一致——物理代理模型常因动态 shape 触发重复编译或回退到慢路径。

一个典型现象是：在 512×512 输入下，某个归一化算子耗时从 0.3 ms 跳到 4 ms，而 FLOPs 只增加 4 倍。这通常不是算力问题，而是该算子在大张量上退化为多次小 kernel 启动，或触发了非合并访存。此时应结合 profile_memory 输出的分配峰值，判断是否需要改用融合实现或调整内存布局（如 NHWC）。

三、用 Nsight Systems 看时间线与并行度

PyTorch Profiler 给的是算子级聚合，Nsight Systems 给的是时间线：CPU 线程、CUDA 流、kernel 起止、H2D/D2H 拷贝、NCCL 通信在一条时间轴上如何交错。对物理代理模型而言，最值得看的是"GPU 是否有空洞"。

nsys profile \
  --trace=cuda,nvtx,osrt,cudnn,cublas \
  --cuda-memory-usage=true \
  --output=proxy_infer_profile \
  python bench_infer.py --shape 512 --batch 1

采集后在时间线上重点观察三处。其一，kernel 之间是否存在明显间隙：若每个小 kernel 之间都有几微秒空隙，说明 CPU 提交速度跟不上 GPU 消费速度，属于启动受限，应减少 Python 侧循环、使用 CUDA Graph 或算子融合。其二，是否存在同步点：torch.cuda.synchronize()、.item()、.cpu() 会在时间线上形成 CPU 等待，把本可重叠的流水打断。其三，多流并行是否真的生效：若模型被拆到两条流上，但时间线上两条流的 kernel 仍串行，说明存在隐式依赖或流同步未正确设置。

建议在代码中用 NVTX 打点，把"数据准备""前向""后处理"三段标出来，这样时间线上能直接看到各阶段占比，而不必靠猜。例如在推理前后插入 torch.cuda.nvtx.range_push("forward") 与 range_pop()，再配合 --trace=nvtx 采集。

四、把两类数据合成一张判断表

单看任一工具都容易误判。正确做法是把 Profiler 的算子聚合与 Nsight 的时间线对齐：Profiler 告诉你"哪个算子贵"，Nsight 告诉你"贵的时候 GPU 在不在干活"。据此可形成如下判读逻辑（示例阈值仅作教学参考，需按实际硬件标定）：

现象	Profiler 侧证据	Nsight 侧证据	判断与动作
端到端延迟高但 GPU 利用率低	头部算子耗时占比分散，elementwise 多	kernel 间频繁空洞，CPU 提交密集	启动受限：算子融合、CUDA Graph、减少 Python 循环
某算子耗时随形状超线性增长	该算子 CUDA Mem 列高，Self CUDA 低	该 kernel 占用时间长且无重叠	访存受限：改内存布局、合并访存、分块
批量增大吞吐不升反降	显存分配峰值接近上限	出现显存回收或拷贝等待	显存压力：降精度、梯度检查点不适用推理，改分块推理
多流并行无收益	各流算子耗时正常	两条流 kernel 仍串行	隐式依赖：检查流同步与事件等待

五、实操中的常见坑

第一，忘记 torch.cuda.synchronize() 就计时，测到的是 CPU 提交时间而非 GPU 完成时间，结果偏小且不稳定。第二，把首次运行计入统计，cuDNN 算法搜索与显存池首次分配会让首轮慢数倍，必须预热。第三，Profiler 本身有开销，active 步数不宜过多，且不要在同一进程里既做长稳测试又开 Profiler。第四，Nsight 采集时若 --cuda-memory-usage=true 会显著增加开销，只在对齐显存问题时开启。第五，物理代理模型的输入常带边界掩码，掩码若在每轮重新生成，会把无关计算混入前向耗时，应预生成并缓存。

建议把基准脚本固化为可传参的命令行工具：--shape、--batch、--precision、--warmup、--iters、--profile。每次优化前后用同一组参数跑两遍，把中位数延迟、P95、峰值显存、头部算子占比四项记录到同一张表里，才能形成可比较的优化基线。

完成本节后，你应当能独立搭起一个形状可控、变量固定的推理基准，并用两类工具分别回答"算子贵在哪""访存是否受限""并行度是否用满"。这三类证据是后续算子融合、显存优化与并行调度的直接输入，缺少任何一类，优化都容易变成盲调。