1. 一个让显存悄悄泄漏的经典陷阱
如果你在用 PyTorch 写多流（multi-stream）的推理或训练代码，大概率见过record_stream这个 API，也大概率在某个深夜被一个诡异的现象折磨过：明明张量已经没人引用了，torch.cuda.memory_allocated()却一直不降，显存像沙子一样从指缝里漏掉，跑几百个 iteration 之后直接 OOM。更气人的是，单流跑得好好的，一上多流就出问题。
这个坑我自己踩过不止一次，也见过同事在 code review 里被反复追问。核心结论先摆在这里：record_stream只负责告诉分配器"这块显存被别的流用着，别急着回收"，它不负责同步。真正要保证跨流的数据依赖正确，必须配合wait_event（或者wait_stream）来建立事件依赖。只写record_stream不写wait_event，你得到的是"显存不泄漏了，但结果偶尔算错"；只写wait_event不写record_stream，你得到的是"结果对了，但显存可能被提前复用导致数据被覆盖"。两个都得写，顺序还有讲究。
这篇东西不是教科书式的 API 文档翻译，而是把我自己在实际项目里踩过的坑、排查过程、以及最后稳定下来的写法完整记录下来。适合正在做多流推理加速、CUDA Graph 捕获、或者自定义算子跨流调度的同学。如果你只是单流跑模型，这篇可以先收藏，等哪天要压榨 GPU 利用率了再翻出来。
2. 先把 CUDA 流和显存分配器的关系讲透
2.1 流是命令队列，不是数据容器
CUDA stream 本质上就是一个任务队列，你往里面塞 kernel、memcpy、event 记录等操作，GPU 按顺序执行。默认流（default stream / legacy stream）有隐式同步语义，而你自己torch.cuda.Stream()创建出来的流之间是完全异步的。这意味着流 A 上的 kernel 和流 B 上的 kernel 可以同时跑，也可以乱序完成，谁先谁后没有任何保证。
问题就出在这里。PyTorch 的 CUDA 缓存分配器（caching allocator）管理着一大块显存池，它回收显存的判断依据是"这块内存对应的张量引用计数归零"。但引用计数是 CPU 侧的概念，它根本不知道 GPU 上还有没有 kernel 正在读这块内存。单流情况下这不是问题，因为同一个流上的操作天然有序，分配器只要保证"下一个用到这块内存的分配发生在当前流的所有操作之后"就行。多流情况下，流 B 上的 kernel 可能还在读这块内存，而流 A 上已经有人申请了同样的地址，数据就被写花了。
2.2 record_stream 到底做了什么
tensor.record_stream(stream)的作用是：把这块显存标记为"被 stream 使用中"，让分配器在回收时知道要等这个 stream 上的工作完成。它内部维护的是一个 per-stream 的待释放列表，当这个 stream 执行到某个同步点时，才真正把内存还给池子。
关键点在于：record_stream不产生任何 GPU 侧的同步操作。它纯粹是 CPU 侧给分配器的一个提示。所以如果你写了：
with torch.cuda.stream(s2): y = some_kernel(x) # x 是在 s1 上产生的 x.record_stream(s2)这段代码并没有保证some_kernel执行时x的数据已经就绪。如果x是在 s1 上刚算出来的，s2 上的 kernel 完全可能先跑，读到的是旧数据或者未初始化数据。这就是"结果偶尔算错"的来源，而且因为 GPU 调度有随机性，这种 bug 往往几百次才复现一次，极难定位。
2.3 wait_event 补上缺失的同步
正确的做法是在使用跨流数据之前，先让当前流等待数据生产流：
ev = torch.cuda.Event() with torch.cuda.stream(s1): x = some_kernel_a() ev.record(s1) with torch.cuda.stream(s2): s2.wait_event(ev) # 关键：建立依赖 y = some_kernel_b(x) x.record_stream(s2) # 关键：防止提前回收wait_event会在 s2 上插入一个等待操作，s2 后续的所有工作都要等 ev 被 s1 记录之后才能开始。这样数据依赖就正确了。而record_stream保证在 s2 用完 x 之前，分配器不会把 x 的显存分给别人。
注意：
wait_event和record_stream解决的是两个正交的问题，前者管"数据什么时候能读"，后者管"显存什么时候能回收"。少任何一个都会出问题，只是表现不同。
3. 为什么单流代码迁移到多流就翻车
3.1 分配器的"流感知"边界
PyTorch 的 caching allocator 有一个stream参数概念。当你在一块显存上调用record_stream(s)，分配器会把这块 block 加入 s 的_blocks_pending_free或者类似的待释放结构。当 s 上发生一次同步（比如s.synchronize()或者一个 event 被 query 到完成），分配器才会真正释放。
但这里有个容易被忽略的细节：分配器只对"显式调用过 record_stream 的流"做延迟回收。如果你在 s2 上用了 s1 产生的张量，却忘了 record_stream，分配器会认为这块内存在 s1 的视角下已经可以回收了（因为引用计数归零），于是可能在 s2 的 kernel 还没跑完时就把它分给了新的分配。这就是数据被覆盖的经典场景。
3.2 一个真实的翻车现场
我之前做一个多流并行的预处理 + 推理 pipeline，结构大概是：流 A 做图像解码和归一化，流 B 做模型前向。为了省显存，我在流 B 用完输入张量后调用了record_stream，但忘了在流 B 开头wait_event。测试集小的时候一切正常，因为 GPU 负载低，流 A 总是先跑完。上了生产环境大 batch 之后，偶发精度下降，排查了两天才定位到是流 B 读到了流 A 还没写完的 buffer。
后来我把这个模式抽象成一个 checklist：任何跨流传递的张量，生产侧 record event，消费侧 wait event，消费侧用完 record_stream。三件事缺一不可。
3.3 和 CUDA Graph 的交互
如果你在用 CUDA Graph 捕获多流代码，情况会更复杂。Graph 捕获期间，record_stream的行为和 eager 模式不完全一样，因为 Graph 内部的分配是静态的。这时候wait_event反而更重要，因为 Graph 重放时不会重新执行 Python 侧的 record_stream 逻辑。我一般建议在 Graph 捕获前就把流依赖用 event 显式串好，捕获后不要再依赖 record_stream 做动态回收。
4. 手把手写一个正确的多流示例
4.1 场景设定
假设我们要做这样一件事：流 A 上生成一个张量x，流 B 上对它做一次逐元素运算得到y，最后在主流程里汇总。这个例子足够简单，但包含了所有关键要素。
4.2 错误写法对照
先看错误写法，很多人第一版就是这么写的：
import torch s1 = torch.cuda.Stream() s2 = torch.cuda.Stream() with torch.cuda.stream(s1): x = torch.randn(1024, 1024, device='cuda') with torch.cuda.stream(s2): y = x * 2 # 危险：没有等待 s1 x.record_stream(s2) # 只做了回收保护 torch.cuda.synchronize() print(y.sum())这段代码在x的生成很快、s2启动稍慢时能跑对，但没有任何同步保证。x * 2完全可能在torch.randn完成前就开始读x。
4.3 正确写法
import torch s1 = torch.cuda.Stream() s2 = torch.cuda.Stream() ev = torch.cuda.Event() with torch.cuda.stream(s1): x = torch.randn(1024, 1024, device='cuda') ev.record(s1) # 在 s1 上记录事件 with torch.cuda.stream(s2): s2.wait_event(ev) # s2 等待 s1 完成 y = x * 2 x.record_stream(s2) # 保护 x 的显存不被提前回收 torch.cuda.synchronize() print(y.sum())4.4 参数与顺序的讲究
ev.record(s1)必须放在x的生产操作之后，否则事件记录的是"生产开始前"的状态，等待就失去意义。s2.wait_event(ev)必须放在消费操作之前。x.record_stream(s2)放在消费操作之后即可，但为了可读性我习惯紧跟消费操作。
还有一个细节：torch.cuda.Event()默认enable_timing=False，如果你不需要测时间，保持默认就行，开销更小。如果需要复用 event，可以用ev.record()覆盖，但要注意不要在还没被 wait 完就覆盖，否则等待的语义会乱。
实操心得：在循环里反复创建 Event 对象会有一定开销，我一般会在循环外创建一个 event 池，或者直接复用同一个 event 变量。实测在 1000 次迭代的场景下，复用 event 能省下几毫秒的 CPU 时间，虽然不多，但积少成多。
5. 排查这类问题的实用套路
5.1 用环境变量逼出错误
PyTorch 提供了CUDA_LAUNCH_BLOCKING=1，它会让所有 CUDA 操作同步执行。这能帮你确认"是不是异步导致的错误"，但注意它会让多流退化成串行，所以只能用来定位，不能用来验证多流正确性。另一个有用的环境变量是PYTORCH_NO_CUDA_MEMORY_CACHING=1，关掉缓存分配器后如果问题消失，基本可以确定是显存复用相关的 bug。
5.2 用 compute-sanitizer 抓 race
NVIDIA 的compute-sanitizer --tool racecheck能检测共享内存和全局内存的竞争，但对跨流的显存复用问题覆盖有限。更实用的是--tool memcheck，它能抓到"读未初始化内存"这类错误，而跨流提前复用往往表现为读到脏数据。我一般会先用 memcheck 跑一遍小规模复现，如果报 invalid read 或者 uninitialized read，就往流同步方向查。
5.3 常见问题速查表
| 现象 | 可能原因 | 排查方向 |
|---|---|---|
| 显存缓慢增长不释放 | 只 wait_event 没 record_stream | 检查消费侧是否漏了 record_stream |
| 结果偶发错误、精度下降 | 只 record_stream 没 wait_event | 检查消费侧开头是否有 wait_event |
| 单流正常多流出错 | 流依赖没建立 | 用 CUDA_LAUNCH_BLOCKING 验证 |
| Graph 重放结果不对 | 捕获期依赖没串好 | 捕获前显式 event 同步 |
| OOM 但 nvidia-smi 显示有空闲 | 分配器碎片或 pending free 堆积 | 检查 record_stream 是否过度使用 |
5.4 一个容易忽略的坑：record_stream 用过头
有人为了保险，在每个流上都对每个张量调 record_stream，结果导致显存迟迟不回收，pending free 列表越堆越长，最后反而 OOM。record_stream 应该只用在"确实跨流使用"的张量上，同流内的张量不需要。我见过一个项目里，作者在数据加载流和计算流之间对所有中间张量都 record_stream，显存占用直接翻倍。
6. 把模式固化下来：封装与工程实践
6.1 封装一个跨流传递的小工具
与其每次手写 event 和 record_stream，不如封装成一个上下文管理器：
from contextlib import contextmanager @contextmanager def cross_stream_tensor(tensor, producer_stream, consumer_stream): ev = torch.cuda.Event() with torch.cuda.stream(producer_stream): ev.record(producer_stream) with torch.cuda.stream(consumer_stream): consumer_stream.wait_event(ev) yield tensor tensor.record_stream(consumer_stream)用起来就是：
with cross_stream_tensor(x, s1, s2) as t: y = t * 2这样依赖和回收都自动处理了，代码可读性也好很多。当然这个封装假设 tensor 已经在 producer_stream 上生产完毕，更复杂的场景需要调整。
6.2 在真实 pipeline 里的落地
我现在的多流 pipeline 一般是这样组织的：数据预处理一个流，H2D 拷贝一个流，计算一个流，三个流之间用 event 串成链。每个跨流张量都走上面的封装。实测在 A100 上，相比单流，端到端吞吐能提升 20% 到 35%，具体取决于预处理和计算的比例。如果预处理很轻，提升就有限，这时候多流反而增加调度开销，不如单流。
6.3 什么时候不该用多流
不是所有场景都适合多流。如果你的 kernel 本身已经占满 GPU，多流不会带来任何加速，只会增加同步复杂度。如果数据依赖是严格串行的，多流也没意义。我一般会先 profile 一下，看 GPU 利用率是不是有明显空隙，有空隙才考虑多流。另外，小 batch 场景下多流的 event 开销可能超过收益，这时候老老实实单流更稳。
提示：多流优化的第一步永远是 profile，不是上来就改代码。用
torch.profiler或者 nsight systems 看清楚时间花在哪，再决定要不要引入多流。
7. 几个我踩过的具体坑和最终写法
7.1 坑一：在错误的流上 record event
有一次我把ev.record()写在了with torch.cuda.stream(s1)外面，结果 event 记录到了默认流上，等待语义完全错了。这种 bug 不会报错，只会让结果偶尔不对。后来我养成了习惯：event 的 record 一定紧跟在生产操作后面，且明确指定流。
7.2 坑二：忘记 record_stream 导致的数据覆盖
这个前面提过，表现是精度下降。定位方法是把 batch 调小、把流调成串行，如果问题消失，基本就是流同步问题。修复就是在消费侧补 record_stream。
7.3 坑三：Event 复用导致的等待失效
在循环里复用同一个 event 时，如果上一轮的 wait 还没完成就 record 了下一轮，等待的语义会乱。我现在的做法是每轮用独立的 event，或者用 event 池轮转。虽然多几个对象，但省心。
7.4 最终稳定下来的模板
def run_pipeline(x, s_prod, s_cons): ev = torch.cuda.Event() with torch.cuda.stream(s_prod): x = preprocess(x) ev.record(s_prod) with torch.cuda.stream(s_cons): s_cons.wait_event(ev) y = model(x) x.record_stream(s_cons) return y这个模板我在好几个项目里用过，只要生产侧和消费侧的流对应正确，基本不会再出问题。核心就是那句话：wait_event 管同步，record_stream 管回收，两个都要，顺序别错。
最后分享一个小技巧：如果你不确定某段多流代码是否正确，可以先用torch.cuda.synchronize()在关键位置插桩，确认逻辑对了再逐步去掉同步看性能。这样能把"逻辑错误"和"性能问题"分开排查，效率高很多。
