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, devicecuda) 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, devicecuda) 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_blockingTrue) 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_blockingTrue的配合。H2D 拷贝要真正异步必须满足两个条件源数据是 pinned memory且non_blockingTrue。如果源数据是普通 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_CONFexpandable_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 结果去掉不必要的调用。多调几次的开销远小于排查一个偶发崩溃的时间。