捕获CUDA Graph与循环异步内存分配时的错误排查

我正在尝试实现一个CUDA Graph实验，包含三个按顺序执行且存在依赖关系的kernel：kernel_0、kernel_1和kernel_2，目前仅对kernel_1相关逻辑进行捕获。但使用compute-sanitizer运行时，触发了CUDA_ERROR_INVALID_VALUE错误，报错点在cuMemFreeAsync调用。

实验代码 #

#include <stdio.h>
#include <chrono>
#include <cuda.h>
#include <cuda_runtime.h>
#include <iostream>

#define N 50000
#define NSTEP 1000
#define NKERNEL 20

using namespace std::chrono;

static const char *_cudaGetErrorEnum(cudaError_t error) {
  return cudaGetErrorName(error);
}

template <typename T>
void check(T result, char const *const func, const char *const file,
           int const line) {
  if (result) {
    fprintf(stderr, "CUDA error at %s:%d code=%d(%s) \"%s\" \n", file, line,
            static_cast<unsigned int>(result), _cudaGetErrorEnum(result), func);
    exit(EXIT_FAILURE);
  }
}

#define checkCudaErrors(val) check((val), #val, __FILE__, __LINE__)

__global__ void shortKernel_0(float * out_d, float * in_d){
    int idx=blockIdx.x*blockDim.x+threadIdx.x;
    if(idx<N) {
        in_d[idx] = 1.0;
        out_d[idx]=1 + in_d[idx];
    }
}

__global__ void shortKernel_1(float * out_d, float * in_d){
    int idx=blockIdx.x*blockDim.x+threadIdx.x;
    if(idx<N) out_d[idx]=2*in_d[idx];
}

__global__ void shortKernel_2(float * out_d, float * in_d){
    int idx=blockIdx.x*blockDim.x+threadIdx.x;
    if(idx<N) {
        out_d[idx]=3*in_d[idx];
    }
}

void test(){

    size_t size_bytes = N * sizeof(float);
    void * in_d_0;
    void * out_d_0;
    void * out_d_1;
    void * out_d_2;

    int threads = 128;
    int blocks = (N+threads)/threads;
    int iter = 10;
    cudaStream_t stream;
    cudaStreamCreate(&stream);

    CUmemoryPool pool_;
    cuDeviceGetDefaultMemPool(&pool_, 0);
    uint64_t threshold = UINT64_MAX;
    cuMemPoolSetAttribute(pool_, CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, &threshold);

    cudaGraph_t graph;
    cudaGraphExec_t instance;
    bool graphCreated=false;

    for (int i =0; i < iter; i++){
        cuMemAllocFromPoolAsync(reinterpret_cast<CUdeviceptr*>(&in_d_0), size_bytes, pool_, stream);
        cuMemAllocFromPoolAsync(reinterpret_cast<CUdeviceptr*>(&out_d_0), size_bytes, pool_, stream);
        shortKernel_0<<<blocks, threads,0, stream>>>(reinterpret_cast<float *>(out_d_0), reinterpret_cast<float *>(in_d_0));
        if (!graphCreated){
            cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
            cuMemAllocFromPoolAsync(reinterpret_cast<CUdeviceptr*>(&out_d_1), size_bytes, pool_, stream);
            cuMemFreeAsync(reinterpret_cast<const CUdeviceptr&>(in_d_0), stream);
            shortKernel_1<<<blocks, threads,0, stream>>>(reinterpret_cast<float *>(out_d_1), reinterpret_cast<float *>(out_d_0));
            cudaStreamEndCapture(stream, &graph);
            checkCudaErrors(cudaGraphInstantiate(&instance, graph, NULL, NULL, 0));
            checkCudaErrors(cudaGraphUpload(instance, stream));
            graphCreated = true;
        }else{
            checkCudaErrors(cudaGraphLaunch(instance, stream));
        }
        cuMemFreeAsync(reinterpret_cast<const CUdeviceptr&>(out_d_0), stream);
        cuMemAllocFromPoolAsync(reinterpret_cast<CUdeviceptr*>(&out_d_2), size_bytes, pool_, stream);
        shortKernel_2<<<blocks, threads,0, stream>>>(reinterpret_cast<float *>(out_d_2), reinterpret_cast<float *>(out_d_1));
       
        cuMemFreeAsync(reinterpret_cast<const CUdeviceptr&>(out_d_1), stream);
        cuMemFreeAsync(reinterpret_cast<const CUdeviceptr&>(out_d_2), stream);
    }
   
    cudaDeviceSynchronize();        
    printf("With async malloc done!");
    cudaStreamDestroy(stream);
    cudaGraphDestroy(graph);
    cudaGraphExecDestroy(instance);
}

int main() {
    test();
    return 0;
}

错误信息 #

========= Program hit CUDA_ERROR_INVALID_VALUE (error 1) due to "invalid argument" on CUDA API call to cuMemFreeAsync.
=========     Saved host backtrace up to driver entry point at error
=========     Host Frame: [0x2ef045]
=========                in /usr/local/cuda/compat/lib.real/libcuda.so.1
=========     Host Frame:test() [0xb221]
=========                in /opt/test-cudagraph/./a.out
=========     Host Frame:main [0xb4b3]
=========                in /opt/test-cudagraph/./a.out
=========     Host Frame:__libc_start_main [0x24083]
=========                in /usr/lib/x86_64-linux-gnu/libc.so.6
=========     Host Frame:_start [0xaf6e]
=========                in /opt/test-cudagraph/./a.out

1. 指针转换错误导致无效参数
   代码中cuMemFreeAsync调用时使用了错误的指针转换：reinterpret_cast<const CUdeviceptr&>(in_d_0)。in_d_0是void*类型，这里的引用转换会导致传递给cuMemFreeAsync的不是设备指针本身，而是void*指针的引用地址，完全不符合API要求，直接触发invalid argument错误。正确的转换应该是reinterpret_cast<CUdeviceptr>(in_d_0)，直接将void*转成CUdeviceptr类型。

2. 图内分配的内存被图外重复释放
   out_d_1是在流捕获阶段分配的，属于CUDA Graph的一部分。每次启动图时，Graph会自动处理内部的内存分配逻辑，但当前代码在图外的循环中执行cuMemFreeAsync(out_d_1, stream)：

   • 第一次迭代时out_d_1有效，但后续迭代中图启动时会重新分配新的out_d_1指针，原指针已失效，此时图外的释放操作尝试处理旧的无效指针，引发参数错误。
   • 正确做法是将out_d_1的释放操作也纳入Graph管理，或者让Graph完全控制其生命周期，图外不再手动释放。

3. 指针生命周期管理混乱
   后续迭代中graphCreated为true时，代码直接启动图，但out_d_1并未被重新赋值为图中新分配的指针，此时shortKernel_2使用的是第一次迭代的旧指针，图外的释放操作也针对该旧指针，不仅会触发释放错误，还会导致kernel访问无效内存。

• 修正指针转换：将所有reinterpret_cast<const CUdeviceptr&>(xxx)改为reinterpret_cast<CUdeviceptr>(xxx)，确保传递正确的设备指针给Driver API。
• 将内存释放纳入Graph管理：在流捕获阶段，添加cuMemFreeAsync释放out_d_1的操作到Graph中，让Graph完全管理该内存的生命周期。
• 调整迭代逻辑：后续迭代中不要直接复用第一次捕获的out_d_1指针，需通过Graph节点获取新指针，或调整数据依赖让kernel_2能正确访问Graph内kernel_1的输出。

内容的提问来源于stack exchange，提问作者kingwales