PyTorch GenAI 内核基准测试实战指南：eager、torch.compile、Quack 与 Liger 四种实现横向评测

GenAI（生成式 AI）模型推理与训练中，CrossEntropy、Softmax、RMSNorm、LayerNorm 是出现频率最高、对整体吞吐影响最大的几个算子。PyTorch 仓库的 benchmarks/dynamo/genai_layers 目录提供了一个开箱即用的内核（kernel）级基准测试套件，可以在同形状、同输入下同时跑通 eager、torch.compile 编译版、Quack 以及 Liger 四种实现，自动完成精度互检、延迟测量、显存带宽换算与几何平均加速比报告。读完本文，你将掌握这套工具的全部 CLI 用法、每种被测内核与后端的实现差异、精度校验机制，以及如何从源码层面理解并复现一次可信的 GenAI 内核评测。

一、基准套件在仓库中的位置与设计目标

该目录完整位于 benchmarks/dynamo/genai_layers，自包含四个文件，职责划分清晰：

文件	作用
benchmark.py	命令行入口，负责参数解析、benchmark 注册表与调度
kernels.py	定义 8 个被测内核类，每个类实现 eager/compiled/quack/liger 各后端
utils.py	基准框架核心：BenchmarkKernel 基类、计时、精度校验、可视化
requirements.txt	被测第三方内核库与绘图依赖

从脚本顶部 docstring 可以确认其目标（见 kernels.py 与 benchmark.py）：为多种 kernel 实现提供统一的基准运行器，覆盖 CrossEntropy、Softmax、RMSNorm、LayerNorm 的前向与反向，从而对比 PyTorch eager、PyTorch 编译器（torch.compile）、Quack、Liger 四类方案的性能。

需要说明的是，Quack 与 Liger 并非 PyTorch 内置模块，而是通过 requirements.txt 引入的第三方 CUDA 内核库（quack-kernels 与 liger-kernel）。因此本套件的定位是"在 PyTorch 仓库内做第三方 fused kernel 的横向对比"，而不是纯 PyTorch 自测。

二、环境准备与依赖安装

目录 README 明确要求假设 PyTorch 已安装，随后补齐依赖：

pip install -r requirements.txt

查看 requirements.txt 可以看到它实际拉取的内容：

quack-kernels
liger-kernel
nvidia-cutlass-dsl==4.1.0.dev0
matplotlib

解读每个依赖的用途（依据源码用法，见下文各节）：

• quack-kernels：提供 quack.cross_entropy、quack.softmax、quack.rmsnorm、quack.layernorm 等融合 CUDA 内核入口；
• liger-kernel：提供 LigerCrossEntropyLoss、LigerSoftmax、LigerRMSNorm、LigerLayerNorm 及底层 layer_norm_backward；
• nvidia-cutlass-dsl==4.1.0.dev0：Quack 内核依赖 CUTLASS DSL，同时代码中直接以 cutlass、cutlass.torch 的 BFloat16 dtype 构造输入张量；
• matplotlib：仅用于 --visualize 绘图。

由于所有被测内核都在 CUDA 上运行（输入统一构造在 device="cuda" 上，且 Quack 使用 cutlass dtype），运行本套件需要具备 CUDA 环境的 GPU 机器。脚本启动时会先打印环境信息，便于确认版本组合：

Environment information:
  Python version: ...
  PyTorch version: ...
  CUDA version: ...

实现见 benchmark.py 的 show_environment_info()。

三、命令行用法：从列清单到跑全套

README 给出了最常用的四条命令，这是全部功能的入口：

python benchmark.py --list                    # List all available benchmarks
python benchmark.py --all                     # Run all benchmarks
python benchmark.py cross_entropy_forward     # Run specific benchmark
python benchmark.py softmax_forward softmax_backward  # Run multiple benchmarks

注册表中的 8 个可用 benchmark 名称（源码见 benchmark.py）：

benchmark 名称	对应内核类	覆盖方向
cross_entropy_forward	CrossEntropyForward	CrossEntropy 前向
cross_entropy_backward	CrossEntropyBackward	CrossEntropy 反向
softmax_forward	SoftmaxForward	Softmax 前向
softmax_backward	SoftmaxBackward	Softmax 反向
rmsnorm_forward	RMSNormForward	RMSNorm 前向
rmsnorm_backward	RMSNormBackward	RMSNorm 反向
layernorm_forward	LayerNormForward	LayerNorm 前向
layernorm_backward	LayerNormBackward	LayerNorm 反向

另外再补充两点源码级细节：

1. 未给任何参数直接运行会报错并打印帮助（见 benchmark.py）；
2. 传入未知名称会提示 Error: Unknown benchmark ... 并建议 --list（见 benchmark.py）。

四、进阶参数：比 README 更完整的 CLI 选项

README 只点到了 --visualize，但实际命令行工具还内置了多个用于精度控制、编译模式切换与自定义编译配置的开关。完整参数解析见 benchmark.py：

参数	类型/取值	默认值	说明
benchmarks	位置参数，可多个	无	要运行的 benchmark 名称，如 cross_entropy_forward softmax_backward
--list	开关	False	仅列出可用 benchmark 后退出
--all	开关	False	依次运行注册表中全部 8 个 benchmark
--visualize	开关	False	每跑完一个 benchmark 生成内存带宽对比图
--compile-mode	default / max-autotune-no-cudagraphs	max-autotune-no-cudagraphs	传给 torch.compile 的优化模式
--tolerance	float	None	精度校验的自定义容差（atol/rtol 同时生效）
--exit-on-accuracy-failure	开关	False	任一后端精度校验失败时立即以错误码退出
--print-benchmark-result	开关	False	打印原始逐点 profiling 结果，便于无 GUI 服务器快速查看
--custom-compile-name	str	None	自定义编译配置在图表上的曲线标签名
--custom-compile-options	str（JSON）	None	自定义 torch._inductor.config 补丁的 JSON 字符串

值得单独讲解的两个能力：

1. 自定义编译配置实验。 传 --custom-compile-options 时脚本会先用 json.loads 解析，并强制要求同时给出 --custom-compile-name，否则抛错（见 benchmark.py）。实际执行时，会先 torch._dynamo.reset() 强制重新编译，再在 torch._inductor.config.patch(...) 上下文内以自定义配置跑 compiled 后端（见 benchmark.py 与 utils.py）。也就是说，你可以在不改任何源码的前提下对比"某条 inductor 配置开关"对性能的影响，适合做编译调优扫描。

2. 无 GUI 环境快速取数。 --print-benchmark-result 会把每个 shape 下各后端的 Performance 对象列表直接打印到终端，而 --visualize 需要 matplotlib 绘图，两者结合可用于服务器上的快速验收。

五、被测内核的实现细节：eager 基准与动态 batch 语义

套件在设计上有几个值得关注的工程决策，全部可以从 kernels.py 源码验证。

5.1 统一的计算/输入约定

• 输入 dtype：主张量统一使用 cutlass_torch.dtype(cutlass.BFloat16)（见 kernels.py），因为 Quack 内核要求 cutlass dtype；
• Norm 类权重 w 为 float32（如 kernels.py）；
• target/标签为 int64，反向的梯度 dloss/dy 与对应 dtype 对齐。

5.2 eager 参考实现

各内核的 eager 后端是精度对照的"金标准"，其定义直接决定了公平性：

• CrossEntropy：F.cross_entropy(x, target, reduction="none")，见 kernels.py；
• Softmax：F.softmax(x, dim=-1)，见 kernels.py；
• RMSNorm：由于 F.normalize 无现成简化调用，套件自己用 PyTorch 算子拼出 float32 中间精度的 rms_norm_ref，公式为 x * rsqrt(mean(x^2, dim=-1) + 1e-6) * w，见 kernels.py；
• LayerNorm：F.layer_norm(x_f32, w.shape, w, None, eps) 后转回原 dtype，见 kernels.py。

5.3 compiled 后端与 batch 维度动态化

编译版在四种内核上普遍采取两个手法（参见 CrossEntropyForward.compiled、SoftmaxForward.compiled 等实现，如 kernels.py）：

1. 用 torch._dynamo.mark_dynamic(x, 0)（以及 CrossEntropy 的 target）把 batch 维标为动态，以模拟真实推理中变化的 batch，避免"形状定死导致编译过度特化"；
2. 通过 lambda 包装后交给 torch.compile(..., mode=self.compile_mode, fullgraph=True)。代码注释特别提醒：必须用 lambda，否则 torch.compile 不会 trace 该函数。

脚本顶部还做了两个全局性 Dynamo 配置（见 benchmark.py）：

torch._dynamo.config.automatic_dynamic_shapes = False
# Needed since changing args to function causes recompiles
torch._dynamo.config.recompile_limit = 1000000

其中关闭自动动态 shape、调高重编译上限，是为了让 sweep 多个 shape 时不因参数变化触发无谓的 guard 失败与重编译，从而保证测得的延迟是稳定的编译后内核延迟。

5.4 反向评测的构造方式

所有 backward 内核都没有直接调 "一个融合反向 kernel"，而是前向 + torch.autograd.grad 的方式构造梯度计算（如 kernels.py 的 eager 反向）。唯一例外是 LayerNormBackward.liger，它直接调用 liger_kernel.ops.layer_norm.layer_norm_backward 并传入 fp32 的 mean/rstd，注释说明了原因：经 LigerLayerNorm + autograd 的反向会把 mean/rstd 保存成低精度，可能过不了精度校验（见 kernels.py）。Quack 的 RMSNormBackward 同样手写了 _rmsnorm_bwd 加 dw_partial.sum(dim=0) 的归约（见 kernels.py）。

5.5 各内核可用后端清单与差异

从各内核类的 available_backends（见 kernels.py 各 __init__）可以整理出以下对照表：

内核	eager	compiled	quack	liger
CrossEntropy Forward/Backward	✅	✅	✅	✅
Softmax Forward/Backward	✅	✅	✅	✅
RMSNorm Forward	✅	✅	✅	✅
RMSNorm Backward	✅	✅	✅	✅
LayerNorm Forward	✅	✅	✅	✅
LayerNorm Backward	✅	✅	❌	✅

需要注意的兼容性细节：

• Quack 的 LayerNorm 不支持 bias（代码注释 quack layernorm does not support bias，见 kernels.py），因此 LayerNormBackward 干脆不注册 quack 后端；
• RMSNorm 的 quack 前向只支持 float32 权重（注释 only supper weight with float32 dtype，见 kernels.py）；
• Liger RMSNorm 反向使用 casting_mode="gemma"（见 kernels.py）。

这些差异说明：横向对比必须逐内核确认后端是否可用，缺失的格子会自动跳过，而不是报错中断。

5.6 被测 shape 矩阵

前向/反向绝大多数内核默认 sweep 同一组 11 个 shape（见 kernels.py）：

(32768, 256)   (32768, 512)   (32768, 1024)  (32768, 2048)
(32768, 4096)  (32768, 8192)  (32768, 16384) (32768, 32768)
(32768, 65536) (16384, 131072) (8192, 262144)

即固定行数 32768（batch×seq_len 语义）并指数级放大最后一维宽度，直逼 262144 的超大词表。Norm 类内核还会叠加 4 个"内部模型更常用"的 shape（见 kernels.py）：(1152*500, 384)、(1152*500, 512)、(1152*1000, 384)、(1152*1000, 512)，即模拟长度 500/1000、隐藏维 384/512 的常见生成式模型配置。个别超大 shape 在 H100 上会 OOM，代码中以 # OOM for ... 注释显式裁掉了最大档位（如 kernels.py），这也是评测时需要注意的硬件内存边界。

六、框架机制：计时、精度校验、带宽换算与几何平均

所有被测逻辑都落在 utils.py 的 BenchmarkKernel 基类及其辅助函数上，理解它能帮你正确解读输出的每一行。

6.1 单点延迟测量

benchmark_kernel_in_milliseconds（utils.py）执行两步：

1. 先跑 5 次 warmup；
2. 在 torch.compiler.set_stance("fail_on_recompile") 保护下，调用 torch._inductor.runtime.benchmarking 的 benchmarker.benchmark_gpu 测得 GPU 延迟并返回（毫秒）。

fail_on_recompile stance 保证一旦被测函数在基准阶段发生重编译就会失败，避免把"重编译时间"误算进内核延迟。

6.2 指标换算：从延迟到内存带宽

每个 shape 的计时结果会封装成 Performance dataclass（utils.py），字段包括：

• latency（毫秒）：实测延迟；
• memory_bytes：由各内核重写的 get_memory_bytes 估算的访存字节数；
• memory_bandwidth：memory_bytes / (latency/1000) / 1e9，即 GB/s；
• compute_intensity：FLOPs/byte（当前默认 0，为扩展预留）。

每个内核类都手写了访存模型，例如：

• CrossEntropy 前向 = 读 x(M×N) + 读 target(M) + 写 loss(M)，见 kernels.py；
• Softmax 前向 = 读 x + 写 y，即 2×M×N，见 kernels.py；
• RMSNorm 前向 = 读 x + 写 y + 读 w（2×M×N + N），见 kernels.py。

由于这几个算子都是典型的内存受限算子，用"内存带宽 GB/s"而不是裸延迟做横轴比较更能说明内核质量——带宽越接近硬件上限，说明融合与访存越优。控制台每跑完一个点都会打印一行，例如 Performance 的 __str__ 输出 setting/latency/memory bandwidth（utils.py）。

6.3 精度校验（accuracy check）

每次 benchmark_single_shape 在测完所有后端后默认执行 check_accuracy（utils.py），流程是：

1. 对每个后端 clone_inputs 克隆一份输入（保留 requires_grad），见 utils.py；
2. 以 eager 输出为 gold，逐一用 torch.testing.assert_close 对比其余后端；
3. 若输出需要梯度，还会进一步校验反向后的 .grad；
4. 逐后端打印绿色 ✓ succeed 或红色 ✗ failed。

若指定 --tolerance，会把该值同时作为 atol/rtol 传入；若开启 --exit-on-accuracy-failure，任一后端失败即 sys.exit(1)（utils.py），适合接入 CI。值得一提的坑是 Quack 的 cross_entropy 只返回 float32 的 loss，比较前需 .to(gold.dtype) 对齐，否则会误报失败（见 kernels.py）。

6.4 汇总报告与可视化

每个 benchmark 跑完后调用 report_geomean_speedup()（utils.py）：以 eager 各 shape 延迟为分母，逐个后端计算 shape 级加速比，再用 scipy.stats.gmean 求几何平均加速比，例如输出 quack N data points, 2.31x speedup 这样的摘要行，避免个别 shape 的离群值扭曲结论。

--visualize 触发 visualize()，调用 visualize_comparison（utils.py）用 matplotlib 绘制各后端的 内存带宽对比曲线：横轴为 shape（自动去掉 shape: [] 前缀便于阅读），纵轴为 GB/s，不同后端使用统一的颜色方案（eager 蓝、compiled 橙、quack 绿、liger 红等，见 utils.py），图标题带 GPU 名。图像以 300 dpi 保存到 pics/<KernelName>_bench.png（utils.py），方便归档进实验报告。

七、一次完整评测流程参考

综合以上内容，推荐在具备 NVIDIA GPU 的机器上按如下流程操作：

# 1. 进入 benchmark 目录并安装依赖（假设 PyTorch 已装）
cd benchmarks/dynamo/genai_layers
pip install -r requirements.txt

# 2. 确认当前环境与可用 benchmark
python benchmark.py --list

# 3. 先单点验证精度与性能
python benchmark.py rmsnorm_forward --print-benchmark-result

# 4. 全量评测并输出对比图
python benchmark.py --all --visualize --print-benchmark-result

# 5. 精度敏感场景收紧容差；CI 场景失败即中断
python benchmark.py softmax_backward --tolerance 1e-2 --exit-on-accuracy-failure

# 6. 对比自定义 inductor 编译配置（会额外生成一条独立曲线）
python benchmark.py cross_entropy_forward \
  --custom-compile-name my_config \
  --custom-compile-options '{"max_autotune": true}'

注意事项与限制（均由源码可证）：

• 全程需要 CUDA GPU，且 Quack 路径强依赖 nvidia-cutlass-dsl 与 BFloat16；
• LayerNormBackward 无 quack 后端、RMSNormBackward 在超大 shape 下已由作者裁掉 OOM 档位，属预期行为，不要误判为故障；
• eager 的 Norm 参考实现使用 float32 中间精度，精度校验天然偏严，非 eager 融合内核在低精度下需关注 atol/rtol 设置。

八、小结：如何把该套件用起来

genai_layers 基准套件为 GenAI 高频算子提供了一条低摩擦的横向评测路径：一条 pip install -r requirements.txt 命令完成依赖，四种后端在统一 shape 矩阵下自动完成计时、带宽换算、精度互检与几何平均汇总。无论是评估融合内核（Quack/Liger）的收益、扫描 torch.compile 的 inductor 配置，还是追踪某次算子改动带来的回归，都可以直接复用这套注册表 + 基类的扩展模式——在 kernels.py 中新增内核类并注册进 benchmark.py 的 BENCHMARK_REGISTRY，即可纳入现有评测闭环。需要深入源码时，可重点研读 utils.py 的计时与精度逻辑，这是保证结论可信度的关键所在。