如何在单个GPU上利用CUDA Streams实现模型推理的异步计算与性能优化？

在现代AI基础设施中，优化推理延迟和提高GPU利用率是核心挑战。即使在单个GPU上，如果不进行适当的调度，许多操作（如数据传输和计算）也会串行执行，导致计算资源闲置。CUDA Streams是解决这一问题的关键技术，它允许我们在GPU上调度并发的、异步的任务序列，实现数据传输和计算的重叠（Overlap）。

本文将通过一个PyTorch示例，演示如何在单个GPU上创建并使用非默认CUDA Streams，以实现两个独立任务的并行处理和性能提升。

1. 为什么需要CUDA Streams？

默认情况下，所有的CUDA操作（包括

和内核启动）都在“默认流”（Stream 0）中执行。Stream 0是一个隐式的同步流，这意味着GPU在执行流0中的下一个操作前，必须等待前一个操作完成。这使得Host-to-Device (H2D) 数据传输、内核执行 (Kernel) 和 Device-to-Host (D2H) 数据传输无法并行。

CUDA Streams 是GPU上的任务队列。创建多个非默认流允许GPU调度器将属于不同流的任务视为独立的、可以并发执行的任务，从而实现H2D、Kernel和D2H操作的重叠。

2. 实践：利用PyTorch Streams实现异步重叠

我们将模拟两个独立的模型推理任务A和B。每个任务包含数据上传、计算（矩阵乘法）和数据下载。

环境准备

确保安装了PyTorch和CUDA环境。

1
pip install torch

2.1 基准测试：串行执行 (使用默认流)

我们首先运行基准测试。即使我们连续启动两个任务，由于它们都在默认流上，它们也会依次执行。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import torch

import time



# 设置设备

device = torch.device('cuda')



# 模拟数据大小 (例如，一个中等大小的批量输入)

MATRIX_SIZE = 4096



# 模拟计算任务：矩阵乘法

def simulate_task(data, matrix_size):

    # 1. Host -> Device (H2D)

    input_tensor_device = data.to(device)



    # 2. Kernel Execution (Computation)

    M = torch.randn(matrix_size, matrix_size, device=device)

    result = torch.matmul(input_tensor_device, M)



    # 3. Device -> Host (D2H)

    result_host = result.cpu()

    return result_host



# 创建两个独立的数据块

data_A = torch.randn(MATRIX_SIZE, MATRIX_SIZE)

data_B = torch.randn(MATRIX_SIZE, MATRIX_SIZE)



# --- 串行执行 --- 

torch.cuda.synchronize() 

start_time = time.time()



# 任务 A

result_A = simulate_task(data_A, MATRIX_SIZE)



# 任务 B

result_B = simulate_task(data_B, MATRIX_SIZE)



torch.cuda.synchronize()

end_time = time.time()



print(f"串行执行总耗时: {end_time - start_time:.4f} 秒")

# 预期输出: 串行执行总耗时: ~1.2000 秒 (取决于硬件)

2.2 优化方案：使用CUDA Streams实现异步重叠

现在我们创建两个非默认流

和

，并将任务A的操作分配给前者，任务B的操作分配给后者。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
# 创建两个非默认流

stream_A = torch.cuda.Stream()

stream_B = torch.cuda.Stream()



# 注意：PyTorch操作必须显式地被放入流的上下文管理器中



def simulate_task_streamed(data, matrix_size, stream):

    # 将数据和计算分配给指定的流

    with torch.cuda.stream(stream):

        # 1. H2D (异步)

        input_tensor_device = data.to(device)



        # 2. Kernel Execution (异步)

        M = torch.randn(matrix_size, matrix_size, device=device)

        result = torch.matmul(input_tensor_device, M)



        # 3. D2H (异步)

        result_host = result.cpu()



        # 返回设备上的结果引用，等待同步后才能安全访问

        return result_host



# --- 异步执行 --- 

torch.cuda.synchronize() 

start_time = time.time()



# 启动任务 A (在 stream_A 上)

result_A_ref = simulate_task_streamed(data_A, MATRIX_SIZE, stream_A)



# 启动任务 B (在 stream_B 上)

result_B_ref = simulate_task_streamed(data_B, MATRIX_SIZE, stream_B)



# 必须等待所有流完成，确保D2H操作确实完成，结果安全地在CPU内存中

torch.cuda.synchronize()

end_time = time.time()



print(f"异步执行总耗时: {end_time - start_time:.4f} 秒")

# 预期输出: 异步执行总耗时: ~0.6500 秒 (接近单个任务的耗时，因为大部分时间重叠了)

3. 性能分析与结论

通过使用CUDA Streams，我们成功地将两个独立任务的H2D、Kernel和D2H操作进行重叠。如果假设两个任务耗时相同（T），那么异步执行的总耗时将接近T，而不是2T，显著提升了吞吐量和降低了端到端延迟。

4. 关键注意事项

1. 流同步 (

   1	torch.cuda.synchronize()

   ): 这是一个全局同步点，它会等待所有流上的所有操作完成。在实际部署中，通常使用

   1	stream.synchronize()

   来仅同步特定的流，避免不必要的等待。

2. 内存分配: 为了在不同流之间安全地共享Host内存或在异步操作中重用内存，需要使用固定内存 (Pinned Memory/Page-Locked Memory)。在PyTorch中，可以通过

   1	tensor.pin_memory()

   来实现，这样H2D传输可以与计算并行进行。

3. 依赖性: 如果任务B需要任务A的输出，则必须使用事件 (

   1	torch.cuda.Event

   ) 或

   1	stream.wait_stream(other_stream)

   来强制同步，确保依赖关系满足。

在模型部署场景中，CUDA Streams通常用于处理多batch请求的并行推理，或在推理核心计算的同时，并行处理下一批数据的CPU预处理和GPU数据上传。