cuda graph学习笔记

cuda graph是CUDA10.x发布的一个功能，用于把kernel和拷贝操作记录成一张图，之后每次launch就能整体交给GPU，省掉kernal之间的间隙

捕获规则

• cudagraph 最好在cudaStreamNonBlocking stream中捕获，不能在null stream中捕获
• 捕获之后，graph可以在任何stream中运行（包括nullstream）
• 如果输入输出的数据的地址改变，则需要重新捕获，这一点很重要，否则出来的数据不对，而且可能不会报错！！

具体示例

// 访存极轻的 kernel：n 小的时候执行时间远小于 launch 开销
__global__ void axpb_kernel(float* x, int n, float a, float b) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) x[i] = x[i] * a + b;
}

static void launch_step(float* x, int n, int k, cudaStream_t s) {
    const int block = 256;
    const int grid = (n + block - 1) / block;
    for (int j = 0; j < k; ++j) {
        axpb_kernel<<<grid, block, 0, s>>>(x, n, 0.999f, 0.001f);
    }
}

static Result run_graph(cudaGraphExec_t exec, int steps, cudaStream_t s,
                        cudaEvent_t ev0, cudaEvent_t ev1) {
    CHECK(cudaStreamSynchronize(s));
    CHECK(cudaEventRecord(ev0, s));
    auto t0 = Clock::now();
    for (int i = 0; i < steps; ++i) CHECK(cudaGraphLaunch(exec, s));
    double cpu_ms = ms_since(t0);
    CHECK(cudaEventRecord(ev1, s));
    CHECK(cudaEventSynchronize(ev1));
    float gpu_ms = 0.f;
    CHECK(cudaEventElapsedTime(&gpu_ms, ev0, ev1));
    return {gpu_ms, cpu_ms};
}

        // ---- graph：capture 一个 step，实例化一次，之后反复 launch ----
        cudaGraph_t graph;
        cudaGraphExec_t exec;
        auto t0 = Clock::now();
        CHECK(cudaStreamBeginCapture(s, cudaStreamCaptureModeThreadLocal));//如果有多个线程调用cuda，这里最好写threadLocal，否则cuMemAlloc会失败
        launch_step(d_graph, n, k, s);  // 只记录，不执行
        CHECK(cudaStreamEndCapture(s, &graph));
        // WithFlags 版本在 CUDA 11.4+ 与 12.x 签名一致
        CHECK(cudaGraphInstantiateWithFlags(&exec, graph, 0));
        double inst_ms = ms_since(t0);

        CHECK(cudaGraphLaunch(exec, s));  // warmup（首次 launch 会上传图）
        CHECK(cudaMemcpyAsync(d_graph, h_init.data(), bytes, cudaMemcpyHostToDevice, s));
        Result g = run_graph(exec, steps, s, ev0, ev1);

测试数据

GPU: NVIDIA GeForce RTX 3090 (sm_86), runtime 11080, driver 13020
kernels/step = 20, steps = 1000, total kernels = 20000

         N | eager GPU ms eager CPU ms | graph GPU ms graph CPU ms  instantiate | eager us/k graph us/k |  speedup check
      1024 |       47.105       46.767 |       20.697        1.493        0.195 |       2.36       1.03 |    2.28x OK
     16384 |       35.354       34.406 |       22.034        1.591        0.036 |       1.77       1.10 |    1.60x OK
    262144 |       54.633       51.851 |       38.841        1.481        0.038 |       2.73       1.94 |    1.41x OK
   4194304 |      851.643      808.167 |      830.995        1.491        0.068 |      42.58      41.55 |    1.02x OK

可以发现 N<=1M * sizeof(float) 的数据加速会很明显

cudagraph实测时间线

可以看到cudagraph是一整块，没有加载时间，而其他kernel之间有1-2us左右的加载时间

cudagraph适用场景

适用场景：

大量小kernel

不适用场景：

• kernel本身就很大，比如prefill场景，
• shape变化大，这时shape变化了 graph就要重新计算
• 只跑几次的kernel