使用Cupy与TensorRT时CUDA流设置non_blocking=True出现同步异常

问题根因 #

核心问题是流同步时机错误，具体原因如下：

• cp.cuda.Stream(non_blocking=True)的语义是创建的CUDA流不会和CUDA默认流做隐式同步，所有提交到该流的任务（显存拷贝、kernel执行）都是完全异步的，只有显式调用synchronize()、等待事件等操作才会阻塞宿主线程等待流任务完成。而non_blocking=False的流会和默认流存在隐式同步，相当于默认流会自动等待该流的任务完成，所以你的错误写法在该模式下被隐式同步掩盖，碰巧能拿到正确结果。
• 你的代码把stream.synchronize()放在了所有输出处理逻辑之后：在异步模式下，你执行output_images = cuda_outputs[0].copy()、cp.split、cp.squeeze这些操作时，TensorRT的推理kernel还在执行，输出显存里的数据还未更新为最新的推理结果，读取到的就是残留的随机值，最终得到错误输出。

修复方案 #

调整同步指令的位置，把stream.synchronize()放在流上所有任务提交完成之后、访问输出结果之前，修改后的代码如下：

# Select stream
stream.use()
# Copy cupy array to the buffer
input_images = cp.array(batch_input_image)
cp.copyto(cuda_inputs[0], input_images)
# Run inference.
context.execute_async(bindings=bindings, stream_handle=stream.ptr, batch_size=len(batch_input_image))
# Copy results from the buffer
output_images = cuda_outputs[0].copy()
# 等待流上所有任务（输入拷贝、推理、输出拷贝）全部执行完成
stream.synchronize()
# Split results into batch
list_output = cp.split(output_images, indices_or_sections=len(batch_input_image), axis=0)
# Squeeze output arrays to remove axis of length one
list_output = [cp.squeeze(array) for array in list_output]

额外注意事项 #

• 如果你的batch_input_image是CPU侧的numpy数组，cp.array(batch_input_image)属于主机到设备的异步拷贝，需要确认主机内存是页锁定内存才能保证异步拷贝的正确性，否则会退化为同步拷贝。
• 如果你希望后续的split、squeeze等后处理操作也异步执行，可以将这些操作也提交到当前流上，只需要在最终需要把结果传回CPU或者访问结果数值前做同步即可。

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