CUDA Graphs

一、引入

近年来，随着 GPU 架构的不断更迭，其性能也越来越好，那些个 GPU 操作(比如 CUDA Kernel 或者 Memory Copy)所花费的时间习惯都以微秒为单位进行测量。而不可忽视的一个问题就是 kernel launch 的时间，即将每个操作提交给 GPU 也会产生开销，通常也是微秒级的，如果在实际应用程序中存在大量 GPU 操作，那么 kernel launch 时间将不可忽略，如此必然会影响整体端到端性能。

CUDA Graphs 的设计初衷就为了解决以上问题，它允许将工作任务定义为图而不是单个操作，提供了一种通过单个 CPU 操作启动多个 GPU 操作的机制来解决上述问题，从而减少开销。

在传统的 GPU 编程中，每次要让 GPU 执行任务时，CPU 必须发送一个指令，并等待 GPU 完成任务后再发送下一个指令。这种“一个一个发送”的方式效率较低，因为 CPU 和 GPU 之间来回通信的开销会导致时间浪费。

CUDA Graph 通过把多个任务组成一个图结构，一次性提交给 GPU，这样 GPU 可以自己安排这些任务的执行顺序，而不需要每次都等待 CPU 的新指令，从而大大提高了效率。

二、组成

CUDA Graph 主要由以下两部分组成：

• 节点（Node）：图中的每个“节点”表示一个任务，比如一个计算任务或一个数据传输任务。
• 边（Edge）：图中的每条“边”表示任务之间的依赖关系。例如，任务 A 必须在任务 B 之前完成，那么就会有一条从 A 到 B 的边。

通过这种图结构，GPU 能够按顺序执行任务，并且可以在不影响依赖关系的前提下并行执行多个任务。

CUDA Graph 的基本流程如下：

1. 创建图：我们先定义好需要 GPU 执行的所有任务，并确定它们之间的顺序和依赖关系。
2. 录制图（Recording）：把这些任务记录成一个图结构。
3. 执行图（Execution）：把这个图提交给 GPU，让 GPU 按照图中的任务顺序来执行。

三、使用

在 CUDA 中使用 CUDA Graph 大致可以分为以下几步：

• 创建和开始录制图：定义一个 CUDA Graph 对象，准备好图的录制。
• 将操作加入到图中：在图的录制过程中，将各种 GPU 操作（比如内存拷贝、核函数调用等）加入到图中。
• 结束图的录制：结束录制并生成一个可执行的图对象。
• 执行图：将图提交给 GPU 执行。
• 清理资源：释放图和相关资源。

实例： 用 CUDA Graph 加速一个向量加法 C = A + B

#include <cuda_runtime.h>
#include <iostream>

// 核函数定义：向量加法
__global__ void vectorAdd(const float* A, const float* B, float* C, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        C[i] = A[i] + B[i];
    }
}

int main() {
    int N = 1 << 20;  // 向量长度
    size_t size = N * sizeof(float);

    // 分配并初始化主机内存
    float *h_A = (float*)malloc(size);
    float *h_B = (float*)malloc(size);
    float *h_C = (float*)malloc(size);
    for (int i = 0; i < N; i++) {
        h_A[i] = 1.0f;
        h_B[i] = 2.0f;
    }

    // 分配设备内存
    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, size);
    cudaMalloc(&d_B, size);
    cudaMalloc(&d_C, size);

    // 定义 CUDA Graph 对象和图实例
    cudaGraph_t graph;
    cudaGraphExec_t instance;

    // 开始录制 CUDA Graph
    cudaStream_t stream;
    cudaStreamCreate(&stream);
    cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);

    // 将操作加入到图中
    cudaMemcpyAsync(d_A, h_A, size, cudaMemcpyHostToDevice, stream);
    cudaMemcpyAsync(d_B, h_B, size, cudaMemcpyHostToDevice, stream);

    int blockSize = 256;
    int numBlocks = (N + blockSize - 1) / blockSize;
    vectorAdd<<<numBlocks, blockSize, 0, stream>>>(d_A, d_B, d_C, N);

    cudaMemcpyAsync(h_C, d_C, size, cudaMemcpyDeviceToHost, stream);

    // 结束录制并生成可执行图
    cudaStreamEndCapture(stream, &graph);
    cudaGraphInstantiate(&instance, graph, NULL, NULL, 0);

    // 执行图
    cudaGraphLaunch(instance, stream);
    cudaStreamSynchronize(stream);

    // 清理资源
    cudaGraphDestroy(graph);
    cudaGraphExecDestroy(instance);
    cudaStreamDestroy(stream);

    // 验证结果
    bool success = true;
    for (int i = 0; i < N; i++) {
        if (h_C[i] != 3.0f) {
            success = false;
            break;
        }
    }
    std::cout << (success ? "Success" : "Failure") << std::endl;

    // 释放内存
    free(h_A);
    free(h_B);
    free(h_C);
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);

    return 0;
}

kernel launch 有固定开销，具体数字随驱动、CPU、功耗状态和采集方式变化。LLM 推理一个 token 要触发很多小 kernel 时，这部分开销会在 timeline 上形成可见空洞，对小 model 和低 batch decode 尤其明显。

CUDA Graphs 思路：**把一连串 launch “录制”**成一个图对象，重放时只付一次 launch 开销。