扒开 PyTorch 源码:一行简单的 a + b 到底在 C++ 核心里走了多少弯路?

扒开 PyTorch 源码:一行简单的 a + b 到底在 C++ 核心里走了多少弯路?
引言a + b。一行 Python,一个加法。你写下它的时候大概只关心两件事:a 和 b 的 shape 能不能广播,结果是留在 GPU 上还是搬回 CPU。但在 PyTorch 的 C++ 核心里,这行代码触发了一条至少 5 层深的 dispatch 链——Autograd 层记录梯度图、Autocast 层检查是否该降精度、BackendSelect 层确定具体走哪个后端、最后才落到 CPU 或 CUDA 的真实 kernel 上执行那次加法。每一层做完自己的事,把剩余的 DispatchKeySet 交给下一层,下一层再从里面按位选出最高优先级的 key,跳到对应的 kernel。整条链走完,framework 自己消耗的时间在 5 到 25 微秒之间。这个数字放在一次 1024×1024 矩阵乘法旁边不值一提——kernel 本身就要跑几百微秒。但如果你的张量是 1×1 的标量,kernel 执行不到 1 微秒,dispatch 的固定开销反而占了总延迟的 90% 以上。在 2024 年字节跳动的 veScale 论文里,MoE(Mixture of Experts)模型训练中 dispatch overhead 占总训练时间的 15-30%,因为 MoE 的特点就是大量小算子密集调用。PyTorch GitHub issue #72746 里一个用户的 profile 结果更直接:CPU 端 70% 的时间花在aten::函数的框架代码里,GPU kernel 反而在等 CPU 把 dispatch 做完。小张量训练慢,常见的第一反应是"算力不够"或"数据搬运太多"。但 profiler 一打开,瓶颈往往在 dispatch—