JanJan

我们如何（试图）准确地对 GPU 内核进行基准测试

如果你之前读过任何其他关于内核基准测试的文章，那么这篇博客中的很多信息对你来说可能已经是老生常谈了。事实上，我们想首先感谢那些在我们之前编写的许多内核基准测试指南中所付出的血汗和泪水，这些指南帮助我们编写了更好的基准测试代码，也为这篇博客的创作提供了指导。

在 Menlo，我们最近购入了一些 RTX PRO 6000 Blackwell 工作站版（在新标签页中打开），并正在努力让像 vLLM（在新标签页中打开） 这样的 LLM 推理引擎在其上运行得更快。我们一直在为 RTX PRO 6000 专门编写我们自己的内核，看看是否能改善我们硬件上的推理时间。

这篇博客将更详细地介绍我们机器学习效率团队是如何识别我们基准测试代码中的问题，以及我们如何借鉴各种优秀的基准测试指南来迭代改进它！闲话少说，让我们从简单的程序开始，逐步深入到 GPU 内核的基准测试。

内核与基准测试简介

对于刚接触 GPU 编程的读者来说，内核（在新标签页中打开）是程序员编写的一段 CUDA 代码，用于在 GPU 上执行所需的操作序列。这些内核一旦启动，就会由并行运行的线程执行，我们通常从一个线程块网格（在新标签页中打开）启动这些内核，它会在整个 GPU 的多个流式多处理器 (SMs)（在新标签页中打开）上执行我们的内核。

基准测试是高性能计算的一个基本方面。它使我们能够量化地比较不同问题规模下的内核性能，并理解各种超参数如何影响执行速度。对于 GPU 内核开发，基准测试有助于我们迭代优化内核，使其更好地利用 GPU。

话虽如此，**准确的内核基准测试**更为重要，因为对在 GPU 上运行的内核进行基准测试可能变得非常复杂，如果在编写基准测试脚本时不加小心，很容易掉入各种陷阱。一个很好的替代方案是使用 NVIDIA 通过其CUDA 工具包（在新标签页中打开）提供的工具，例如Nsight CUDA 分析工具接口（在新标签页中打开）（cupti）或使用Nsight Compute CLI（在新标签页中打开）（ncu），它们为基准测试内核的各种特性提供了精确的测量。对我们来说，我们想使用 Python，因为它能更方便地快速遍历不同的问题形状和内核，但这也意味着我们必须从头学习如何正确地对内核进行基准测试。

我们将展示一些如何在 GPU 上对内核进行基准测试的例子。此外，我们的大部分基准测试代码都选择了 Python，因为我们自己的代码库大部分是 Python，这样可以简单地集成进来。

对 CUDA 程序进行基准测试

Pytorch 提供了一个非常基础的 API 来帮助计时 torch 程序，可以参考这个教程（在新标签页中打开）。

我们可以看到一个基本的实现可能像这样简单：


def batched_dot_mul_sum(a, b)
'''通过乘法和求和计算批量点积'''
return a.mul(b).sum(-1)
num_threads = torch.get_num_threads()
print(f'在 {num_threads} 个线程上进行基准测试')
t0 = benchmark.Timer(
stmt='batched_dot_mul_sum(x, x)',
setup='from __main__ import batched_dot_mul_sum',
globals={'x': x},
num_threads=num_threads,
label='多线程批量点积',
sub_label='使用 mul 和 sum 实现')

在对内核进行基准测试时，有几个技巧我们应该遵循，以确保我们准确地测试我们的内核。

1. 始终不要使用你机器上的设置来对代码进行基准测试，而要使用**用户会看到的设置**。

如果你将在 H100 DGX 节点上部署你的模型，那么在 3090 上测试你的内核速度是毫无意义的。始终在你计划部署的硬件上对你的内核进行基准测试是一个好主意。

2. **预热你的内核**

看看教程里的这个片段。


mul_sum(x, x): 27.6 μs
mul_sum(x, x): 25.3 μs
bmm(x, x): 2775.5 μs
bmm(x, x): 22.4 μs

第一个用于 bmm 的内核运行时间要长得多。这是因为第一次运行时，大部分时间都花在了加载 cuBLAS（在新标签页中打开） 内核上。

预热你的内核可以很简单，就是在计时前先运行一次内核。这有助于预先加载这些内核，这样我们只测量内核运行所需的时间。

3. torch.cuda.synchronize 和 CUDA 事件

现在我们还将介绍一个新的 API，这是对内核进行基准测试的标准方法。CUDA 事件（在新标签页中打开）因多种原因而出色。最简单的原因是它从 GPU 的角度测量时间，而 time.time() 和 time.perf_counter() 是从 CPU 的角度测量时间。

此外，其简洁的 API 允许你像这样调用基准测试代码：


steps = 10
start_events = [torch.cuda.Event(enable_timing=True) for _ in range(steps)]
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(steps)]
for i in range(steps)
start_events[i].record()
run_kernel()
end_events[i].record()
torch.cuda.synchronize()
times = [s.elapsed_time(e) for s, e in zip(start_events, end_events)]

torch.cuda.synchronize 告诉 CPU 等待 GPU 上的工作完成，以便在同步后计算经过的时间，如下图所示：

图片 图 1：插图取自 https://www.speechmatics.com/company/articles-and-news/timing-operations-in-pytorch（在新标签页中打开）

4. 清空你的 L2 缓存

什么是 L2 缓存

当从 HBM 或 GDDR（在新标签页中打开） 读取或写入数据时，它首先会经过 L2 缓存（在新标签页中打开），这个缓存由所有流式多处理器（SM）（在新标签页中打开）共享。L2 缓存会缓存对本地和全局内存的数据访问，并帮助重用数据，而不是再次将其加载到共享内存（这可能很慢！）。

此外，与每个 SM 上都有的 L1 缓存不同，所有 SM 共享同一个 L2 缓存！

我们为什么需要清空 L2 缓存

根据这篇指南（在新标签页中打开），如果你之前已经预热或运行过内核，一些中间数据可能存储在 L2 缓存中，这意味着内核的速度可能会被误导性地加快。

然而，在实际应用场景中，你想测量内核运行的真实时间，而且在运行大型模型时，你通常会运行不止一个内核。这意味着你的缓存可能会频繁地被替换（thrash），并且不会为重用而存储特定内核的数据。因此，为了模拟这种行为，我们会事先清空 L2 缓存，以消除 L2 缓存的“帮助”。

此外，这在计算内核的数据重用时也变得容易得多，因为现在任何 L2 缓存的使用都独立于其他内核或运行。

不清空 L2 缓存的例子

之前，当我们最初对内核进行基准测试时，我们犯了一个小错误，没有清空 L2 缓存。

图片 图 2：我们的 SOL%（观测到的最大速度百分比）对于形状为 [2, 19456, 2560] 的行超过了 100%。

如何清空 L2 缓存

要清空它，我们应该添加以下代码行：


l2_size = torch.cuda.get_device_properties().L2_cache_size
cache = torch.empty(l2_size, dtype=torch.uint8, device="cuda")
#<你的基准测试代码在这里>
cache.zero_() # 清空 L2 缓存
# 如果你多次重复相同的过程，你应该在基准测试代码内部清空 L2 缓存

这会实例化一个与 L2 缓存大小相同的数据，通过就地清零，我们调用一个写入操作，该操作会经过 L2 缓存并将其清空。

清空 L2 缓存后，我们得到了一个更合理的结果：

图片 图 3：清空 L2 缓存后，新的 SOL% 所有值现在都在 100% 以下。

5. 对短时内核进行计时

最初，我们使用 Triton 的（在新标签页中打开） do_bench（在新标签页中打开） 进行基准测试，因为它已经完成了我们上面提到的所有事情，比如预热、CUDA 事件和清空 L2 缓存。然而，我们发现在对较小形状的内核进行精确基准测试时存在一个问题。在较小的形状上，内核可能太快，以至于它可能在 CPU 在 Python 中发出 CUDA 结束事件之前就完成了。

图片 图 4：取自 Speechmatics（在新标签页中打开），内核比 CUDA 事件结束启动更快，因此没有记录到内核的真实计时。

这导致内核看起来非常慢：

图片 图 5：Python 基准测试延迟与 ncu 计时（右）的并排比较，形状为 [2, 19456,2560]。ncu 记录的 71.36 μs 比 Python 的 103.9 μs 快得多。

为了解决这个问题，我们编写了一个自定义的 do_bench_cuda()，它在对每个形状进行基准测试之前插入一个虚拟的、不计时的 FP32 矩阵乘法，这样 CPU 就有足够的时间将 CUDA 结束事件加入队列。

这为我们的小 M 内核带来了更准确的延迟。

图片 图 6：插入虚拟矩阵乘法后，SOL% 有了显著改善。

然后，我们还在 5 份输入/输出数据副本上对每个形状重复基准测试函数，以使 CUDA 事件持续时间更长。

最后，这是我们用来对内核进行基准测试的 do_bench_cuda 函数。


import statistics
import torch
def do_bench_cuda(f, n_warmup: int = 10, n_repeats: int = 20)
l2_size = torch.cuda.get_device_properties().L2_cache_size
cache = torch.empty(l2_size, dtype=torch.uint8, device="cuda")
# 以防 CUDA 事件过短，先进行一次矩阵乘法
A = torch.randn(4096, 4096, dtype=torch.float32, device="cuda")
B = torch.randn(4096, 4096, dtype=torch.float32, device="cuda")
A @ B
# L2 缓存清空 + 预热
for _ in range(n_warmup)
cache.zero_()
f()
start_list = [torch.cuda.Event(enable_timing=True) for _ in range(n_repeats)]
end_list = [torch.cuda.Event(enable_timing=True) for _ in range(n_repeats)]
torch.cuda.synchronize()
for start, end in zip(start_list, end_list)
cache.zero_() # 清空 L2 缓存
A @ B # 添加一个重型任务来填满 GPU 流水线
start.record()
f()
end.record()
torch.cuda.synchronize()
timings = [start.elapsed_time(end) for start, end in zip(start_list, end_list)]
return statistics.median(timings)

6. 时钟速度

这是一个隐藏的问题，很难发现它正在对我们的内核造成问题。我们最初发现在对形状 [2048, 19456, 2560] 进行分析时，ncu 的延迟（676.64 μs）和 do_bench_cuda 的延迟（535 μs）之间存在差异，因为 do_bench 报告的时间比 ncu 的延迟快了约 140 μs。

可以看出，尽管我们大部分的内核基准测试代码库是 Python，但开发人员可能会犯错，因此拥有一个准确的内核计时参考点总是好的。Nsight Compute CLI（简称 ncu）是一个可以帮助我们准确测量内核延迟的工具，从中获得的值是检查我们自己基准测试代码是否合理的好参考。

6.1 时钟速度

首先，我们怀疑时钟速度可能是导致 ncu 的计时和我们自己的基准测试代码之间差异的原因之一。时钟速度会影响基准测试时间，因为它是 GPU 处理单元运行的速率，更高的时钟速度意味着每秒更多的操作，这既可能加速也可能减慢内核，具体取决于其实现方式。

图片 图 7：取自 GPU Mode 讲座 56（在新标签页中打开）。我们可以看到时钟速度影响内核性能。对于问题形状为 1024，提高时钟速度后变得更快，而对于问题形状为 384，提高时钟速度后变得更慢。

查看这个论坛帖子（在新标签页中打开），我们意识到导致差异的原因之一是 ncu 默认将时钟速度锁定到 GPU 的基础时钟速度。我们尝试通过将时钟速度锁定到基础时钟速度进行调查，也尝试使用 nvidia-smi -ac=<memClk>,<smClk> 将其锁定到最大时钟速度。根据 GPU Mode 讲座，这不是一个合适的解决方案。

这是由于以下原因：

  • 锁定到最大时钟速度没有帮助，因为它只是设定了 GPU 性能的上限，我们的 GPU 总是可以回到约 2287 Hz 的基础时钟速度，而不是 2617 Hz 的增强时钟速度。

  • 锁定到基础时钟速度也没有意义，因为它不能正确反映用户在我们的内核上获得的性能和体验，而这些内核在最佳情况下会以增强时钟速度运行。

然而，我们确实发现我们应该将 ncu 的 --clock-control 设置为 None，这样它就不会将自己限制在基础时钟速度。这有助于将 ncu 上的延迟从 676.64 μs 提高到 575 μs，当在相同的问题形状 [2048, 19456, 2560] 上进行分析时。

6.2 clock-control 后的差异

在撰写本文时，我们观察到 ncu 有时在相同的基准测试代码和相同的问题形状上会给出不同的延迟结果。其原因是，当我们将 clock-control 速度设置为 None 时，GPU 的时钟速度是随机的，因此会影响测量的内核延迟。一个更全面的方法是在不同的固定时钟速度下也对内核进行基准测试。图片 图 8：在相同的基准测试代码和问题形状上，我们可以看到持续时间有巨大的偏差，这是由 SM 频率的差异引起的。这与图 7 中显示的图表相呼应。

因此，ncu 的计时和我们自己的基准测试计时可能存在一些差异。要确定你的差异是否由 SM 频率引起，你可以使用 FLOPS 与 SM 时钟成正比的关系，因此它们的持续时间成反比。

在我们的案例中：544 / 2.14 (575 μs 内核的 SM 频率) * 2.28 (544 μs 内核的 SM 频率) = ~579，所以大部分差异来自 SM 频率的不同。

我们最终使用的命令是：

ncu -s 5 -k $kernel_name --clock-control none python3 benchmarks/bench_mm.py --profile 2048 19456 2560

参数解释：-s：跳过的内核数量；-k：要分析的内核名称；--clock-control：是否控制时钟速度。

以下是在进行了所有调整后，ncu 的基准测试延迟与我们的脚本的并排比较。

图片 图 9：上述 ncu 命令（左）（测量形状 [2048,19456,2560]）与我们自己的 Python 基准测试脚本（右）的并排比较。我们可以看到 ncu 中的 持续时间 (Duration) 和我们基准测试脚本的 延迟 (us) (Latency (us)) 测量之间最多有 10us 的差异。

结论与 TLDR (内容概要)

TLDR（太长不看，一句话总结），在进行基准测试时：

  1. 确保使用你打算部署的硬件
  2. 在对内核进行基准测试前进行预热
  3. 使用 CUDA 事件
  4. 清空你的 L2 缓存
  5. 使用虚拟的矩阵乘法使短时内核的计时更准确
  6. 确保你的时钟速度不会导致不一致的读数

我们希望这能帮助任何有兴趣对自己的内核进行基准测试的人，或者对 GPU 内核如何进行基准测试感兴趣的人。祝你基准测试愉快！

致谢与相关资源：

我们要感谢并归功于我们在探索如何最好地在我们的 GPU 上对内核进行基准测试的过程中所使用的许多资源和指南，如果没有这些出色的指南，我们的许多工作可能无法完成。