你写c = a + b的时候,GPU上真正跑起来的其实是个叫“内核”的小程序。搞清楚它是怎么回事,读显存的代价又在哪里,你就能看懂性能瓶颈,不再把GPU当黑盒子。这跟操作系统内核没关系,单纯是GPU干活的方式。
内核就是GPU接到的活儿
你肯定在Python里写过无数次c = a + b。换成PyTorch的张量也一样,只不过张量这玩意儿就是一堆数组成的数组。把张量搬到GPU上,意思就是这堆数住进了显存里,不在普通内存待着了。当a和b俩张量都在GPU上时,这一行代码跑得飞快,快到你根本不会多想。
现在把事儿想得更细一点。假设a和b就是两个浮点数,孤零零地蹲在显存里。还是那行c = a + b。芯片上到底跑了啥?
答案就是内核。在这个场景下,内核就是个在GPU上处理一小块数据的小程序。它不是电脑开机时加载的操作系统内核,也不是线性代数课本里说的那个数学核。这个词被到处借用来用去,搞混了真不怪你。在GPU的世界里,内核指的就是:一小段函数,GPU被叫去立刻执行它,并行地,在你给的数据上跑。
把这套逻辑弄明白之后,随便拿一段PyTorch代码,你都能数清楚GPU会跑多少个内核。这听着像个小把戏,但它确实是第一步,能帮你摆脱那种把GPU当黑盒子的感觉——就是那种模型跑得慢,你还不知道为啥的憋屈感。但凡跟GPU性能沾边的问题,追根溯源都得回到“跑了哪些内核,它们到底在干啥”。所以,咱就从这儿开始。
内核里是谁在干活
把a和b弄大一点,这回弄成8个元素的张量。还是一行PyTorch代码:c = a + b。
你运行这行代码的时候,CPU(就是真正跑Python代码的那个家伙)会告诉GPU:嘿,去把这个程序跑在这堆数据上。这个指令就叫发射。发射出去的东西就是内核:一个准备好运行的独立程序。发射动作本身挺便宜,几微秒的事儿。真正花钱的是发射周围的事(把数据送到GPU,再把结果取回来),这些才是我们要数的。
内核里面,真正干活的是叫线程的小工人。一块GPU有成千上万个线程随时待命。对于这个长度为8的加法,8个线程领了活儿:线程0负责第0个元素,线程1负责第1个,一直到线程7。每个线程跑的是同一个迷你程序:读一个a里的元素,读一个b里对应的元素,加起来,把结果写进c。
(实际情况里,GPU会把线程组织成固定大小的组,N卡上叫线程束,一组固定32个线程。如果你的数组大小不能被32整除,它会屏蔽掉多余的线程。现在先不用管这个细节。)
所以,一行PyTorch代码,一次发射,一个内核,8个线程做了8次加法。现在来数数芯片上到底搬了多少数据。为了做加法,每个线程都需要它自己的那个a元素和b元素。那就是读了8次a,读了8次b。然后每个线程把结果写进c,又是8次写入。
这些读写操作,都是在GPU芯片旁边那个大块显存里进出的。数据中心用的显卡(A100、H100)上,这块显存叫HBM(高带宽显存)。普通消费级显卡(RTX 4090)或者大家在Colab上常用的T4,它叫GDDR。不管是啥,都是芯片旁边的高速显存,我们就直接叫它显存。速度确实快,但不是免费的,每次访问都得付出点代价。
一个内核,就是一次发射,就是对数据的一次完整处理。内核体内部干啥先不管,它在边界上的读写操作(去显存取输入,把输出写回显存)才是花钱的大头。整个格局就是这样。
张量变大,这套逻辑也不变。同样的PyTorch行,同一个内核,只是线程变多了。a和b要是各有一百万个元素,GPU就用更多线程启动同样的内核。计算量按比例增加,搬运的数据量也按比例增加,但思维模型没变。一行代码,一个内核。
两个操作之间藏着啥
c = (a + b).relu()
你懂Python,这显然是两个操作,先加法再ReLU,串起来的。在解释器里,这是两次函数调用。在GPU上,在即时执行模式的PyTorch里,这就是两次内核发射:一次给加法,一次给ReLU。到这儿都还好理解。
真正有意思的是在两个内核之间发生了什么。
加法跑完,它得把结果放到某个地方去。那个地方就是显存。加法把一整个中间数组(就叫它tmp吧)写进了显存。紧接着,ReLU内核启动了,它干的第一件事就是把刚才那个tmp数组从显存里再读回来。它把整个数组读完,每个元素过一遍ReLU,把结果写进c。
来数数这两个内核造成的显存流量:
- 加法:读a,读b,写tmp。三次数组规模的传输。
- ReLU:读tmp,写c。两次。
总共五次数组规模的传输。想想上一节里单独那个长度为8的加法,才三次。在代码后面加个.relu(),代价不仅仅是ReLU本身的运算。它还让你的数组在显存里多跑了个来回,因为tmp必须先被写出去,下一个内核才能再把它读回来。
这中间没有任何缓存。tmp没能留在寄存器或者局部缓存里。它直接去了显存(又慢又远的那种),然后又马上被读回来。这两个内核互相不认识对方。它们只能通过一个双方都会用的介质来交接数据:显存。
PyTorch为啥要这么干?因为在即时执行模式下,你写a + b的时候,PyTorch立刻就跑了。它不知道你下一句要调用.relu()。每个操作在Python那行执行到的瞬间就被分派出去了。没有计划,不看前面也不看后面。每个操作都独立运行,产生一个实实在在的数组,通过显存交给下一步。
记住这个模式。PyTorch代码里的每一个中间值,都会实实在在地被写进显存,再被下一个操作读回来。每一个都这样。“内核数量”真正衡量的就是这个。多一个内核,就是你的数据在显存里多跑一趟来回。
把两个活儿合成一个干
想象有一个内核,把这一切一步到位:读它自己的那部分a,读对应的b,加起来,对结果做ReLU(所有这些全在内核内部搞定,用每个线程自己那点儿从不离开芯片的临时空间),最后才把最终值写进c。中间的(a + b)还是存在,但它只活在内核里头,在每个线程自己的临时空间里。它从来没被写进显存。tmp这个实实在在的数组,压根就没存在过。
现在来数传输次数。读a:每元素一次。读b:每元素一次。写c:每元素一次。总共三次数组规模的传输。跟两个内核版本的数学计算量一样,但少了两趟来回。
长度是8的时候,这点差别可以忽略不计。没人会在意。但长度到一百万,或者一亿,多出来的那两趟来回就成了运行时间里相当大的一块,墙上挂钟的读数会明明白白地告诉你。*为啥*显存流量会这么要命,那是这个系列下一篇文章的主题,今天就不展开“为啥”了。现在只要记住:数学计算量一样,来回次数少了,实际跑起来就快了。
这个招数,把本来该分开跑的几个内核合并成一个,让中间结果不用去显存里溜达一圈,它有个名字,叫融合。就这一个词。就是这么个事儿。
尴尬的地方来了。手写这样一个融合内核,加ReLU看着挺简单。就俩操作,中间一行“计算”。但真正的PyTorch代码里,几十个操作串在一起,各有各的形状、数据类型、广播规则。手写一个能正确处理所有这些情况的融合内核,那是实打实的工程活。这种常规的逐元素内核,一般没人会自己动手写。
好消息是:PyTorch自带了一个工具,能自动帮你重写代码,专门处理这类情况。它叫torch.compile。
你肯定见过这么用torch.compile:
python
model = torch.compile(model)
就一行。网上有人告诉你这能让代码变快。它到底干了啥,用大白话说就是:不像即时执行模式那样一次一个操作地跑,torch.compile会先把你的函数里用到的张量操作都抓下来,看看有没有机会把它们合并,然后生成优化后的代码。以后调用如果符合同样的条件,就能复用这份代码。
我们刚才纸上谈兵做的那个融合(加法和ReLU共享一个内核,tmp从来不碰显存),就是torch.compile会自动给你的代码做的事儿,只要操作够简单。大家说torch.compile“让PyTorch变快”,这里头很大一部分功劳就来自于此。
碰上torch.compile自己没法融合的情况(它不认得的自定义操作、特殊的规约运算、奇怪的内存布局),那就还是得有人手写内核。这是Triton和CUDA这类工具的用武之地。那是另一篇文章的事了。
自己动手数一次
前面全是在纸上数内核。现在来真格的,在真正的GPU上数一次。只要手头有台带NVIDIA GPU的机器(工作站、Colab笔记本、云主机都行),几分钟就能自己跑一遍。
要用到的工具是torch.profiler。它是PyTorch自带的。干的活儿就是把你代码跑的时候GPU到底做了啥记录下来,然后给你一张表看。
第一步:即时执行版本。
把那两行操作包在一个函数里,方便调用:
python
import torch
from torch.profiler import profile, ProfilerActivity
def add_relu(a, b):
return (a + b).relu()
a = torch.randn(1_000_000, device="cuda")
b = torch.randn(1_000_000, device="cuda")
with profile(activities=[ProfilerActivity.CUDA]) as prof:
add_relu(a, b)
torch.cuda.synchronize()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
torch.cuda.synchronize()就是用来确保GPU干完活了,我们再读时间数据的。GPU干活是异步的,不加同步的话,有时候量到的是发射开销,不是真正的内核干活时间。
第二步:看输出结果。
实际的性能分析器输出会比预想的多很多行。会混进来一堆显存分配和PyTorch内部记账的行。我们要关注的是CUDA内核,也就是GPU真正跑的那些函数。找名字里带kernel的行。这俩看着大概长这样:
text
vectorized_elementwise_kernel<...BinaryFunctor_add...> 1 12us
vectorized_elementwise_kernel<...threshold_kernel...> 1 9us
具体名字在不同PyTorch版本里会变(ReLU经常显示成threshold,因为底层操作就是阈值操作,加法有时候显示成CUDAFunctor_add)。别费劲去解析整行。数行数就行。两行。两个内核。一个给加法,一个给ReLU。跟上节说的一模一样。
第三步:编译版本。
改一行就行。把函数用torch.compile包起来:
python
compiled = torch.compile(add_relu)
这里有个坑,跑之前最好知道一下。*第一次*调用torch.compile包过的函数会慢,有时候慢得离谱,因为那是torch.compile在干活:分析你的代码,琢磨哪些能融合,生成融合后的内核。要是把第一次调用也放在性能分析里头,你量到的是编译时间,不是内核时间。所以标准操作是:先调用一次做预热,结果丢掉,*然后*再开始性能分析。
text
compiled(a, b) # 预热,结果扔掉
torch.cuda.synchronize()
with profile(activities=[ProfilerActivity.CUDA]) as prof:
compiled(a, b)
torch.cuda.synchronize()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
第四步:再看一次输出。
这次:
text
triton_poi_fused_add_relu_0 1 14us
一行。一个内核。名字甚至直接告诉你它干了啥:融合的加法和ReLU。数学计算量跟之前一样,但发射从两次变成了一次。
你刚刚亲手做了这篇文章一直在聊的事,一句话概括就是:你让PyTorch把两个操作合并成了一个内核,盯着性能分析器看,确认内核数量从二变成了一。融合,在真实环境里,在你自己的机器上。
想看得更明显点,可以换几个不同的张量尺寸试试。长度100的时候,两个版本都快得没影,差别都淹没在噪声里了。长度到一千万或者一亿,编译版本的优势就开始明显了,因为我们省掉的那趟来回在这么大的数据量下已经是实实在在的一笔开销了。
数内核不再是纸上谈兵的建议了。你现在有办法亲自验证了。
收个尾
把整个事儿串起来说一遍。
你的PyTorch代码在GPU上跑的时候,会变成一串内核。每个内核就是一次发射,对数据做一次完整处理,去显存里取输入、写输出,跑一趟来回。简单的操作就是一个内核。串起来的操作默认情况下是*每个操作*一个内核,中间的临时结果在显存里来来回回地倒腾。torch.compile能帮你把简单的操作链融合起来,让那些中间结果压根不用碰显存。内核少了,通常意味着显存流量少了。显存流量少了,通常意味着跑得更快。
到头来,GPU不关心你的模型设计得多精巧,它就认内核和显存读写这两件事。调优的本质,说白了就是让数据少在显存里遛弯。