★ FEATURED ARTICLE 前两天调一套多路视频推理服务遇到一个非常折磨人的Bug单条CUDA stream运行稳如老狗两条stream一起跑就开始间歇性报错有时候是illegal memory access有时候干脆没有报错只是输出结果里混入上一批数据。一开始我怀疑是kernel内部越界花了大半天时间打日志、查内存最后定位到是一个显存复用相关的老朋友——record_stream准确点说是多CUDA stream并发时没有正确使用wait_event导致的跨流依赖缺失。如果你也在用PyTorch自定义stream做多路推理、多batch异步流水线或者在纯CUDA里手动管理多stream并发这篇记录值得看完。我会把整个坑的成因、record_stream的真实语义以及如何在多stream场景里把事件依赖补齐一步步说清楚。1. record_stream到底在记录什么1.1 显存池与“假空闲”问题先从PyTorch的显存管理说起。PyTorch的CUDACachingAllocator在底层维护一个大缓存池第一次分配会真的调用cudaMalloc之后所有分配都尽量从池里取只有池里没有合适块时才去问驱动要新内存。这样做的好处是省去反复cudaMalloc/cudaFree的系统开销坏处是“池里的空闲块”并不等于“GPU上的活已经干完”。举个例子你在默认流上分配了一个buf然后把它丢给自定义流s1去写。buf这个Python对象可能很快被del显存块回到池子里变成free状态。可是GPU上s1队列里的kernel可能还压根没执行完。如果这时候另一个streams2需要分配显存分配器看到池子里有这么一块“空闲”内存顺手就给出去了。结果就是s1的kernel还在往这个地址写数据s2的kernel也开始往同一个地址写数据两边互踩轻则结果错乱重则直接非法内存访问。这个问题的本质是显存块的使用生命周期不等于Python对象的生命周期也不等于分配时所在stream的生命周期。我一开始想不通为什么单stream没问题。后来想明白了单stream是串行队列块回收后再次分配时同一队列里的后续kernel天然排在前面kernel后面GPU硬件会保证顺序。但多stream之间默认是并发、无顺序约束的一旦内存被错分谁先写谁后写完全看调度心情。用生活里的例子类比显存缓存池像一个仓库管理员他知道货物什么时候被借走、什么时候归还。现在有两条流水线都在用同一件设备管理员只盯着第一条流水线的登记记录以为第一条线完工就能把设备租出去可实际上第二条流水线还在用。两拨人同时操作一台设备必然出问题。1.2 record_stream的官方语义与实现思路torch.Tensor.record_stream(stream)这个接口就是用来解决“跨stream使用但分配器不知道”的问题。它的官方语义是让分配器知道这个tensor还会在指定的stream上继续被使用因此它的底层显存块在该stream上的所有工作完成之前不能安全回收复用。调用了record_stream之后分配器会在底层为这个显存块记录一个和指定stream关联的事件通常对应cudaEventRecord。当这个块进入空闲列表分配器检查到它身上有关联的未完成事件就不会把它分配给新的分配请求直到事件真正完成块才会重新进入可复用状态。这个机制是GPU侧的异步等待不会立刻阻塞CPU线程但会有效推迟显存的归还。这里有个容易被忽略的关键点record_stream必须由用户在实际使用这个tensor的stream上下文里调用或者明确指定正确的stream。如果调用位置不对它记录的stream和你实际跑kernel的stream不是同一个那等于白记。我踩过的坑就发生在这一步。我用多stream做推理时习惯在一个stream上下文里创建中间张量又在另一个stream上下文里用。代码里虽然写了buf.record_stream(s1)但因为调用时的current stream已经不是s1了导致分配器理解的是“这块内存在当前流上有依赖”而不是“在s1上有依赖”。这个错误非常隐蔽因为单stream时不会暴露一旦并发跑起来就随机翻车。为了把行为说清楚我把三种情况列在下面场景分配器感知结果不调用record_stream认为显存块只和分配时所在stream相关跨stream复用时大概率数据竞争调用record_stream但调用位置/stream错误记录的是错误流上的依赖并发场景仍可能提前复用正确调用record_stream 正确跨流等待显存块在所有相关流上都标记了依赖缓存复用安全跨流顺序有保障2. 多stream并发为什么必须见event2.1 CUDA事件在多流协作里的地位CUDA stream是GPU上的任务队列每个队列里的kernel和memcpy按顺序执行但不同stream之间默认是互相独立的。想让两个stream之间产生先后关系唯一的正规方式就是用事件event建立依赖。流程很简单先在stream A上cudaEventRecord(evt, streamA)把这个事件标记插入A的队列然后在stream B上调用cudaStreamWaitEvent(streamB, evt)让B的后续任务等在evt之后。这样B虽然还是异步提交但GPU调度器会保证B里排在wait_event之后的kernel不会早于A里record事件之前的任务完成。如果你接触过纯CUDA编程对这个应该不陌生。但在PyTorch里很多人会忽视这一点因为PyTorch的自动内存管理把这些底层操作封装得太圆滑了。with torch.cuda.stream(s1):这种写法只是把当前执行环境切到s1并不会自动把s1的依赖关系同步给其他stream。s2该不等还是不等。我后来总结了一句话stream是流水线event是流水线之间的交接信号。没有交接信号两条流水线各跑各的共享的物料就一定会抢。2.2 为什么只调record_stream还不够有朋友可能会问既然record_stream内部会给显存块挂事件那分配器在复用前会检查事件完成状态是不是我只要调了它就行不用再调wait_event这个想法对了一半但恰恰是另一半导致了问题。record_stream管的是内存池的复用安全。它的意思是这块显存虽然没人引用了但某个stream上还有活没干完先别把它发给别人。它不会主动去阻塞另一个stream的kernel提交也不会改变两个业务流的调度顺序。看一个典型场景两个stream从同一个缓存池里拿到同一块物理显存。分配器在s1分配时调用了record_stream(s1)所以s2要分配时理论上会等s1的事件完成才把这块给出。但如果你不是通过分配器拿到这块内存而是自己维护了一个buffer pool提前把同一块显存分别交给了s1和s2那么record_stream根本管不到你的业务调度。此时s2上的kernel该写入还是会写入s1上的kernel可能还没写完。就算用的是分配器record_stream也不是银弹。分配器能保证“这块内存在事件完成前不再分给新的分配请求”但它不能保证“你新分配的这块内存不会和另一个还活着的tensor重叠”。后者靠的是分配器对活动块的追踪前者靠的是record_stream的延迟归还。两者职责不同容易混为一谈。所以在多stream并发里你真正需要在代码层面补上的是显式的跨流等待把s1上的完成事件挂到s2的队列头部让s2后续所有共享显存的操作都排在s1的事件后面。这正是标题里说的“要用wait_event”。2.3 一条完整的安全依赖链把事件依赖画出来整个逻辑其实很清晰在stream1上执行写buffer的kernel。在stream1上evt.record(stream1)给“写buffer完成”打个标记。在stream2上stream2.wait_event(evt)让stream2的后续任务等待这个标记。只有evt完成后stream2上读/写buffer的kernel才会真正执行。同时在buffer可能被释放前调用buffer.record_stream(stream1)确保显存池不会提前把这个块分给别人。注意第3步和第5步是两件事。第5步保护的是“不能让其他分配拿到这块地址”第3步保护的是“即使是同一个bufferstream2也不能抢跑”。两者配合才算完整。我顺手总结了一个通俗版依赖时刻表顺序stream1stream2T0写buffer的kernel入队执行wait_event(evt)阻塞后续任务T1evt.record()标记写完成点继续等待T2做其他与buffer无关的任务前序事件完成buffer读取kernel开始执行3. 从报错到修复一个真实的多流复现3.1 只看record_stream的坏代码为了把问题复现清楚我写了一个简化版本。这个例子不去追求模型复杂度只保留多stream共享显存的核心矛盾import torch import torch.cuda as tc def bad_case(): s1 torch.cuda.Stream() s2 torch.cuda.Stream() # 在默认流上分配一块显存后续两个流都要使用 buf torch.empty(4096, devicecuda) with torch.cuda.stream(s1): # 模拟stream1向buffer写入数据 buf.fill_(1.0) # 告诉分配器这块内存在s1上还有依赖 buf.record_stream(s1) with torch.cuda.stream(s2): # 模拟stream2立刻向同一块buffer写入数据 # 这里没有任何事件等待s2可能抢在s1的fill_完成前执行 buf.fill_(2.0) torch.cuda.synchronize() # 最终buf是1还是2完全不确定。 print(buf[:8])这段代码的bug很明显s2上的fill_(2.0)是异步提交的它入队时s1的fill_(1.0)不一定执行完。两个kernel操作同一块显存谁能先跑完全看GPU调度。就算调用了record_stream也没能让s2排在s1后面。3.2 用wait_event补上跨流依赖修复方式不复杂核心是在with torch.cuda.stream(s2):之前让s2等待一个由s1记录的事件def fixed_case(): s1 torch.cuda.Stream() s2 torch.cuda.Stream() buf torch.empty(4096, devicecuda) evt torch.cuda.Event() with torch.cuda.stream(s1): buf.fill_(1.0) # 在s1上插入完成标记 evt.record(s1) # 保护显存池别提前复用 buf.record_stream(s1) # 关键让s2等待s1上的事件然后s2的所有后续任务排在evt之后 s2.wait_event(evt) with torch.cuda.stream(s2): buf.fill_(2.0) torch.cuda.synchronize() print(buf[:8])这里evt.record(s1)的作用是给s1队列里“fill_完成之后”的位置做一个标记。s2.wait_event(evt)则把这个标记写成s2队列的一个前置条件。所以s2的fill_不会真正执行直到s1的fill_执行完毕。结果必然先写入1.0再写入2.0程序行为确定下来。需要注意的是这个例子里buf本身是同一个tensor所以我能用事件等待解决。如果两个流分配的是两块地址相同、但来自缓存池不同分配路径的显存那就要同时检查分配器日志确认物理地址没有重叠。3.3 为什么不用全局synchronize有人会说直接在两个流之间加一次torch.cuda.synchronize()不也能解决问题吗同步确实能解决正确性问题但代价是破坏并发。synchronize()会阻塞CPU并且让GPU执行队列里的所有任务都跑完才返回。在多stream流水线里这意味着stream1的全部工作做完之前stream2完全闲着CPU也在空等。如果stream1有十段任务只有第三段和buffer相关你就被迫等全部十段完成浪费掉后面七段的并行能力。wait_event的好处是只建立“第三段完成”这一点上的依赖stream1的后续任务可以继续排队执行stream2也可以提前做和buffer无关的准备工作。我把这个思路称为细粒度同步这是多stream优化里最有价值的地方。实际项目里流之间的依赖往往不是一条直线。比如stream1处理视频帧Astream2处理视频帧B中间只共享一个后处理buffer那只需要在“共享buffer被写完后”和“共享buffer被读取前”这两个点上做事件依赖其他时间两个流完全可以各自并行跑。3.4 多消费者场景怎么处理再往复杂一点说。如果同一块显存被三个stream依次使用比如stream1预处理写数据stream2推理读数据stream3后处理再读或再写那么最好的做法是每一步都维护事件依赖s1 torch.cuda.Stream() s2 torch.cuda.Stream() s3 torch.cuda.Stream() evt1 torch.cuda.Event() evt2 torch.cuda.Event() with torch.cuda.stream(s1): data.fill_(1.0) evt1.record(s1) data.record_stream(s1) s2.wait_event(evt1) with torch.cuda.stream(s2): result model(data) evt2.record(s2) data.record_stream(s2) s3.wait_event(evt2) with torch.cuda.stream(s3): postprocess(data, result)这里s2等待evt1保证看到的是s1写完整的数据s3等待evt2保证看到的是s2模型推理完成后的数据。每条依赖都是按需建立的没有把一个流的全部任务都同步掉。这个模式在纯CUDA C里也一样只是API换成了cudaEventRecord和cudaStreamWaitEvent。理解了事件依赖的原理换语言只是换个写法。4. 常见故障排查与避坑速查4.1 三个典型报错长什么样多stream显存冲突的症状并不是每次都一样我这次项目里就遇到了三种第一种是偶发CUDA error: an illegal memory access was encountered。这种最吓人因为报错地址随机换一次运行可能发生在不同kernel上。排查半天如果发现kernel逻辑没问题大概率就是显存块被两个stream同时用了。第二种是misaligned address。有时候不是地址本身没对齐而是因为显存块被覆盖你读到的是一个被其他数据篡改过的指针值或张量metadata。多stream复用时这类错也很常见。第三种最坑没有任何报错但推理结果里夹杂上一批数据。原因是上一轮kernel还没写完这一轮就已经开始复用同一块地址导致一部分数据是旧的。这种问题在torch.cuda.synchronize()隐藏很深的pipeline里特别难发现。4.2 定位多流显存冲突的几条路子遇到上述症状第一反应不要急着改kernel先怀疑“谁和谁共享了同一块显存”。这里有几个有效的排查方法按成本从低到高排打开CUDA_LAUNCH_BLOCKING1让所有kernel变为同步启动。如果问题变成必定复现或消失方向基本锁定在异步并发调度上。在分配前后打印关键tensor的data_ptr()。如果两个stream拿到的tensor地址相同证明它们共享了同一块物理显存。临时把显存缓存分配关掉或调整比如设置PYTORCH_NO_CUDA_MEMORY_CACHING1或不断调用torch.cuda.empty_cache()观察问题是否变化。用compute-sanitizer老工具叫cuda-memcheck跑一遍它通常能直接指出第一次非法访问发生在哪个kernel、哪个地址。在代码里临时插入torch.cuda.synchronize()如果同步后问题消失基本可以确认是跨流依赖缺失。我自己最常用的手段是前两个因为它们不需要额外工具打几行日志就能定位。先确认data_ptr()是否重叠如果重叠再看两个stream之间有没有事件依赖大部分问题立刻浮出水面。4.3 多流编程的避坑清单吃了一次大亏之后我给自己列了一张多流编程的避坑清单现在写代码都会对着过一遍record_stream只能解决显存池的延迟归还问题不能解决业务流之间的顺序问题。跨流共享数据时必须显式用wait_event或wait_stream建立依赖。record_stream要放在正确的stream上下文中调用最好显式传入实际使用该tensor的stream不要依赖当时的current stream。一个tensor被多个stream使用时最好对每个使用它的stream都调用record_stream。漏掉任何一个分配器都可能在那条漏掉的流上提前复用显存。不要对大量小tensor频繁调用record_stream。事件对象本身有开销记录太多会导致显存迟迟不回收反而把显存占用量顶上去。多线程场景下CPU线程各自操作stream时要确保对应的stream对象和事件对象生命周期合法PyTorch的Stream和Event默认不跨线程自动同步。我用一个表把这些要点收拢起来方便快速查阅症状最可能原因排查重点偶发illegal memory access两块活动tensor物理地址重叠打印data_ptr检查分配日志输出混入上一轮数据上一轮kernel未完成就开始复用buffer补record_stream wait_event加了record_stream仍出问题调用位置/stream不正确或缺少跨流等待检查调用上下文补wait_event显存占用异常上涨事件记录过多块迟迟不复用删除不必要的record_stream或显式同步并发下性能反而变差同步粒度过大比如频繁synchronize改用细粒度event只同步必要依赖点4.4 一个能保命的小习惯我后来养成的习惯是在写任何多stream代码之前先在纸上画一张“资源-流”依赖图。横轴画stream纵轴画共享资源哪个stream在哪个时刻读、哪个时刻写一目了然。画完图再写代码基本能提前发现八成以上的依赖缺失问题。举个例子如果图上出现“同一块bufferstream1在T2要写stream2在T1要写”那你就知道必须在T1和T2之间加一个事件依赖。如果图上只有stream1自己碰这块buffer那record_stream就够用了不需要额外的事件等待。这个习惯看着简单但比记一堆API细节靠谱得多。这次排错最后只改了几行代码但查错过程花了两天。表面上是record_stream和wait_event两个API配合的事背后却是对CUDA流调度模型的理解。多stream的并行度很诱人可一旦共享资源没有明确的依赖关系它回报给你的就是这种随机崩溃和脏数据。希望这篇记录能帮你省掉那两个熬夜排查的晚上。
阅读完成 · 觉得有帮助？
