黑威尔Blackwell架构把GPU矩阵乘法从"一群人一起算"变成了"一个人下单、一群人收快递",张量核心(Tensor Core)自带独立内存、异步执行、还能内存泄漏,NVIDIA B200正在重写整个CUDA编程范式。
以前128个线程挤在一起等结果,现在一个人扔个任务就跑了。但收账的时候又得喊回来一整个连队。这操作像极了食堂打饭——你一个人冲过去把碗一放说“我要饭”,然后跑回座位玩手机,等饭好了得四个哥们一起去端。
问题是如果碗不还,下一波人直接饿死在窗口。
寄存器文件锁死了整整七年
GPU芯片里有个叫寄存器文件(register file)的仓库,专门给线程们临时存数据。从2017年开始,这个仓库的容量就被焊死在了256KB,整整七年没涨过。七年是什么概念?够一个人从初一读到高三毕业,够iPhone从7换到15,够比特币从一千美元飙到六万再跌回来。但NVIDIA的寄存器文件就卡在那儿,一动不动,像被施了定身咒。
与此同时,张量核心(Tensor Core)这匹野兽却在疯狂长肉。从2017年到现在,它的算力大概翻了十八倍。十八倍是什么感觉?以前你骑一辆共享单车,现在突然给你换了一辆高铁。问题是高铁得吃电,张量核心得吃数据。为了让这头十八倍大的野兽吃饱,程序员们不得不把每次喂进去的分块(tile)越切越大。分块就像喂老虎的肉块,肉块小了老虎懒得张嘴,算力就浪费;肉块大了才能塞满它的牙缝。
所以这就形成了一个荒诞的局面:仓库容量冻住了,但送货量翻了十八倍。程序员们只能拼命把每块肉切得更大,希望一次喂饱。可肉块越大,临时占用的仓库空间就越多。这就像一个快递员,面包车容量没变,但客户要求一次送十八倍的货。他唯一能做的就是每次少跑几趟、每趟多塞一点,直到车门彻底关不上。
一块矩阵吃掉了半个仓库
到了Blackwell这一代,荒诞变成了灾难。一个128乘256的FP32(单精度浮点)分块,光结果数据就占128KB。128KB是什么概念?它直接吞掉了单个流式多处理器(SM)寄存器文件的一半。一半啊,兄弟们。你开个奶茶店,房租占一半收入还能活;但一个临时计算结果就占你仓库的一半,这生意还怎么做?
更惨的是,你连双缓冲(double buffer)都玩不起了。双缓冲是GPU编程里的经典套路:一边算着上一块,一边准备着下一块,两班倒不歇火。但现在一块肉就占半个仓库,你放两块就直接爆仓。程序员们站在仓库门口,手里拿着第三块肉,看着满当当的货架,只能干瞪眼。寄存器文件彻底没房间了,连地板上都堆满了。
NVIDIA的工程师们肯定也急眼了。他们看着这个256KB的紧箍咒,看着张量核心饿得嗷嗷叫,看着程序员们把分块切得越来越大却塞不进去。这就像一个餐厅,厨房越来越大、厨师越来越快,但冰箱还是七年前的老型号。菜炒得再快,没地方放也是白搭。到Blackwell这里,冰箱门彻底被肉块顶开了。
NVIDIA给张量核心开了独立账户
既然仓库塞不下,那就再建一个仓库。NVIDIA的解决方案简单粗暴:给张量核心(Tensor Core)单独开了一张银行卡。他们把计算结果从线程们的寄存器文件里抽出来,转给每个张量核心配了一块专属内存:每个流式多处理器(SM)上划了256KB,独立的地址空间,只有张量核心能读写,普通线程连门都摸不着。
这个操作相当于什么呢?以前公司所有员工共用一个储物柜,现在给技术部单独配了一个保险柜。保险柜里的东西只有技术部能碰,其他部门的人站在旁边干看着。这个保险柜和储物柜是两套系统,互不相通。张量核心算完的东西往保险柜里一扔,线程们该干嘛干嘛,不用集体守在那儿当仓库管理员。
一旦结果不再赖在线程们的寄存器里,事情就发生了质变。以前发射一次矩阵乘法(matmul)需要128个线程排排站,像拔河一样一起发力。因为结果要存在他们的寄存器里,所以必须凑够人手才能把结果兜住。现在结果直接进保险柜了,那还要那么多人干嘛?发射矩阵乘法需要的线程数,直接从128个塌缩到了1个。对,一个人,一个线程,点一下按钮,完事。
一个线程就能发射矩阵乘法
这就是Blackwell上矩阵乘法的全新打开方式。一个线程走上前来,先malloc一块张量核心内存,就像去保险柜申请一个抽屉。然后它掏出一条tcgen05.mma指令,啪地拍下去,矩阵乘法就发射了。这条指令是异步(asynchronous)的,什么意思?它跳过傻站着等结果的环节,直接把任务单往张量核心脸上一拍,扭头就走,继续干自己的活。
结果开始源源不断地涌进那块专属内存,在整个内层循环里堆着、叠着、攒着,全程不碰任何寄存器。线程们该搬数据的搬数据,该做别的做别的,张量核心在保险柜里默默干活。等整个分块(tile)算完了,那个线程再过来,free掉这块内存,抽屉清空,收工下班。整个过程就像点外卖:你下单、骑手送餐、你吃饭、扔垃圾。你不需要全程盯着厨房看厨师炒菜。
这种异步模型彻底改写了GPU编程的直觉。以前你觉得矩阵乘法是一群线程齐心协力的大工程,现在它变成了一条后台指令。以前结果要过线程们的手,现在它直接走专线。以前算完要大家一起收拾残局,现在一个人点个free就搞定。这个转变之大,就像从集体劳动包产到户,突然变成了一个人承包一片地。
矩阵乘法也会内存泄漏
但新权力带来新麻烦。既然内存是手动分配、手动释放的,那就一定会有人忘记释放。在Blackwell上,如果一个线程算完矩阵乘法,忘了free那块张量核心内存,会发生什么事?下一个被调度到同一个流式多处理器(SM)上的线程块(thread block)就会直接卡死,像一辆车开进死胡同,前面堵着一辆抛锚的卡车,退也退不得,进也进不得。
这是货真价实的内存泄漏(memory leak),而且发生在矩阵乘法这种底层运算里。以前内存泄漏是C++程序员的专利,是malloc了没free、new了没delete的那种高级失误。现在NVIDIA告诉你,跑个矩阵乘法也能漏,而且漏得更隐蔽、更致命。因为这压根儿不属于普通内存的范畴,它归张量核心专属,普通检测工具根本看不见。
NVIDIA自己也知道这事有多离谱,所以他们直接ship了一个检测器(sanitizer)。打开这个开关,运行时一旦发现有内存没释放,立刻trap住,精确告诉你哪一列、哪一块、哪个位置漏了。这跟普通代码里的use-after-free检测器一个套路,只不过现在抓的对象早就换了,从指针乱飞换成了矩阵乘法的尾巴没收拾干净。NVIDIA给矩阵乘法写内存泄漏检测器,这画面本身就够surreal的。
最诡异的不对称出现了
如果说前面的变化已经够让人脑壳疼,那这部分简直就是行为艺术。在Blackwell上,发射一个128乘256的矩阵乘法只需要一个线程。但要把结果从张量核心内存里读出来,需要128个线程。一个人点火,一百二十八个人灭火。这不对称得令人发指,就像你一个人按了电梯按钮,需要全楼的人一起才能把电梯门拉开。
原因在于张量核心内存被切成了四等份,每个线程束(warp)只能看到自己的那一份。一个线程束有32个线程,四份就是128个线程。所以draining一个结果需要一个完整的warpgroup(线程束组,128线程)。你发射的时候像个独行侠,收结果的时候得像合唱团一样列队进场。这种设计让人想起那些古老的寓言:一个人种下一棵树,一百个人才能把它砍倒运走。
这种asymmetric属于feature级别的设计,NVIDIA故意这么干的。NVIDIA故意这么设计的,因为张量核心内存的带宽和访问模式跟寄存器文件完全不同。它更像一个专用的高速缓存,跟通用仓库已经划清界限。所以读取它需要特定的队列、特定的宽度、特定的人数。你一个人能把任务踢出去,但要把成果搬回家,必须组团。这就像寄快递,下单一个人就行,但收件得全家出动搬大件。
六个线程束演起了流水线话剧
把这些碎片拼起来,一个高速核函数(kernel)在Blackwell上长什么样?它已经蜕变为六个线程束(warp)各司其职的流水线话剧,不再是三十二个线程跑同一套代码的整齐方阵。第一个线程束负责把数据搬进张量核心内存,像仓库的进货员;第二个线程束专门发射tcgen05.mma指令,像按按钮的操作员;剩下四个线程束负责把结果从张量核心内存里drain出来,像打包发货的工人。
六个演员,三种剧本。进货员跑的是搬运代码,操作员跑的是发射代码,工人们跑的是读取代码。他们通过硬件邮箱(hardware mailboxes)互相传纸条:货到了、可以算了、结果好了、该搬了。没有任何两个线程束在执行同一行代码,没有任何一个线程从头到尾跑完整流程。这跟你印象中的GPU编程完全两码事。以前GPU是"众人拾柴火焰高",现在是"流水线分工各干各的"。
这种架构的精妙之处在于,它把张量核心从"大家一起调用的功能"变成了"独立运转的单元"。线程们早就改了站队方向,从围着张量核心转改成了围着它服务。它有了自己的内存、自己的指令、自己的工作节奏。线程们变成了它的秘书、司机和保洁,跟它同事这层关系已经翻篇了。这个权力结构的翻转,是Blackwell最深层的设计哲学。
张量核心从指令变成了协处理器
归根结底,张量核心(Tensor Core)在Blackwell上完成了一次身份蜕变。它已经告别了线程束(warp)里一条指令的身份,程序员写一行代码就能召唤它的日子也过去了。它变成了一个坐在线程旁边的独立协处理器(coprocessor),有自己的内存空间、自己的指令集、自己的执行节奏。程序员要手动给它分配内存、手动给它dispatch任务、手动给它清理现场。
这个转变的意义远超一次架构升级!它意味着GPU内部出现了两级计算实体:
- 线程们负责通用计算和协调
- 张量核心负责专用加速。
两者之间通过内存分配和硬件邮箱通信,绕开了寄存器共享的老路子。这就像一台电脑里突然多了一块独立显卡,只不过这块"显卡"是嵌在GPU芯片内部的,而且比主处理器还挑食、还难伺候。
NVIDIA用Blackwell告诉全世界:矩阵乘法(matmul)已经升级为"养"出来的,"算"这个字已经罩不住它了。你要先给它建窝、喂数据、发指令、等它算完、再帮它收拾。它从仆从变成了主子,从工具变成了搭档。而这个搭档还有洁癖:算完不收拾它就闹脾气,下一位客人直接堵在门口。这种娇贵又强大的存在,就是Blackwell张量核心的真面目。