1. 一个让我排查到凌晨三点的显存报错
先说结论：在多 CUDA streams 场景下，record_stream和wait_event必须配对使用，缺一个都可能让你在某个不经意的时刻收到CUDA error: an illegal memory access was encountered或者更隐蔽的显存数据错乱。这个坑我在一个多流推理管线的项目里踩过，当时现象是单测全过、压测偶发崩溃、定位花了大半天，最后发现根因就是漏了一个wait_event。
这篇内容适合两类人看：一类是正在用 PyTorch 做多流并行、想搞清楚record_stream到底在干什么的工程师；另一类是被显存复用、异步执行、stream 同步问题折磨过，想系统梳理一遍 CUDA stream 语义的开发者。我会从 PyTorch 的显存分配器机制讲起，把record_stream和wait_event各自的职责拆开，再给出一套可以直接抄的多流代码模板，最后附上我实际踩过的几个坑和排查手法。
需要提前说明的是，本文涉及的 CUDA 版本行为以 CUDA 11.x / 12.x 和 PyTorch 2.x 为参考，不同版本在分配器实现细节上可能有差异，但核心语义是一致的。如果你用的是比较老的 PyTorch 1.x，部分 API 行为需要自行对照官方文档确认。
2. 为什么多流场景下显存会被提前复用
2.1 PyTorch 显存分配器的“缓存复用”逻辑
要理解这个坑，得先理解 PyTorch 的 CUDA 缓存分配器（CUDACachingAllocator）是怎么工作的。它并不是每次torch.empty都向驱动申请新显存，而是维护了一个按块大小分级的空闲池。当你释放一个 tensor（引用计数归零）时，它占用的显存块不会立刻还给驱动，而是被放回空闲池，等待下一次分配复用。
这个设计本身是为了性能——向驱动申请显存是昂贵操作，缓存复用能把分配开销降到接近零。但问题就出在“释放”这个动作的判定上：PyTorch 认为一个 tensor 的引用计数归零，就可以回收它的显存块了。在单流同步执行的场景下，这个假设是成立的，因为所有 CUDA 操作都在同一个流里按顺序执行，tensor 用完了就是真的用完了。
可一旦引入多流，情况就变了。假设你在 stream A 上产生了一个 tensor，然后把它传给 stream B 上的某个 kernel 使用。此时如果 stream A 上的后续操作让这个 tensor 的引用计数归零，分配器会立刻把这块显存标记为空闲，可能马上分配给 stream A 上的另一个新 tensor。但 stream B 上的 kernel 可能还没执行完，它读到的就是一块已经被覆写的数据——这就是典型的 use-after-free，表现为数据错乱或者 illegal memory access。
2.2 一个具体的翻车场景
我用一个简化例子说明。假设你在做多流并行的数据预处理加推理：
s1 = torch.cuda.Stream() s2 = torch.cuda.Stream() with torch.cuda.stream(s1): data = torch.randn(1024, 1024, device='cuda') data = data * 2 # stream A 上的操作 # 把 data 交给 stream B 用 with torch.cuda.stream(s2): result = some_kernel(data) # 此时 data 在 Python 侧可能已经不再被引用 del data如果del data之后，stream A 上又分配了新 tensor，分配器很可能就把原来data那块显存给了新 tensor。而 stream B 上的some_kernel还在读data的旧地址，读到的就是新 tensor 的数据。这种 bug 最恶心的地方在于：它不一定每次都崩，取决于两个流的执行时序，可能跑一万次才崩一次，让你误以为是随机硬件故障。
2.3 record_stream 到底做了什么
record_stream就是为解决这个问题而生的。它的语义是：告诉分配器，这个 tensor 的显存在指定 stream 上还有未完成的使用，不要提前回收。调用方式很简单：
data.record_stream(s2)这行代码的意思是：在 stream B 上，data这块显存还有待执行的操作，分配器必须等到 stream B 执行到相应位置后，才能把这块显存放回空闲池。PyTorch 内部会记录一个 stream 到 tensor 的映射，当 stream B 上的操作推进到某个同步点后，才真正释放。
但这里有个关键点很多人会忽略：record_stream只保证显存不被提前复用，它不保证数据依赖的正确性。也就是说，它解决的是“显存生命周期”问题，不是“执行顺序”问题。如果你的 stream B 需要等 stream A 上的某个操作完成才能开始，那是wait_event的职责，record_stream管不了。
3. wait_event 与 record_stream 的分工
3.1 两者解决的是不同维度的问题
我把这两个 API 的职责用一张表说清楚，这是我在实际排查中总结出来的，比官方文档更直白：
| API | 解决的问题 | 作用对象 | 不解决什么 |
|---|---|---|---|
record_stream | 显存生命周期，防止提前复用 | 显存块 | 不保证执行顺序，不保证数据就绪 |
wait_event | 执行顺序，保证跨流依赖 | 流之间的同步点 | 不管理显存回收 |
很多人第一次遇到多流 bug 时，会以为加一个record_stream就万事大吉，结果发现还是偶发崩溃。原因就是：record_stream只保证显存不被复用，但如果 stream B 的 kernel 在 stream A 的生产操作完成之前就启动了，它读到的就是还没写完的数据。这时候必须用wait_event让 stream B 等待 stream A 上的某个事件。
3.2 正确的配对写法
一个完整的多流数据传递应该长这样：
s1 = torch.cuda.Stream() s2 = torch.cuda.Stream() # 在 stream A 上生产数据 with torch.cuda.stream(s1): data = torch.randn(1024, 1024, device='cuda') data = data * 2 # 记录一个事件，标记数据已就绪 ev = torch.cuda.Event() ev.record(s1) # 在 stream B 上消费数据 with torch.cuda.stream(s2): # 先等 stream A 的数据就绪 s2.wait_event(ev) result = some_kernel(data) # 告诉分配器：data 在 s2 上还有用，别提前回收 data.record_stream(s2) # 此时可以安全地释放 Python 引用 del data注意顺序：wait_event要在使用data之前调用，record_stream要在使用data之后调用。这两个顺序搞反了都会出问题。wait_event放后面等于没等，record_stream放前面虽然也能记录，但语义上应该在确认使用之后。
3.3 为什么 record_stream 要放在使用之后
这里有个细节值得展开。record_stream的调用时机其实比较灵活，放在使用前或使用后都能记录，但放在使用后更符合语义，也更安全。原因是：如果你在使用前调用record_stream，然后中间发生了异常导致实际没用到这个 tensor，那这块显存会被多占用一段时间，虽然不影响正确性，但浪费了显存。放在使用后调用，能确保只有真正被跨流使用的 tensor 才被标记。
不过要注意，record_stream必须在 tensor 的 Python 引用还存在的时候调用。如果你先del data再想record_stream，那就晚了，引用已经没了。所以标准写法是：使用完立刻record_stream，然后再让引用自然消失。
4. 完整的多流推理管线实操模板
4.1 场景设定与整体结构
我拿一个实际项目里的多流推理管线来演示。场景是：一个模型分两段，前段在 stream A 上跑，后段在 stream B 上跑，中间需要传递中间激活值。同时还有一个独立的拷贝流负责 H2D 和 D2H 传输。这种结构在需要流水线并行的推理服务里很常见。
整体思路是：用事件（Event）串联三个流的依赖关系，用record_stream管理跨流 tensor 的显存生命周期。下面是我实际用的代码骨架，去掉业务逻辑后保留核心同步部分。
4.2 代码实现与逐行说明
import torch class MultiStreamPipeline: def __init__(self, model_a, model_b): self.model_a = model_a self.model_b = model_b # 三个流：前段计算、后段计算、数据传输 self.stream_a = torch.cuda.Stream() self.stream_b = torch.cuda.Stream() self.stream_copy = torch.cuda.Stream() # 复用事件对象，避免频繁创建 self.ev_a_done = torch.cuda.Event() self.ev_copy_done = torch.cuda.Event() def forward(self, host_input): # 1. 在拷贝流上做 H2D with torch.cuda.stream(self.stream_copy): gpu_input = host_input.to('cuda', non_blocking=True) self.ev_copy_done.record(self.stream_copy) # 2. 前段计算，等拷贝完成 with torch.cuda.stream(self.stream_a): self.stream_a.wait_event(self.ev_copy_done) mid = self.model_a(gpu_input) # 标记 mid 在 stream_b 上还有用 mid.record_stream(self.stream_b) self.ev_a_done.record(self.stream_a) # 3. 后段计算，等前段完成 with torch.cuda.stream(self.stream_b): self.stream_b.wait_event(self.ev_a_done) out = self.model_b(mid) return out逐行拆解几个关键点。第一，gpu_input在拷贝流上产生，前段计算流通过wait_event等它就绪，这是执行顺序的保证。第二，mid在前段流产生，后段流通过wait_event等前段完成，同时mid.record_stream(self.stream_b)保证mid的显存在后段流用完之前不被回收。第三，ev_a_done和ev_copy_done是复用的，避免每次 forward 都创建新 Event 对象，减少开销。
4.3 参数选择与性能权衡
这里有几个可以调的参数。第一个是流的数量。流不是越多越好，每个流都有调度开销，而且流之间的同步点会引入等待。我实测下来，对于两段式模型，三个流（计算、计算、拷贝）是比较平衡的配置。如果模型分段更多，可以按段数增加计算流，但要注意同步点的数量会线性增长。
第二个是 Event 的复用策略。Event 对象本身创建开销不大，但频繁创建销毁会有 GC 压力。我习惯在__init__里预创建一组 Event，循环使用。如果你的管线是动态深度的，可以用一个 Event 池来管理。
第三个是non_blocking=True的配合。H2D 拷贝要真正异步，必须满足两个条件：源数据是 pinned memory，且non_blocking=True。如果源数据是普通 pageable memory，即使设了non_blocking，驱动也会在后台做一次隐式同步，异步就失效了。这一点在做数据加载优化时特别重要。
5. 那些年我踩过的坑与排查手法
5.1 坑一：以为 record_stream 能替代 wait_event
这是我最早踩的坑。当时看到record_stream的文档说“防止显存提前回收”，就想当然地以为它也能保证数据依赖。结果代码跑起来偶发数据错乱，排查了很久才发现：record_stream只保证显存不被复用，但 stream B 的 kernel 可能在 stream A 的生产操作完成前就启动了，读到的是半成品数据。
排查手法：在怀疑有依赖问题的两个流之间插入torch.cuda.synchronize()，如果加上同步后 bug 消失，基本可以确定是执行顺序问题，需要补wait_event。这个方法虽然粗暴，但定位效率极高。
5.2 坑二：record_stream 调用时机太晚
有一次我在一个 helper 函数里做跨流传递，tensor 传进去之后在函数内部用了，但record_stream是在函数返回后才调的。结果函数内部有个临时变量让 tensor 的引用计数提前归零，显存被回收，record_stream调用时对象已经失效。这种 bug 很隐蔽，因为 Python 的引用计数是运行时行为，静态看代码不容易发现。
排查手法：用torch.cuda.memory_summary()观察显存分配曲线，如果发现某块显存在预期之外的时间点被复用，就重点检查record_stream的调用位置是否在引用消失之后。
5.3 坑三：多流下的显存碎片化
多流场景下，不同流的分配和释放交错进行，容易造成显存碎片化。我遇到过一次：单流下能跑通的 batch size，改成多流后 OOM 了。原因是碎片化导致没有足够大的连续块。解决办法是设置PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True，让分配器支持可扩展段，能显著缓解碎片问题。
5.4 常见问题速查表
| 现象 | 可能原因 | 排查方向 |
|---|---|---|
| 偶发 illegal memory access | 显存提前复用 | 检查 record_stream 是否遗漏 |
| 数据错乱但无报错 | 执行顺序错误 | 检查 wait_event 是否遗漏 |
| 多流下 OOM | 显存碎片化 | 开启 expandable_segments |
| 异步拷贝没生效 | 源数据非 pinned | 用 pin_memory 包装 |
| 性能不如单流 | 同步点过多 | 减少流数量或合并操作 |
6. 几个容易被忽略的细节
6.1 Event 的 disable_timing 选项
创建 Event 时有个enable_timing参数，默认是 False。如果你不需要测量时间间隔，保持默认就好，因为开启计时会有额外开销。我见过有人为了“以后可能要用”就开着计时，结果在热路径上创建大量 Event，性能掉了一截。Event 应该预创建复用，而不是每次 forward 都 new 一个。
6.2 默认流与其他流的关系
PyTorch 的默认流（default stream）有特殊的同步语义，它和显式创建的流之间的交互规则比较绕。我的建议是：在多流代码里，尽量全部使用显式创建的流，不要混用默认流。默认流的隐式同步行为会让你的依赖分析变得复杂，容易出意料之外的等待或非等待。如果必须用默认流，务必查清楚当前 PyTorch 版本的默认流语义。
6.3 跨流传递时的 dtype 和 device 一致性
record_stream只关心显存块，不关心 tensor 的 dtype 和 device。但跨流传递时，如果两个流在不同的 device 上（多卡场景），record_stream的行为需要额外注意。多卡跨流传递涉及 P2P 访问，同步逻辑更复杂，建议先用torch.cuda.synchronize()验证正确性，再逐步去掉同步优化性能。
6.4 用 profiler 验证同步点
PyTorch Profiler 能可视化流之间的依赖关系。我习惯在优化多流代码时开一次 profiler，看时间线上各个流的执行区间和等待区间。如果发现某个流大量时间在等待，说明同步点设置得太保守，可以尝试放宽依赖。如果发现流之间有重叠但结果不对，说明依赖漏了。这个工具比盲猜高效得多。
7. 我个人的一点经验
多流编程的核心难点不在于 API 本身，而在于脑子里要同时维护两条时间线：显存的生命周期线和操作的执行顺序线。record_stream管前者，wait_event管后者，两者缺一不可。我现在的习惯是：只要一个 tensor 跨流使用，就条件反射地检查这两件事有没有做全。
另外，多流代码的调试成本远高于单流，所以我的策略是：先用单流把逻辑跑通，确认数值正确，再逐步引入多流并保持每一步都可回退。每次只加一个流，加完立刻用 profiler 和数值对比验证，不要一次性把整个管线改成多流。这样出问题时，你能快速定位是哪一步引入的。
最后分享一个小技巧：如果你不确定某个 tensor 是否需要record_stream，可以先用一个保守策略——所有跨流传递的 tensor 都加record_stream，跑通之后再根据 profiler 结果去掉不必要的调用。多调几次的开销远小于排查一个偶发崩溃的时间。
