NVIDIA CUDA 中的核心融合：最佳化記憶體流量與啟動開銷

最佳化 GPU 程式碼的方法有很多。在這篇文章中，您將學習如何利用核心融合來改善記憶體頻寬並減少核心啟動開銷，以及在 NVIDIA CUDA 程式碼中實作的多種方法…

NVIDIA Technical Blog · Jul 10, 2026

Kernel Fusion in NVIDIA CUDA: Optimizing Memory Traffic and Launch Overhead

最佳化 GPU 程式碼的方法有很多。在這篇文章中，您將學習如何利用核心融合來改善記憶體頻寬並減少核心啟動開銷，以及在 NVIDIA CUDA 程式碼中實作的多種方法。

撰寫 GPU 程式碼時常見的瓶頸在於，GPU 運算速度非常快，即使是高頻寬的裝置記憶體也無法完全利用 GPU 核心。核心融合透過將多個 GPU 運算合併到單一裝置核心中來解決這個問題，因此中間結果不需要來回傳輸至全域記憶體，也不需要分開啟動核心。

本文討論的是融合核心主體，讓中間結果保留在暫存器中，並減少記憶體傳輸。CUDA Graphs 提供了另一種融合方式，但是在不同的層級。Graph 會將一連串的核心啟動、記憶體複製與同步動作擷取成單一可重複使用的物件，主機只需呼叫一次即可派遣。它並不會融合核心主體。

核心融合的實務應用

在 graph 內執行的核心各自獨立運作，並透過全域記憶體傳遞中間結果。將一個簡單的基準實作包裝在 graph 中，可以從主機端省下數微秒的時間，但透過中間緩衝區的一 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::std::span、cuda::launch 與 cuda::make_buffer（這些功能在 CUDA 13.2 中首次推出）。若想進一步了解這些基本元素，請參考 CCCL Runtime: A Modern C++ Runtime for CUDA。

若我們查看完整實作的 NVIDIA Nsight Systems 分析結果，會看到以下畫面：

圖 1. sum(abs(x)) 非融合核心的 Nsight Systems 分析圖。每個核心獨立執行，且需要中間結果

手動核心融合

我們能做的第一件事，就是將這兩個核心重寫成單一的 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]) 是在暫存器中即時計算，而非從另一個核心寫入的中間緩衝區讀取。
grid-stride 迴圈讓每個執行緒將多個元素累加到單一暫存器（thread_sum）中，因此部分總和永遠不會接觸全域記憶體。
cub::BlockReduce 負責在共享記憶體中進行區塊層級的歸約，而 cuda::atomic_ref 則讓每個區塊中的一個執行緒將其部分總和加入全域結果。

若我們查看融合實作的 Nsight Systems 分析結果，會發現它已簡化為單一核心：

圖 2. 手動融合核心的 Nsight Systems 分析圖。單一 sum_abs_kernel 完成整個運算

Metric	Naive	Manual fusion
Kernels Launched	2	1
Intermediate buffer	Yes	No
Bytes moved through global memory	3 GB	1 GB
Time	3.51 ms = 2.28ms + 1.23 ms	1.18 ms
Effective memory bandwidth	855 GiB/s	851 GiB/s
Speedup vs. naive	–	3 times

表 1. 手動融合核心的基準測試結果

這些數據來自 NVIDIA GeForce RTX 4090。記憶體頻寬約為 850 GiB/s，大約是其理論峰值 1,008 GB/s 的 90%。

我們受到頻寬限制，因為 GPU 已盡可能執行記憶體操作。融合版本之所以更快，是因為需要執行的記憶體操作較少。

像這樣手寫的 CUDA C++ 核心有時是正確的解決方案，因為您擁有完整控制權，能盡可能接近峰值效能。缺點是撰寫核心所需的努力，以及高昂的維護成本與跨 GPU 架構的有限可攜性。

裝置端函式庫如 CUB、libcu++、nvmath.device 與 NVSHMEM 可以透過提供 FFT、矩陣乘法、歸約、掃描、排序、通訊等可重複使用的基本元件，簡化手動融合。

隱式核心融合

自己撰寫核心是可行的。在這個案例中，運算本質上是「套用 abs 並進行歸約」。我們不必在 CUDA C++ 中手動表示這個運算，而是可以用純 Python 描述運算，再讓編譯器產生核心。

這就是 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)

透過 torch.compile 包裝 sum_abs 函式，我們可以獲得融合的核心。在 eager 版本（未使用 torch.compile）中，我們會看到預期的結果：兩個核心，一個是逐元素運算（abs），另一個是歸約（sum）。

圖 3. 單次 eager PyTorch 運算的 PyTorch 分析畫面（使用 Perfetto.dev 檢視）。顯示兩個核心，每個運算各一個

在編譯版本中，我們同樣會得到兩個核心：triton_red_fused_abs_sum_0 與 triton_red_fused_abs_sum_1。

原因在於這是 Inductor 編譯器產生的結果。

圖 4. 已編譯 PyTorch 的 PyTorch 分析畫面（使用 Perfetto.dev 檢視）。這些核心是由 PyTorch 產生，即使運算已被編譯器融合

若要了解 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 核心，也沒有逐元素的暫存中間緩衝區。第二個核心在單一區塊中執行，讀取 2,048 個部分結果並將它們歸約成單一最終純量。您可以在這篇 PyTorch 部落格文章中了解更多資訊。

實際上，torch.compile 融合了這個運算，但將歸約分成兩個階段：第一個核心將輸入歸約成一小組每區塊的部分總和，第二個核心將這些部分總和加總成單一數值。在手動版本中，我們做了不同的選擇：使用原子操作將結果累加到單一全域浮點數。兩個版本都只移動一 GiB 的全域記憶體，而非三 GiB。編譯器選擇了一個合理的策略，但與我們會選擇的策略不同。

這有好有壞。編譯器免費提供融合的核心，但我們犧牲了一些可預測性。不同的資料型態、不同的形狀，或是在中間加入一個小運算，都可能產生完全不同的核心組合。有時會得到一個核心，有時會得到四個核心，有時您所依賴的融合會在沒有警告的情況下悄悄停止。

顯式核心融合：回呼、迭代器與結尾運算

如果我們想要 Python 的生產力，同時保有手寫核心的可預測性呢？這就是 cuda.compute 的用途。

cuda.compute 提供了裝置端演算法，例如 reduce、scan、sort 與 transform，以及一組可組合的迭代器，讓您控制這些演算法如何讀取輸入。您不需要直接撰寫核心，而是將運算組合後傳遞給演算法，融合之所以發生是因為您明確要求。

同樣的 sum(abs(x)) 運算在 cuda.compute 中看起來像這樣：

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 Array API 的輸入，例如 CuPy 陣列。

一旦將「資料是什麼」與「如何走訪它」分開，產生元素 i 就變成任意的即時運算：載入、轉換、收集或結構解包。演算法會將這些運算串接成單一過程，而不需要中間儲存。這種解決方案具有確定性、可攜性與可組合性，非常適合函式庫作者與效能導向的應用程式。

若我們使用 Nsight Systems 進行分析，會得到類似由幾行 Python 程式碼產生的單一融合核心：

圖 5. 來自 cuda.compute 的顯式融合 Nsight Systems 分析圖。單一核心名為 NondeterministicDeviceReduceAtomicKernel

cuda.compute 建構於其上的 CUB 函式庫，會產生單一融合核心，並針對我們描述的運算進行最佳化。它提供了經過實戰測試的許多演算法實作，以及 CCCL 提供的保證，例如針對新硬體的最佳化。

與編譯器方法不同，它不會因為微小的程式碼變更、不同的編譯器版本或編譯器決策而悄悄改變融合策略，因為您明確決定應該融合什麼。這讓您能獲得更具確定性的結果，同時仍能控制最終行為，包括浮點數確定性等細節。

Metric	Naive	Manual Fusion	`torch.compile`	`cuda.compute`
Kernels Launched	2	1	2	1
Per-element intermediate buffer	yes	no	no	no
Speedup vs. naive	–	3x	3x	3x
Who decides what gets fused?	–	You	Compiler	You
Determinism	–	Full	Varies	Predictable

表 2. 所有核心融合方法的摘要

所有三種融合方法最終都會達到相同結果：一 GiB 的記憶體流量，以及比簡單版本快約 3 倍的效能。手動融合由您撰寫核心，因此融合是明確且由您控制的。編譯器融合則是由編譯器為您選擇。

使用 cuda.compute，您可以使用 Python 的易用性撰寫組合。與手動融合一樣，它提供明確的保證：轉換會融合到歸約中。額外的好處是這個組合由 CUB 提供支援，因此您可以在底層獲得經過充分測試、高度最佳化的實作，而不需要自己維護與調整自訂核心。

cuda.compute 讓 Python 開發者能夠使用 Thrust 與 CUB 長期以來在 C++ 中提供的迭代器式確定性融合。它讓高階 API 與框架能夠表現得像擁有自訂核心一樣，而不需要開發人員自己撰寫這些核心。

您可以在文件與 Python 的 cuda.compute 以及 cuda-samples 中了解更多資訊。

Read the original source

NVIDIA CUDA 中的核心融合：最佳化記憶體流量與啟動開銷

核心融合的實務應用

手動核心融合

隱式核心融合

顯式核心融合：回呼、迭代器與結尾運算

0 Comments