PyTorch多流编程陷阱：record_stream与wait_event缺一不可

1. 一个让无数人栽跟头的显存回收陷阱

如果你写过 PyTorch 的自定义算子，或者用过 CUDA Graph、多流并行推理，大概率见过record_stream这个 API。它看起来人畜无害，文档里就一句话：告诉分配器这块显存还被某个 stream 用着，别急着回收。但真正踩过坑的人都知道，record_stream和多 stream 配合的时候，如果漏掉了wait_event，程序不会立刻崩，而是会在某个你完全意想不到的时刻，给你一个莫名其妙的数值错误，或者一个illegal memory access。

我最近就在一个多流推理的项目里栽了这个坑。现象很诡异：单流跑完全正常，一开多流，前几百个 batch 结果都对，跑到某个随机位置开始输出乱码，而且每次乱码的位置还不一样。用compute-sanitizer跑，报的是 use-after-free，但代码里明明已经调用了record_stream。折腾了整整两天，最后定位到问题：record_stream只解决了"分配器不要提前复用这块显存"的问题，但它不解决 stream 之间的执行顺序问题。换句话说，record_stream保证的是生命周期，wait_event保证的是可见性和顺序，这两件事必须一起做，缺一不可。

这篇就把这个坑从头到尾拆一遍。不管你是刚接触 CUDA 多流编程的新手，还是已经写过不少自定义算子但总感觉多流下"不太稳"的老手，应该都能从里面找到点有用的东西。我会先讲清楚 PyTorch 的 caching allocator 到底怎么工作、record_stream到底做了什么、为什么它不够，然后给出正确的写法，最后把我踩过的几个变种坑和排查手法一并整理出来。

2. 先把 caching allocator 和 stream 的关系捋清楚

2.1 PyTorch 的显存分配器为什么这么快

要理解record_stream，得先理解 PyTorch 的 CUDA caching allocator。原生cudaMalloc和cudaFree是同步操作，而且开销很大，动辄几百微秒。如果一个训练 step 里要反复申请释放几十上百块显存，光分配开销就能把 GPU 利用率拖垮。PyTorch 的做法是：一次性向驱动申请一大块显存，然后在用户态自己维护一个空闲块池。你调用torch.empty的时候，它从池子里挑一块合适的给你；你释放的时候，它不还给驱动，而是标记为空闲，放回池子。

这个设计的关键在于"标记为空闲"这个动作。它默认认为：只要 Python 侧不再引用这块 tensor，GPU 上就不会再有人用它了。这个假设在单流场景下是成立的，因为所有 kernel 都在同一个 stream 上按顺序执行，Python 侧释放 tensor 的时候，前面用到它的 kernel 早就执行完了（或者至少排在同一个队列里，不会乱序）。

但多流场景下这个假设直接崩了。你在 stream A 上跑了一个 kernel 用了 tensor X，然后 Python 侧把 X 释放了，allocator 立刻把这块显存分给了 stream B 上的另一个 tensor Y。问题是 stream A 和 stream B 是并发的，A 上的 kernel 可能还没执行完，B 上的 kernel 就已经开始往同一块显存里写数据了。这就是经典的 use-after-free。

2.2 record_stream 到底做了什么

record_stream就是用来堵这个洞的。当你调用tensor.record_stream(s)的时候，allocator 会在这块显存的元数据里记一笔：stream s 还在用这块显存。等到这块显存被释放回池子的时候，allocator 不会立刻把它分给别人，而是先检查所有记录过的 stream，确认这些 stream 上的工作都完成了，才真正放回可用池。

具体实现上，allocator 用的是 CUDA event。它会在记录的 stream 上插入一个 event，然后延迟回收，直到这个 event 被 query 到已经完成。所以record_stream本质上是一个"延迟回收"机制，它保证的是：这块显存不会在你记录的 stream 完成之前被复用。

听起来很完美对吧？问题就出在这里。record_stream保证的是"不被复用"，但它不保证你的 kernel 能看到正确的数据。这两件事听起来很像，其实是两个完全不同的问题。

2.3 生命周期和可见性是两码事

打个比方。你有一个共享的白板，A 同事在上面写数据，B 同事要读。record_stream相当于告诉管理员："这块白板 A 还在用，别擦。"它保证白板不会被擦掉。但是，如果 B 在 A 写完之前就去读，读到的还是旧数据或者半截数据。record_stream管不了这个，它只保证白板还在，不保证白板上的内容已经写好了。

要让 B 读到正确的数据，你需要的是同步：让 B 等到 A 写完再读。在 CUDA 里，跨 stream 的同步靠的是 event：在 stream A 上记录一个 event，然后让 stream B 等待这个 event。这就是wait_event的作用。

所以正确的多流写法是：既要record_stream保证显存不被提前复用，又要wait_event保证数据依赖的执行顺序。只做前者，你会遇到偶发的数据错误；只做后者，你会遇到 use-after-free。两个都做，才是安全的。

3. 一个最小可复现的翻车案例

3.1 场景设定

我构造一个最简单的场景来说明问题。假设我们有两个 stream：producer_stream负责往一块显存里写数据，consumer_stream负责读这块数据做后续计算。中间通过一个 Python 侧的 tensor 传递。

import torch producer_stream = torch.cuda.Stream() consumer_stream = torch.cuda.Stream() def buggy_pipeline(): # 在 producer stream 上分配并写入 with torch.cuda.stream(producer_stream): x = torch.empty(1024, 1024, device='cuda') x.fill_(1.0) # 记录 stream，防止显存被提前复用 x.record_stream(producer_stream) # 在 consumer stream 上读取 with torch.cuda.stream(consumer_stream): y = x * 2.0 # 这里也记录一下 x.record_stream(consumer_stream) return y

这段代码看起来"很规范"：两个 stream 都调用了record_stream。但它有一个致命问题：consumer_stream上的x * 2.0可能在producer_stream上的fill_完成之前就开始执行了。因为两个 stream 之间没有任何同步，GPU 调度器完全可能先跑 consumer 的 kernel。

结果就是y里可能有一部分是 2.0（读到了正确的 1.0），一部分是 0.0（读到了未初始化的显存），甚至可能是垃圾值。而且这个错误是非确定性的，取决于两个 stream 的实际调度顺序，跑一百次可能只有几次出错。

3.2 为什么 record_stream 救不了你

很多人会想：我不是调了record_stream吗？allocator 应该知道这块显存还在用啊。问题在于，record_stream记录的是"这块显存被哪些 stream 用过"，allocator 据此延迟回收。但在这个例子里，x根本没有被释放，它一直活到函数返回。所以record_stream在这里其实没起到任何作用，它防的是"x 被释放后显存被复用"，而这里的问题是"x 还没被释放，但 consumer 读得太早"。

这就是最迷惑人的地方：你调了record_stream，代码看起来没问题，但错误照出。因为record_stream和wait_event解决的是两个正交的问题，前者管生命周期，后者管执行顺序。

3.3 正确写法

修法很简单，在 consumer 读之前，等 producer 写完：

def correct_pipeline(): with torch.cuda.stream(producer_stream): x = torch.empty(1024, 1024, device='cuda') x.fill_(1.0) x.record_stream(producer_stream) # 在 producer stream 上记录一个 event ev = torch.cuda.Event() ev.record(producer_stream) with torch.cuda.stream(consumer_stream): # consumer 等待 producer 完成 consumer_stream.wait_event(ev) y = x * 2.0 x.record_stream(consumer_stream) return y

关键就是那两行：ev.record(producer_stream)和consumer_stream.wait_event(ev)。event 在 producer stream 上记录，表示"这个点之前的所有工作都完成了"；consumer stream 等待这个 event，表示"在 event 完成之前，我不执行后面的 kernel"。这样就建立了跨 stream 的依赖关系。

注意：wait_event是让当前 stream 等待 event，不是让 event 等待 stream。方向别搞反了。consumer_stream.wait_event(ev)的意思是 consumer 等 producer，不是反过来。

4. 深入 allocator：record_stream 的延迟回收机制

4.1 延迟回收是怎么实现的

前面提到 allocator 用 event 来实现延迟回收，这里展开讲一下，因为理解了这个机制，你才能明白为什么有些写法"看起来多余但必须写"。

当你调用tensor.record_stream(s)时，allocator 会做两件事：第一，把 stream s 加入这块显存的"使用中 stream 集合"；第二，在 stream s 上插入一个 event。当这块显存被释放（Python 引用计数归零）时，allocator 不会立刻把它放回空闲池，而是把它放进一个"待回收"队列，同时检查所有记录的 event。只有当所有 event 都 query 到完成状态，这块显存才真正回到空闲池。

这个机制有个重要含义：record_stream的调用时机必须在显存被释放之前。如果你在 tensor 已经被释放之后才想起来调record_stream，那已经晚了，allocator 可能已经把显存分给别人了。所以正确的做法是在使用这块显存的 stream 上下文里，尽早调用record_stream。

4.2 为什么有时候不调 record_stream 也没事

很多人会问：我写多流代码从来没调过record_stream，也没出过问题啊。这有几种可能。

第一种，你的 tensor 生命周期足够长，长到所有 stream 都用完了它还没被释放。比如你在一个长函数里持有 tensor 引用，函数返回前所有 stream 的工作都做完了，那确实不需要record_stream。但这是靠"运气"和"代码结构"保证的，一旦有人重构代码，把 tensor 的释放提前了，坑就埋下了。

第二种，你的多流其实是"伪多流"。比如你用了torch.cuda.stream()上下文，但里面全是同步操作，或者你每次都在 stream 之间调了torch.cuda.synchronize()。这种情况下所有 stream 实际上被串行化了，自然不会有并发问题，但也就失去了多流的意义。

第三种，你的显存足够大，allocator 从来没触发过复用。比如你有一块 80G 的卡，跑一个只用 10G 的任务，allocator 池子里永远有空闲块，不会去复用正在被使用的显存。这种情况下record_stream缺失不会暴露问题，但一旦你把 batch size 调大，或者换到小显存的卡上，问题立刻出现。

实操心得：判断你的代码是否真的需要record_stream，最可靠的方法是开PYTORCH_NO_CUDA_MEMORY_CACHING=1环境变量跑一遍。这个变量会禁用 caching allocator，每次分配都走原生cudaMalloc，虽然慢，但能暴露所有 use-after-free 问题。如果禁用缓存后程序崩了，说明你的代码有生命周期问题。

4.3 record_stream 的性能代价

record_stream不是免费的。每次调用都会插入一个 event，event 的创建和 query 都有开销。在高频调用的路径上（比如每个 kernel 都调），累积起来可能相当可观。我实测过一个场景，在每层都调record_stream的情况下，端到端延迟增加了大约 3% 到 5%。

所以实践中要权衡：不是所有 tensor 都需要record_stream，只有那些"跨 stream 使用且生命周期可能短于 stream 执行时间"的 tensor 才需要。比如一个中间激活值，在 stream A 上产生，在 stream B 上消费，然后 Python 侧立刻释放，这种就必须调。而一个贯穿整个 forward 的输入 tensor，生命周期很长，通常不需要。

5. 多流编程的正确姿势与完整实操

5.1 建立依赖的三板斧

多流编程里，建立正确的依赖关系靠三样东西：event、wait_event、record_stream。我把它们的分工整理成一张表，方便对照。

这三者的关系是：record_stream管"显存还在不在"，event+wait_event管"数据对不对"。缺任何一个，多流代码都不安全。

5.2 一个完整的多流推理 pipeline

下面给一个更贴近实际的例子：一个两阶段的推理 pipeline，stage1 在 stream A 上做预处理，stage2 在 stream B 上做模型推理，中间通过一个中间 tensor 传递。

import torch class MultiStreamPipeline: def __init__(self): self.stream_a = torch.cuda.Stream() self.stream_b = torch.cuda.Stream() self.events = [torch.cuda.Event() for _ in range(4)] def run(self, input_tensor): # Stage 1: 预处理，在 stream A 上 with torch.cuda.stream(self.stream_a): # 分配中间 tensor mid = torch.empty_like(input_tensor) mid.copy_(input_tensor) mid.mul_(2.0) # 记录 stream，防止 mid 被提前复用 mid.record_stream(self.stream_a) # 记录 event，标记 stage1 完成点 self.events[0].record(self.stream_a) # Stage 2: 推理，在 stream B 上 with torch.cuda.stream(self.stream_b): # 等待 stage1 完成 self.stream_b.wait_event(self.events[0]) # 消费 mid out = mid + 1.0 mid.record_stream(self.stream_b) self.events[1].record(self.stream_b) # 回到默认 stream，等待 stage2 torch.cuda.current_stream().wait_event(self.events[1]) return out

这段代码里有几个细节值得说。

第一，mid.record_stream(self.stream_a)和mid.record_stream(self.stream_b)都调了。因为mid在两个 stream 上都被使用，两个 stream 都要记录。少记一个，allocator 就可能在那一个 stream 还没用完的时候回收显存。

第二，event 的复用。我用了一个 event 列表，而不是每次创建新 event。event 的创建有开销，在热路径上复用 event 能省不少。但要注意，event 被record之后，如果还没被wait就再次record，行为是未定义的。所以复用 event 的前提是确保上一次的 wait 已经完成。

第三，最后的torch.cuda.current_stream().wait_event(self.events[1])。这一步很多人会漏。如果out要在默认 stream 上被后续代码使用，默认 stream 必须等待 stream B 完成。否则默认 stream 上的 kernel 可能在out还没算完的时候就开始读它。

5.3 参数选择与性能调优

多流不是越多越好。stream 数量太多，event 同步的开销会吃掉并行带来的收益。我的经验是：stream 数量控制在 2 到 4 个，超过这个数，同步开销通常大于并行收益。

具体怎么定？看你的 pipeline 有几个真正独立的阶段。如果两个阶段之间有严格的数据依赖，那它们本质上还是串行的，多流帮不上忙。只有当两个阶段可以真正并行（比如处理不同的 batch），多流才有意义。

另外，event 的粒度也要注意。event 记录得太频繁，同步开销大；记录得太稀疏，并行度上不去。一个经验法则是：在每个"数据交接点"记录 event，也就是一个 stream 的输出要交给另一个 stream 消费的地方。其他地方不用记。

6. 踩坑实录：那些 record_stream 救不了的场景

6.1 CUDA Graph 捕获下的 record_stream

CUDA Graph 是另一个重灾区。Graph 捕获的时候，所有操作被记录成一个静态图，包括record_stream和 event 操作。但 Graph 重放的时候，这些操作的语义和普通执行不完全一样。

我遇到过一个 case：在 Graph 捕获模式下，record_stream被捕获进了图里，但 Graph 重放时 allocator 的状态和捕获时不同，导致record_stream记录的 stream 集合和实际使用的 stream 对不上。结果是显存被提前复用，出现随机错误。

解决办法是在 Graph 捕获之前，把所有跨 stream 的 tensor 都预先分配好，并且在捕获时显式地建立 event 依赖。Graph 内部的record_stream尽量少用，因为它的语义在重放时不够可靠。

6.2 自定义 autograd Function 里的多流

如果你在自定义 autograd Function 的 forward 里用了多流，backward 也要对应处理。很多人只在 forward 里调了record_stream和wait_event，backward 里忘了，结果训练时梯度出错。

backward 的多流处理和 forward 是对称的：backward 的每个 stream 也要record_stream，也要建立 event 依赖。而且要注意，forward 保存的中间 tensor 在 backward 时可能已经被释放了，如果 backward 要用，必须在 forward 时就record_stream到 backward 会用到的 stream 上。

6.3 常见问题速查表

避坑技巧：compute-sanitizer是排查这类问题的神器。用compute-sanitizer --tool memcheck python your_script.py跑，它会精确报出哪一行代码访问了已释放的显存。虽然慢，但定位问题非常准。我一般先用它跑一个小 batch，确认没有内存问题，再放大规模跑性能测试。

7. 我个人的几条实战经验

写多流代码这几年，我总结了几条不太写在文档里但很管用的经验。

第一条，先写单流版本，确认正确后再改多流。多流的 bug 极难调试，如果你在单流下都没跑通，多流下只会更乱。我见过太多人一上来就写多流，结果连基本逻辑对不对都判断不了。

第二条，每个 stream 的边界都要有明确的 event。不要依赖"隐式同步"，比如靠 Python 的 GIL 或者 tensor 的引用计数来保证顺序。这些在单流下可能碰巧成立，多流下全是坑。显式地记录 event、等待 event，代码多几行，但省下的是几天的调试时间。

第三条，用torch.cuda.synchronize()做粗粒度验证。在开发阶段，可以在每个 stream 段落后加一个torch.cuda.synchronize()，把多流强制串行化。如果串行化后结果正确，说明逻辑没问题，问题出在同步上；如果串行化后还是错，说明逻辑本身有问题。这是一个快速二分定位的手段。

第四条，关注 allocator 的统计信息。torch.cuda.memory_stats()能看到显存池的分配、复用、延迟回收等数据。如果num_alloc_retries或者num_ooms在增长，说明显存压力大，record_stream缺失的风险就高。定期看一眼这些指标，能提前发现隐患。

最后再分享一个小技巧：如果你不确定某个 tensor 需不需要record_stream，可以先用一个"保守版"——所有跨 stream 的 tensor 都调record_stream，所有跨 stream 的依赖都加wait_event。跑通之后，再用 profiler 看哪些同步是多余的，逐步删减。先保证正确，再优化性能，这个顺序不能反。