GPU 上的两个数字
你已经知道如何在 Python 中执行 c = a + b 了。你做过上千次。这里我们要讨论的是 PyTorch 张量中的情况。张量就是一个数字数组。把它放到 GPU 上意味着该数组驻留在 GPU 的内存中，而不是普通的 CPU 内存中。当 a 和 b 是 GPU 上的两个张量时，这一行代码完成得足够快，以至于你从来不会去思考它。
现在缩小范围。假设 a 和 b 是两个单独的浮点数，都位于 GPU 上。同一行代码：c = a + b。芯片上实际运行的是什么？
答案是 内核。在这个语境下，内核是一个小型程序，GPU 会在某些数据上运行它。不是你笔记本电脑启动时用的操作系统内核，也不是线性代数教材里的数学内核。这个词被重用了很多次，这可不是你的错。在 GPU 领域，内核只意味着：一个被 GPU 告知立即运行的小函数，在它被交给的数据上并行执行。
读完本文后，你将能够看着一段 PyTorch 代码，数出 GPU 会运行多少个内核。这听起来像个小技巧，确实也是，但它是摆脱“GPU 是黑盒”这种感觉的第一步——那种模型变慢却不知原因的感觉。任何关于 GPU 性能的问题，最终都会归结为“哪些内核运行了，它们做了什么”。所以这就是我们开始的地方。
你的第一个内核
让我们把 a 和 b 稍微变大一点。这次是长度为 8 的张量。仍然是一行 PyTorch 代码：c = a + b。
当你运行这行代码时，你的 CPU（实际执行 Python 的机器）会告诉 GPU：嘿，在这份数据上运行这个程序。这个指令被称为 启动。被启动的是一个内核：一个准备好运行的程序。启动本身很便宜，每个只需几微秒。真正昂贵的是 围绕 启动的部分（数据进入 GPU，结果返回），而这正是我们接下来要计数的。
在内核内部，实际工作是由称为 线程 的微小工作者完成的。GPU 有成千上万个这样的线程可用。对于我们的长度为 8 的加法，8 个线程接手工作：线程 0 处理元素 0，线程 1 处理元素 1，以此类推直到线程 7。每个线程运行相同的微小程序：读取 a 的一个元素，读取 b 的对应元素，将它们相加，然后将结果写入 c。
（实际上，GPU 会以固定大小的组（称为 warp，在 NVIDIA 卡上始终是 32 个线程）来启动线程，当你的数组不能被整除时，它会屏蔽掉多余的线程。现在可以忽略这一点。）
所以我们有一行 PyTorch 代码，一次启动，一个内核，8 个线程执行 8 次加法。现在让我们数一下实际跨芯片传输了什么。为了执行加法，每个线程需要它的 a 元素和 b 元素。那就是 8 次 a 的读取和 8 次 b 的读取。然后每个线程将其结果写入 c。那是 8 次写入。
这些读取和写入会进入紧挨着 GPU 芯片的大内存。在数据中心卡（A100、H100）上，这种内存称为 HBM（高带宽内存）。在消费级卡（RTX 4090）和人们很可能实际尝试的 Colab 风格 T4 上，它称为 GDDR。无论哪种方式，它都是紧挨着芯片的快速内存，我们就叫它 GPU 内存。它很快，但不是免费的，每次访问都有代价。
一个内核 = 一次启动 = 一次数据遍历。 无论内核在其内部做什么，它在边缘的读取和写入（去 GPU 内存取输入，写回输出）才是成本所在。这就是全部模式。
当张量变大时，这些都不会改变。同一行 PyTorch 代码，同一个内核，只是更多的线程。如果 a 和 b 各有 100 万个元素，GPU 会用更多的线程启动同一个内核。计算量缩放，字节移动量缩放，但心智模型不变。一行代码，一个内核。
两个操作之间发生了什么
1c = (a + b).relu()
你足够了解 Python，知道这是两个操作：一个加法，然后一个 relu，链式调用。在解释器中，这是两次函数调用。在 GPU 上，在 eager 模式的 PyTorch 中，这是两次内核启动：一次给加法，一次给 relu。到目前为止，没什么令人惊讶的。
真正有趣的是两个内核之间发生了什么。
当加法完成时，它必须把结果放在 某处。那个某处就是 GPU 内存。加法将整个中间数组（称为 tmp）写入内存。然后片刻之后，relu 启动，它的第一个任务是从内存中读回同一个 tmp 数组。它读取整个数组，对每个元素应用 relu，然后将结果写入 c。
数一下这两个内核的内存流量：
- 加法： 读取
a，读取b，写入tmp。三次数组大小的传输。 - Relu： 读取
tmp，写入c。两次。
总共 五次数组大小的传输。相比之下，上一节中单独的长度为 8 的加法只有三次。在链中加上 .relu() 不仅仅花费了 relu 的计算时间。它还花费了数组在 GPU 内存中完整往返一次的时间，因为 tmp 必须被写出来，以便下一个内核能读回去。
没有任何缓存。tmp 没有机会留在寄存器或快速的本地缓存中。它被送到 GPU 内存（那种慢且远的内存）然后立即被读回来。这两个内核彼此陌生。它们必须通过两者都懂得如何通信的媒介——GPU 内存——来传递数据。
为什么 PyTorch 要这样做？因为在 eager 模式下，当你写 a + b 时，PyTorch 会 立即 运行它。它不知道你接下来要调用 .relu()。每个操作在其 Python 行执行时立即被分发。没有计划，没有前瞻。每个操作独立存在，产生一个真正的数组，并通过内存传递给下一个操作。
这就是你要记住的模式。 你的 PyTorch 代码中的每个中间值都会被物理地写入 GPU 内存，然后被下一个操作读回。每一个。这就是“内核计数”真正衡量的东西。每多一个内核，就是你的数据在 GPU 内存中多一次往返。
融合：两个操作，一个内核
想象一个内核一次性完成所有这些：读取它的 a 元素，读取它的 b 元素，将它们相加，对结果应用 relu（全部在内核内部，在微小的每线程暂存空间上完成，从不离开芯片），然后才将最终值写入 c。中间值 (a + b) 仍然存在，但只在内核内部，在每个线程的私有暂存空间中。它从未被写入 GPU 内存。tmp 作为一个真正的数组，根本不存在。
现在数一下传输次数。读取 a：每个元素 1 次。读取 b：每个元素 1 次。写入 c：每个元素 1 次。三次数组大小的传输。 与两个内核版本的计算量相同，但少了两次往返。
在长度为 8 时，这可以忽略不计。没人在意。在长度为 100 万或 1 亿时，那些额外的往返会成为运行时间中的一大块，墙钟时间会反映出来。为什么 内存流量会如此主导，这是本系列第 2 篇文章的全部主题，所以我在这里不讨论“为什么”。现在关键是：相同计算量，更少往返，实际更快。
这种将原本会是独立内核的操作组合成单个内核，使得中间值永远不必访问 GPU 内存的技巧，有一个名字。它叫做 融合。就这个词。就这么个想法。
现在尴尬的部分来了。手工编写那个组合内核对于加法加 relu 看起来很容易。两个操作。中间一行“计算”。但真正的 PyTorch 代码有几十个操作链在一起，每个都有各自的形状、数据类型和广播规则。编写一个能正确处理所有这些的融合内核是一项真正的工程工作。你通常不会手工编写这些常规的逐元素内核。
好消息：PyTorch 附带了一个工具，可以为你自动重写这种代码，正是针对这类情况。它叫做 torch.compile。
你几乎肯定见过 torch.compile 这样使用：
1model = torch.compile(model)
一行代码。互联网上有人告诉你它能让事情变快。这里用通俗英语解释它实际做了什么：torch.compile 不会像 eager 模式那样逐个运行你的操作，而是捕获你的函数执行的张量操作，寻找合并它们的机会，并生成优化后的代码。后续调用如果匹配相同的假设，可以重用该代码。
我们在上面纸上做的融合（加法与 relu，共享一个内核，tmp 从不接触内存）正是 torch.compile 会自动对你的代码做的事情，只要操作足够简单。当人们说 torch.compile “让 PyTorch 更快”时，这占了很大一部分。
对于 torch.compile 无法自行融合的情况（它不认识的自定义操作、不常见的规约、奇怪的内存布局），仍然需要有人手工编写内核。这就是 Triton 和 CUDA 等工具的用途。另文讨论。
亲眼看看
以上所有内容都是纸上谈兵地数内核。现在是时候在真正的 GPU 上数一数了。如果你有任何带 CUDA GPU 的机器（工作站、Colab 笔记本、云实例），你可以在几分钟内自己运行这个。
工具是 torch.profiler。它内置于 PyTorch 中。它所做的就是记录 GPU 在你代码运行时实际做了什么，然后返回一个你可以阅读的表格。
步骤 1：eager 版本。
将两个操作的行包装在一个函数中，这样我们就有东西可以调用了：
1import torch2from torch.profiler import profile, ProfilerActivity34def add_relu(a, b):5 return (a + b).relu()67a = torch.randn(1_000_000, device="cuda")8b = torch.randn(1_000_000, device="cuda")910with profile(activities=[ProfilerActivity.CUDA]) as prof:11 add_relu(a, b)12 torch.cuda.synchronize()1314print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
torch.cuda.synchronize() 只是用来确保 GPU 在读取时间之前已经完成。GPU 工作是异步运行的，没有同步，你有时会测量到启动开销而不是实际内核工作。
步骤 2：读取输出。
你的实际分析器输出可能会比你预期的行数多。会混入一大堆内存分配和 PyTorch 簿记行。我们关心的行是 CUDA 内核，即 GPU 实际运行的函数。查找名称中包含 kernel 的行。这两个行大致如下：
1vectorized_elementwise_kernel<...BinaryFunctor_add...> 1 12us2vectorized_elementwise_kernel<...threshold_kernel...> 1 9us
确切的模板名称会随 PyTorch 版本变化（relu 通常显示为 threshold，因为那是底层操作，而 add 有时显示为 CUDAFunctor_add）。不要试图解析整个内容。只需数行数。两行。两个内核。 一个用于加法，一个用于 relu。正是我们上一节所说的。
步骤 3：编译版本。
改一行代码。将函数包装在 torch.compile 中：
1compiled = torch.compile(add_relu)
现在有一个值得注意的陷阱，在运行之前你应该知道。对 torch.compile 函数的 第一次 调用很慢，有时甚至异常慢，因为那正是 torch.compile 完成其工作的时候：分析你的代码，找出要融合的内容，生成融合后的内核。如果你分析第一次调用，你测量的是编译步骤，而不是内核。所以模式是：先调用一次来预热，丢弃结果，然后 分析。
1compiled(a, b) # 预热，丢弃2torch.cuda.synchronize()34with profile(activities=[ProfilerActivity.CUDA]) as prof:5 compiled(a, b)6 torch.cuda.synchronize()78print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
步骤 4：再次读取输出。
这次：
1triton_poi_fused_add_relu_0 1 14us
一行。一个内核。 名称甚至告诉你它做了什么：融合了加法与 relu。与之前相同的计算量，一次启动而不是两次。
你刚刚完成了一句话就能概括的本文所讨论的事情：你要求 PyTorch 将两个操作合并为一个内核，观察了分析器，并确认计数从两个降为一个。融合，在实际使用中，在你的机器上。
如果你想更戏剧性地看到效果，可以在几种不同的张量大小下尝试。在长度为 100 时，两个版本都运行得很快，差异在噪声中丢失。在长度为 1000 万或 1 亿时，编译版本开始明显领先，因为我们在那个规模下削减的往返是实际工作的重要组成部分。
计数内核不再是抽象的建议。你现在有办法验证了。
总结一下
这里用一句话概括全部内容。
你的 PyTorch 代码在 GPU 上运行时，会变成一系列内核。每个内核是一次启动，一次数据遍历，一次通过 GPU 内存取输入并写输出的往返。简单操作变成一个内核。操作链默认变成每个操作一个内核，中间值在内存中往返。torch.compile 可以为你融合简单的链，使这些中间值从不接触内存。更少的内核通常意味着更少的内存流量。而更少的内存流量通常意味着更快的速度。
