Torch CUDA 异步执行

一次 forward 的 Python 调用已经返回，GPU 算完了吗？为什么连续两行 y = f(x)、z = g(y) 之间不需要加同步？
代码里只多了一行 .item()，吞吐为什么明显下降？
用 time.perf_counter() 包住一次 matmul 只有几微秒，服务延迟却很高——到底测到了什么？

这些问题都来自同一个事实：CUDA 调用的“返回”“排定先后”和“真正完成”是三个不同的时刻。 本篇先建立这个模型，再用它解释同步、跨 stream 依赖、异步搬运、计时和报错位置，最后看 SGLang 的 overlap scheduler 如何利用它。本篇以 CUDA 后端为例；其他设备后端的同步规则需要分别确认。

抽象：CPU 是生产者，stream 是队列

核心的一句话：CPU 不执行 GPU 计算，它只把命令排进队列。 大多数 CUDA 运算在 CPU 侧只是“提交”：把 kernel 和参数放进某个 stream，然后立刻返回，Python 接着往下跑。GPU 从 stream 里按顺序取出工作执行。

于是一个 CUDA 调用有三个需要分开的时刻：

提交（submitted）：Python 调用返回。此时只能说“这项工作已经排上队”。
排序（ordered）：同一 stream 中的工作按提交顺序执行，所以 z = g(y) 天然排在 y = f(x) 之后，两行之间无需同步。不同 stream 之间没有这种顺序，依赖需要显式建立。
完成（completed）：GPU 真正算完，结果可以被主机读取。

这个模型下，性能的关键是让 CPU 始终领先 GPU：只要 Python 不需要知道数据里的值，CPU 就能持续往队列里放工作，GPU 一直有活干。一旦 CPU 需要某个值，它必须停下等待“完成”时刻；等它回来继续提交时，GPU 队列可能已经空了，出现空泡。总体执行语义见 CUDA semantics。

由此得到统一判据：这一行代码只需要“排序”，还是需要“完成”？ 只需要排序的依赖，应该交给 GPU 队列表达；只有真正要读值时，才让 CPU 等待完成。

混淆点一：GPU 等 GPU，还是 CPU 等 GPU

创建多个 stream 不等于让有依赖的计算自动并行；反过来，不同 stream 上有依赖的工作，也不会自动排好先后。建立依赖有两类手段，等待的主体不同：

手段	谁等待	用途
`stream.wait_stream(other)` / `stream.wait_event(event)`	消费 stream 中后续工作等待	建立 GPU 工作之间的依赖，CPU 通常继续执行
`event.synchronize()`	CPU 等待指定事件完成	CPU 即将读取结果
`torch.cuda.synchronize()`	CPU 等待指定设备上所有 stream 完成	调试、整体计时，范围较大

前一类把依赖写进 GPU 队列：调用立即返回，只有消费 stream 里排在其后的工作被推迟。后一类让 CPU 线程阻塞。设计跨 stream 流程时，应尽量用前者表达“谁在谁之后”，把后者留给真正要读值的一处。Event 的记录、等待和计时语义见 torch.cuda.Event。

混淆点二：隐式同步——Python 想知道数据里的值

x.shape、x.stride()、x.dtype 等元数据一般可在 CPU 侧直接获取，不需要等待 GPU。CUDA 张量的值则需要等待设备计算完成并取回主机，例如：

CUDA 标量的 .item()、bool(tensor)；
用 CUDA 标量驱动 Python 的 if；
.tolist() 或打印 CUDA 张量的具体值；
默认阻塞的 GPU 到 CPU 拷贝。

next_token = logits.argmax(dim=-1)
# 若下一步仍在 GPU 使用 next_token，这里可以继续排队。
# 若调用 next_token[0].item()，CPU 就需要等待这个值。

这些写法在代码里看不到“synchronize”字样，却等价于让 CPU 停在“完成”时刻，即上面时间线图的 B 情形。采样完成后向客户端发送 token，终究需要访问结果；优化重点是减少不必要的读取、批量搬运，并把同步放在真正需要结果的位置。不要在每个中间层为了日志读取 CUDA 标量。

默认 .numpy() 不能直接转换 CUDA 张量，通常先调用 .cpu()；同步发生在取回数据的路径，而不是“任何 .numpy() 都会同步 GPU”。

混淆点三：non_blocking 与计算重叠是两件事

non_blocking=True 请求让搬运尽可能不阻塞 CPU，解决的是“CPU 是否等待”；搬运能否与 GPU 计算重叠，是另一个问题，取决于 pinned memory、copy engine、stream 安排以及两项工作之间是否有数据依赖。

即使源张量已经 pinned，如果拷贝和计算都放在同一个 stream，后续计算仍按顺序等待拷贝——这正是“排序”规则。为了重叠，通常要把独立工作放在不同 stream，并建立正确的消费依赖。临时调用 pin_memory() 本身需要开销，不一定比直接拷贝更快。见官方 pinned memory 与 non_blocking 教程。

异步传输把“提交”和“完成”拉开了距离，也就要求在这段时间里遵守缓冲区生命周期：

pinned CPU 源缓冲区在 H2D 完成前不能被覆盖；
异步 D2H 的 CPU 目标缓冲区，在拷贝完成前不能读取；
跨 stream 使用 CUDA 张量，除了计算依赖，还要保证分配器不会过早复用其存储。相关概念见显存管理。

混淆点四：计时测的是哪个时刻

只用 time.perf_counter() 包围异步 CUDA 调用，停表时只到了“提交”时刻，测到的可能主要是 CPU 提交时间。CUDA Event 则是排进 stream 的标记，GPU 执行到它时才记录时间，因此能度量设备上一段工作。

import torch

x = torch.randn(1024, 1024, device="cuda")
for _ in range(5):
    y = x @ x  # 预热
torch.cuda.synchronize()

start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
y = x @ x
end.record()
end.synchronize()
elapsed_ms = start.elapsed_time(end)

这测量两个事件之间的设备时间，不直接代表包含请求排队、CPU 调度、数据处理和网络返回的服务延迟。跨 stream 的工作要纳入事件依赖，不能假定一对事件自动覆盖所有 stream。

测端到端耗时可在测量区间前后同步，再使用 CPU 时钟；但不要把用于测量的同步长期留在生产热路径。torch.profiler 的 CPU/CUDA trace 可以把两条时间线并排画出来，帮助区分提交间隙、搬运、kernel 执行与同步等待。见 PyTorch Profiler。

混淆点五：异步错误为什么出现在后一行

报错同样遵循“完成”时刻：非法设备内存访问可能在后续同步或另一个 CUDA 调用时才报告，因此 Python 堆栈不一定指向最初出错的 kernel。调试时可设置 CUDA_LAUNCH_BLOCKING=1，让执行位置与报错位置更容易对应；它会改变执行时序，不能用于代表真实吞吐的测试。见 CUDA 环境变量。

在 SGLang 中：overlap scheduler 把同步推迟一拍

推理服务每个 decode step 都必须把新 token 取回 CPU，用来判断停止条件、更新请求状态、发给客户端——这个同步躲不掉。SGLang 的 overlap scheduler 做的事，是让这次必要的等待发生在下一批已经入队之后，从而 CPU 在等，GPU 不空转。上面“两种等待”图描绘的正是这条流水线。

主循环是 python/sglang/srt/managers/scheduler.py 中的 Scheduler.event_loop_overlap：先 run_batch() 提交当前批次，把 (batch, result) 放进 result_queue，再弹出上一批的结果处理。几个与本篇对应的细节：

用 stream 依赖代替 CPU 同步。 调度线程的 GPU 工作在 schedule_stream 上；run_batch() 里先 forward_stream.wait_stream(self.schedule_stream) 再执行 forward，结果拷贝前又让 copy_stream.wait_stream(self.forward_stream)。这些都只是“排序”，CPU 不停。代码注释说明了动机：结果 D2H 放在 copy_stream 上，可与下一次 forward 重叠，而不是在 forward_stream 上串行。
异步 D2H 同时处理目标与源的生命周期。 managers/utils.py 中的 _async_d2h() 在 CUDA 上先分配 pinned 的 CPU 目标张量，再 copy_(t, non_blocking=True)，并对源张量调用 record_stream()，防止缓存分配器在拷贝 stream 完成前复用源存储。GenerationBatchResult.copy_to_cpu() 用它搬运 next_token_ids、logprob 等，最后 copy_done.record()。
唯一的硬等待放在读值之前。 managers/scheduler_components/batch_result_processor.py 中，SchedulerBatchResultProcessor.process_batch_result_decode() 先 result.copy_done.synchronize()，之后 _normalize_decode_outputs() 才对 next_token_ids 调用 .tolist()。由于此时读的是已同步的 CPU pinned 张量，.tolist() 不再触发设备同步。

# scheduler.py, run_batch()（节选、简化）
self.copy_stream.wait_stream(self.forward_stream)   # GPU 等 GPU
with self.copy_stream_ctx:
    batch_result.copy_to_cpu(...)                   # pinned D2H + copy_done.record()
# 下一轮循环处理该结果时：
result.copy_done.synchronize()                      # CPU 等 GPU，只在这里

这一设计也有边界：连续两个 prefill 批次时，is_disable_overlap_for_batch() 会关闭重叠以改善首批的 TTFT；ROCm（_is_hip）路径则直接在 forward_stream 上拷贝，注释给出的理由是跨 stream 同步的代价高于它所重叠的小型 D2H。重叠是否值得，仍取决于具体硬件和测量。

小结

把 CUDA 执行看成“CPU 往 stream 队列里提交命令、GPU 按序消费”，就能把提交、排序、完成三个时刻分开：同 stream 的先后天然成立，跨 stream 的先后用 wait_stream/wait_event 写进队列，只有 Python 真正需要值时才让 CPU 等待完成。隐式同步、non_blocking 的边界、计时读数和延迟报错，都是这三个时刻被混淆的不同表现。读推理热路径时依次问：这里需要排序还是需要完成？谁在等谁？异步期间缓冲区由谁保活？再用 profiler 的双时间线验证。