CuPY异步流问题：为何未观测到Kernel并发执行？

我目前正在使用CuPY的RawKernels结合异步流实现大规模矩阵计算的并行化。尽管已指定非阻塞流，但每个RawKernel调用似乎都在等待前一个Kernel执行完成。请问有人知道我哪里出错了吗？以下是复现该问题的简单示例代码，创建32个流，每个流负责将3D输入数组的单个z轴切片复制到输出数组：

import cupy
kernel = cupy.RawKernel(
    '''
    extern "C" __global__ void simple_copy(float* iArr, float* oArr, int rows, int cols, int slice){
        unsigned int col = blockDim.x*blockIdx.x + threadIdx.x;
        unsigned int row = blockDim.y*blockIdx.y + threadIdx.y;
        if(row < rows && col < cols){
            //this for loop is just additional work to see kernel launches in visual profiler more easily
            for(int i=0; i<1000; i++){
                oArr[rows*cols*slice + row*cols + col] = iArr[rows*cols*slice + row*cols + col];
            }
        }
    }
    '''
    , 'simple_copy')
device = cupy.cuda.Device()
# [x, y, z]
iArr1 = cupy.ones((32*32, 32*32, 32), dtype=cupy.float32)
oArr1 = cupy.zeros((32*32, 32*32, 32), dtype=cupy.float32)
n = 32
map_streams = []
for i in range(n):
    map_streams.append(cupy.cuda.stream.Stream(non_blocking=True))
# I want to run kernel on individual z-axis slice asynchronous
for i, stream in enumerate(map_streams):
    with stream:
        kernel((32, 32), (32, 32), (iArr1, oArr1, 32*32, 32*32, i))
device.synchronize()

Hey Martin, let's dig into why your streams aren't running in parallel—this is a super common gotcha with CuPy (and CUDA overall) when working with arrays initialized in the default stream.

核心问题：默认流的隐式同步 #

你用cupy.ones()和cupy.zeros()创建iArr1和oArr1时，这些内存分配和初始化操作是在**CUDA默认流（stream 0）**中执行的。当你后续在自定义非阻塞流中启动kernel时，CuPy会自动插入默认流与自定义流之间的同步操作——这是为了确保默认流中的内存操作完成后，自定义流才能安全访问数组。这种隐式同步直接导致你的kernel只能串行执行，而非并行。

解决方法：让数组操作与自定义流解绑 #

要解决这个问题，你需要避免数组初始化操作强制触发默认流与自定义流的同步。这里有两种可行方案：

1. 异步初始化数组

放弃使用默认的cupy.ones()/cupy.zeros()，先分配未初始化的内存，再通过异步操作填充数据，全程在独立流中完成，避免默认流的干扰：

import cupy

kernel = cupy.RawKernel(
    '''
    extern "C" __global__ void simple_copy(float* iArr, float* oArr, int rows, int cols, int slice){
        unsigned int col = blockDim.x*blockIdx.x + threadIdx.x;
        unsigned int row = blockDim.y*blockIdx.y + threadIdx.y;
        if(row < rows && col < cols){
            for(int i=0; i<1000; i++){
                oArr[rows*cols*slice + row*cols + col] = iArr[rows*cols*slice + row*cols + col];
            }
        }
    }
    '''
    , 'simple_copy')

device = cupy.cuda.Device()
rows = 32*32
cols = 32*32
depth = 32

# 先分配未初始化内存（不会触发默认流的同步操作）
iArr1 = cupy.empty((rows, cols, depth), dtype=cupy.float32)
oArr1 = cupy.empty((rows, cols, depth), dtype=cupy.float32)

# 用单独的流异步初始化数组
init_stream = cupy.cuda.stream.Stream(non_blocking=True)
with init_stream:
    iArr1[:] = 1.0
    oArr1[:] = 0.0

# 创建自定义流
map_streams = []
for i in range(depth):
    map_streams.append(cupy.cuda.stream.Stream(non_blocking=True))

# 显式等待初始化完成，再启动kernel
init_stream.synchronize()

# 并行启动kernel
for i, stream in enumerate(map_streams):
    with stream:
        kernel((32, 32), (32, 32), (iArr1, oArr1, rows, cols, i))

# 等待所有kernel执行完毕
device.synchronize()

2. 显式管理流依赖（保留默认流初始化的情况）

如果必须在默认流中初始化数组，可以手动添加流之间的依赖关系，避免隐式同步导致的串行：

# 在创建自定义流后，添加依赖：让自定义流等待默认流完成数组初始化
for stream in map_streams:
    stream.wait_on_stream(cupy.cuda.stream.get_default_stream())

# 之后再按原代码启动kernel即可

额外检查点 #

• GPU资源限制: 即使流设置正确，如果你的GPU没有足够的SM（流式多处理器）同时运行32个kernel，部分kernel会进入等待队列，但修复后它们会在资源可用时尽可能重叠执行。
• kernel负载: 你添加的for循环很好，它增加了kernel的计算量，能让你在性能分析工具中更清晰地看到并行效果，避免被kernel启动开销掩盖。

用NVIDIA Nsight Systems等工具分析时，你应该能看到kernel并行启动的效果了。

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