NVIDIA CUDA 中的内核融合：优化内存流量并启动开销

针对 GPU 优化代码的方法有很多。在本文中，您将了解内核融合如何提高内存带宽并减少内核启动用度，以及在 NVIDIA CUDA 代码中应用该技术的多种方法。

在编写 GPU 代码时，一个常见的瓶颈是 GPU 计算速度极快，即使是高带宽设备内存也无法完全使用 GPU 内核。内核融合通过将多个 GPU 操作合并到单个设备内核中来解决这一问题，因此中间结果不需要在全局内存中来回切换，也不需要单独的内核启动。

本文介绍的是融合核函数，以便将中间结果保留在寄存器中，并减少内存传输。CUDA 计算图提供另一种融合，但位于不同的层。图形将核函数启动、内存复制和同步的序列捕获到一个可重复使用的对象中，主机只需调用一次即可发送该对象。它不会融合内核主体。

内核融合的实际应用

图形中的内核单独运行，并通过全局内存传递中间结果。在图形中包装原始基准可在几微秒内从主机端删除数据，但通过中间缓冲区的 GiB 往返次数保持不变。这两种方法相辅相成，可以一起使用。

我们来看一个简单的例子：sum(abs(x))。我们读取数组，获取每个元素的绝对值，并返回总和。

原生实现使用两个内核：一个用于将 abs 计算为与输入大小相同的临时中间缓冲区，另一个用于将该缓冲区缩减为单个数字。这虽然有效，但速度很慢，这说明了更广泛的问题。

template <typename Config>
__global__ void abs_kernel(Config config,
                           cuda::std::span<const float> in,
                           cuda::std::span<float> tmp)
{
    /* Base SIMT implementation using fsabs */
}

template <typename Config>
__global__ void sum_kernel(Config config,
                           cuda::std::span<const float> tmp,
                           cuda::std::span<float> out)
{
    /* CUB block-reduce + atomic add into out[0] */
}

int main()
{
    /* device, stream, and buffers via cuda::make_buffer
     * See the CCCL Runtime post for the full setup. */

    auto config = cuda::distribute<BLOCK_THREADS>(in.size());

    cuda::launch(stream, config, abs_kernel<decltype(config)>, in,  tmp);
    cuda::launch(stream, config, sum_kernel<decltype(config)>, tmp, out);
}

本文中介绍的 C++ 示例使用了最新的 NVIDIA CCCL 运行时接口，包括在 CUDA 13.2 中首次推出的 cuda::std::span、cuda::launch 和 cuda::make_buffer。如需进一步探索这些基元，请参阅 CCCL 运行时：适用于 CUDA 的现代 C++ 运行时。

如果我们查看完整实现的 NVIDIA Nsight Systems 配置文件，如下所示：

图 1. 非融合核函数的 Nsight Systems 配置文件，用于求和 (abs ( x)) 。每个内核都独立执行，并且需要中间结果

手动核函数融合

首先，我们可以将这两个内核重写为单个 sum_abs_kernel。这将删除中间缓冲区，该缓冲区仅用于在内核之间传递数据。如果两个操作位于同一内核中，则缓冲区将消失，操作的性能也会提高。

template <typename Config>
__global__ void sum_abs_kernel(Config config,
                               cuda::std::span<const float> x,
                               cuda::std::span<float> out)
{
    using BlockReduce = cub::BlockReduce<float, BLOCK_THREADS>;
    __shared__ typename BlockReduce::TempStorage temp_storage;

    // Grid-stride loop. abs() is computed inline -- no tmp buffer.
    float thread_sum = 0.0f;
    const auto tid    = cuda::gpu_thread.rank(cuda::grid, config);
    const auto stride = cuda::gpu_thread.count(cuda::grid, config);
    for (size_t i = tid; i < x.size(); i += stride) {
        thread_sum += fabsf(x[i]);
    }

    // Block-level reduction in shared memory, via CUB.
    float block_sum = BlockReduce(temp_storage).Sum(thread_sum);

    // One thread per block atomically accumulates into the global result.
    if (threadIdx.x == 0) {
        cuda::atomic_ref<float, cuda::thread_scope_device> r(out[0]);
        r.fetch_add(block_sum, cuda::memory_order_relaxed);
    }
}

int main()
{
    /* device, stream, and buffers -- see the CCCL Runtime post. */

    // Fixed grid of 1024 blocks for the grid-stride pattern.
    auto config = cuda::make_config(cuda::make_hierarchy(
        cuda::grid_dims(1024),
        cuda::block_dims<BLOCK_THREADS>()));
    cuda::launch(stream, config, sum_abs_kernel<decltype(config)>, x, out);
}

• fabsf(x[i]) 在寄存器中内联计算，而不是从另一个内核编写的中间缓冲区中读取。
• 网格跨度循环允许每个线程将多个元素累加到单个寄存器 (thread_sum) 中，因此部分总和永远不会触及全局内存。
• cub::BlockReduce 处理共享内存中的每块归约，而 cuda::atomic_ref 允许每个块中的一个线程将其部分求和添加到全局结果中。

如果我们查看融合实现的 Nsight Systems 配置文件，它会相应地折叠：

图 2. 手动融合内核的 Nsight Systems 配置文件。sum_abs_kernel  是执行整个操作的单个内核n

这些数字适用于 NVIDIA GeForce RTX 4090。显存带宽约为 850 GiB/s，约为理论峰值 1008 GB/s 的 90%。

我们受到带宽限制，因为 GPU 所做的内存工作与物理特性所能实现的一样多。融合版本的速度更快，因为需要处理的内存工作量更少。

像这样的手写 CUDA C++ 内核有时是正确的解决方案，因为您可以完全控制并尽可能接近峰值性能。缺点是编写内核所需的工作量，以及高维护成本和跨 GPU 架构的可移植性有限。

CUB、libcu++、nvmath.device 和 NVSHMEM 等设备端库可以为 FFT、矩阵乘法、归约、扫描、排序、通信等提供可重复使用的基元，从而简化手动融合。

隐式内核融合

编写您自己的内核作品。在这种情况下，操作本质上是“应用 abs，同时减少”。我们可以使用简单的 Python 描述计算，并让编译器生成核函数，而不是在 CUDA C++ 中手动表达。

这就是 torch.compile 的作用。您为其指定了一个函数，它会追踪张量的运作情况，然后 PyTorch 的编译器 Torch Inductor 会生成一个在 GPU 上运行的内核。

import torch

x = torch.randn(N, device="cuda")

def sum_abs(t):
    return t.abs().sum()

compiled = torch.compile(sum_abs)
result = compiled(x)

通过将 sum_abs 函数包装为 torch.compile，我们得到了一个融合的内核。在 eager 版本 (不含 torch.compile) 中，我们得到了预期结果：两个内核、一个元素级 (abs) 和一个 reduce (sum) 。

 

在编译后的版本中，我们还获得了两个内核：triton_red_fused_abs_sum_0 和 triton_red_fused_abs_sum_1。

为什么？因为这是 Inductor 编译器生成的结果。

 

要了解 Inductor 的功能，请查看生成的 Triton 内核：

TORCH_LOGS="output_code" python implicit_fusion.py

第一个内核的有趣部分：

tmp0 = tl.load(in_ptr0 + ...)   # load a chunk of input
tmp1 = tl_math.abs(tmp0)        # apply abs in registers, inline
tmp2 = _tmp3 + tmp1             # accumulate into the running sum

abs 发生在寄存器中，在加载之后，在归约核函数内。没有单独的 abs 内核，也没有每个元素的中间缓冲区。在一个块中运行的第二个内核会读取 2048 个部分结果，并将其缩减为一个最终标量。您可以在这篇 PyTorch 博客文章 中详细了解这一点。

在实践中，torch.compile 融合了运算，但将归约分为两个阶段：一个核函数将输入减少为一组按块计算的小部分求和，第二个核函数将这些部分求和添加到一个数字中。在手动版本中，我们做出了不同的选择：使用原子累加到一个全局浮点运算中。这两个版本通过全局显存移动一个 GiB，而不是 3 GiB。编译器选择了合理的策略，但与我们选择的策略不同。

这既有好的一面，也有坏的一面。编译器免费为我们提供了一个融合内核，但我们牺牲了一些可预测性。不同的 dtype、不同的形状或在中间添加的小运算可以生成一组完全不同的内核。有时你会得到一个内核。有时你会得到四个。有时，您所期待的融合会悄无声息地停止发生，毫无预警。

显式核函数融合：回调函数、迭代器和结语

如果我们希望 Python 的工作效率和手写内核的可预测性，该怎么办？这就是 cuda.compute 的意义所在。

cuda.compute 提供全设备算法 (例如 reduce、scan、sort 和 transform) ，以及一组您可以合成的迭代器，以控制这些算法读取其输入的方式。您无需直接编写核函数。相反，您需要合成计算并将其传递给算法，而融合是因为您明确要求它。

cuda.compute 中相同的 sum(abs(x)) 操作如下所示：

import torch
import numpy as np
from cuda.compute import (
    Determinism, OpKind, TransformIterator, reduce_into,
)

x = torch.randn(N, device="cuda")
out = torch.empty(1, dtype=torch.float32, device="cuda")
h_init = np.array([0.0], dtype=np.float32)

abs_it = TransformIterator(x, lambda a: abs(a))
reduce_into(
    d_in=abs_it, d_out=out, num_items=N,
    op=OpKind.PLUS, h_init=h_init,
    determinism=Determinism.NOT_GUARANTEED,
)

TransformIterator 不会分配任何内容或启动核函数。这是一个延迟视图，显示“当有人要求我输入元素 i 时，请阅读 x[i] 并对其应用 abs”。然后，reduce_into 在整个设备范围内执行归约，通过迭代器提取元素。

abs 在寄存器中内联运行，因为每个元素都是从全局内存加载的。在任何时候都没有针对每个元素的中间缓冲区。另请注意，cuda.compute 适用于 PyTorch 张量。它支持实现 DLPack 或 CUDA 数组 API 的输入，例如 CuPy 数组。

将“数据是什么”与“如何处理”分开后，生成元素 i 将成为任意内联计算：加载、转换、收集或结构解包。该算法会将这些操作线程化为单遍，而无需中间存储。这种解决方案具有确定性、可移植性和可组合性，非常适合库作者和性能应用程序。

如果使用 Nsight Systems 对此进行分析，我们得到的结果类似于由几行 Python 生成的单个融合核函数：

图 5. 用于 cuda.compute 显式融合的 Nsight Systems 配置文件。一个称为 NondeterministicDeviceReduceAtomicKernel 的内核

cuda.compute 所依托的 CUB 库 会生成针对我们所述操作进行优化的单个融合核函数。它为许多算法提供了经过实际测试的实现，以及 CCCL 提供的保证，例如针对新硬件的优化。

与编译器方法不同，该方法不会因为代码更改较小、编译器版本不同或编译器决策而悄无声息地更改融合策略，因为您明确决定了应融合的内容。这可提供更具确定性的结果，同时仍可控制最终行为，包括浮点确定性等细节。

所有三种融合方法最终都位于同一位置，显存流量为 1 GiB，速度大约是原生版本的 3 倍。借助手动融合，您可以编写内核，因此融合是显式的，在您的控制之下。借助编译器融合，编译器可为您选择。

借助 cuda.compute，您可以使用 Python 人体工程学编写合成。与手动融合一样，这明确保证了转换融合到归约中。额外的好处是，此合成由 CUB 提供支持，因此您可以获得经过良好测试、高度优化的内部实现，而不是必须自行维护和调整的自定义内核。

借助 cuda.compute，Python 开发者可以使用 Thrust 和 CUB 长期以来在 C++ 中提供的基于迭代器的确定性融合。它使更高级别的 API 和框架能够像使用自定义内核一样运行，而无需开发者自行编写这些内核。

您可以在文档和 Python 的 cuda.compute 文档中详细了解，以及 cuda-samples 中的 Python 示例。