【笔记】一行代码让大模型推理的 overlap scheduler 失效

1. 先说结论

overlap scheduler 的前提是「GPU 上的推理」和「CPU 上的前后处理」能同时跑。任何一次 CPU 与 GPU 之间的同步等待,都会把这个并行切断 ------只要 cudaStreamSynchronize 一出现,CPU 就得停下来等 stream 空出来。

两种最常见的触发方式:

数据方向 触发写法 为什么会同步
GPU → CPU if args.is_greedy.all(): 要判断 GPU tensor 的值,就得把数据取回 CPU,触发隐式同步
CPU → GPU 同步的拷贝(source 不是 pinned memory) CPU 必须等数据传递完成才能继续

两者的代价并不相同:前者让 overlap 完全消失 ,后者只损失 cudaGraphLaunch 这一小段。原因见第 4 节。

2. 案例 1:在 sample 阶段判断 GPU tensor

阻塞的本质:要判断的东西在 GPU 上,CPU 就必须停下来等它传回 CPU。

sample 里如果有这样一行判断代码:

python 复制代码
if args.is_greedy.all():

它会让 overlap 失效。因为 args.is_greedy 是 GPU 上的 tensor,对它做判断就意味着数据要从 GPU 传到 CPU,会触发隐式的 CUDA synchronization------CPU 会阻塞等待 GPU 上所有排队的操作完成。

sample 是在 model 推理完成之后调用的,此时的同步会把 overlap 的异步模式变成串行执行。这里的阻塞非常明显,所有的 overlap 全部失效

3. 案例 2:CPU → GPU 的同步拷贝

阻塞的本质:CPU 往 GPU 传数据,CPU 必须停下来等传递完成。

效果和案例 1 类似,同样是因为 cudaStreamSynchronize 导致 overlap 失效。

nsys 示意图:

下一轮 forward batch 的 cudaGraphLaunch 无法提前启动,需要等 cudaStreamSynchronize 完成;而 cudaStreamSynchronize 要完成,则需要等上一轮的推理完成、空出 stream。于是整个推理变成:

时序 说明
现状 N 推理 → N+1 cudaGraphLaunch → N+1 推理 中间插进了 cudaGraphLaunch,无法全部 overlap
理想 N 推理 → N+1 推理 → ... 全流程都在执行推理,没有任何中断

注意 :CPU → GPU 的拷贝设置 non_blocking=True,且 source 是 pinned memory 时,不会触发 cudaStreamSynchronize

4. 两个案例的严重性区别

案例 阻塞位置 阻塞在等什么 代价
案例 1 sample 阶段(推理完成之后) 当前计算流里的推理任务全部完成 overlap 完全消失
案例 2 CPU → GPU 拷贝发起处 拷贝本身完成 只损失 cudaGraphLaunch 这一段

案例 2 只会导致 cudaGraphLaunch 这段时间无法 overlap。因为 cudaGraphLaunch 是异步执行 GPU 任务的:这个函数在 CPU 侧执行完之后,GPU 开始执行推理,而 CPU 可以去处理前一个请求的后处理、下一个请求的前处理,GPU 推理和这些处理就重叠上了。

案例 1 是在 sample 阶段引发了 GPU → CPU 方向的阻塞,需要等当前所在计算流里的推理任务完成,才能继续往下走。从 cudaGraphLaunch 到 sampler 阻塞的这段时间,本来正是该做 overlap 的区间,但「前一个请求的后处理、下一个请求的前处理」此时还没有机会执行,这段区间几乎没做任何工作,因此 overlap 完全消失了。

一句话概括 :在 cudaGraphLaunch 执行完之后到出现阻塞的这段期间是可以 overlap 的,这段期间越长,overlap 效果越好

在 overlap 正常执行的情况下,阻塞应该要等到下一轮、等的是上一轮结果完成的时刻:

图中第一个箭头指向的是 cudaGraphLaunch,第二个箭头指向的是时间同步等待,等待的是第一个箭头发起的 decode 推理的完成事件通知。这段期间,计算流的 GPU 在执行推理,而 CPU(或者是非计算流的 GPU)可以执行其他工作,这就实现了 overlap。

5. 小结

  • overlap 的有效窗口 = cudaGraphLaunch 执行完 → 出现阻塞。
  • 窗口越长,overlap 效果越好;窗口被同步切断,overlap 就失效。
  • 阻塞出现得太早(案例 1,sample 阶段)→ overlap 完全消失;阻塞只挡住 cudaGraphLaunch(案例 2)→ 只损失一小段。
  • 正常的 overlap 里,阻塞应该发生在下一轮,等的是上一轮推理的完成事件通知。