不换显卡提速1.68倍:英伟达工程师如何把2.8万亿参数大模型推上每秒423词
一个 2.8 万亿参数的 AI 模型,一秒钟能给一个用户吐出 423 个词。
一秒钟 423 个词,是个什么概念? 人类平时说话,一秒钟大概三四个词。 普通人看书扫得快一点,一秒钟七八个词。 423 个词,相当于眼睛眨一下,大半篇高考作文就直接糊在屏幕上了。
更绝的是什么? 跑出这个速度的,可不是几十亿参数的轻量小模型。 它是月之暗面那个 2.8 万亿参数的庞然大物,Kimi K3。 而且服务器里的显卡一颗螺丝都没换。 硬件一模一样,单纯靠英伟达工程师写的一套新代码,速度直接翻了 1.68 倍。
显卡没加,硬件没变,它怎么就能突然跑得像坐了火箭一样?
先来看这件事有多难。 Kimi K3 采用的是混合专家架构(MoE),总参数高达 2.8 万亿。 哪怕每次只激活其中的一小部分,每个词也要调动大约 1040 亿参数。 大模型回答问题,在生成阶段是一颗词接一颗词往外蹦的。 每蹦出一个词,显卡就必须把上千亿参数从显存里完整搬运一遍。 算力再强,显存带宽不够快,核心就在那里干等。 这就是搞大模型推理的人天天头疼的「内存墙」。
但如果只是显存慢,事情还好办。 真正让速度卡死的,是显卡里的交通大堵塞。
在开源推理框架 vLLM 原来的代码里,模型每生成一个词,要经过十几道甚至几十道工序。 先算注意力,再做路由打分,再分配给对应的专家网络计算,再把十几张显卡上的结果汇总合并。 在显卡里,每道工序都是一个单独的算子。 这就像一个流水线车间。 工人每拧一颗螺丝,就得把零件放下,在纸上打个勾,打个电话通知下一个工位。 下一个工位跑过来,把零件抱走,再拧下一颗螺丝。 十六张英伟达 GB300 显卡连成集群在跑,由于算子切得太碎,加上跨显卡通信的等待,显卡里全是闲置等待的气泡。 大部分时间,价值几十万的计算核心根本没在算题,全在等零件、等指令。
那这次英伟达是怎么把这个瓶颈敲碎的?
英伟达 CentML 的工程师 Benjamin Chislett 干了一件极其硬核的事。 他用英伟达自家的底层语言 CuTe DSL,手写了一个长达一万两千行的超级算子。
什么叫超级算子? 简单来说,就是把原来流水线上几十个散落的工位,彻底焊死在了一条全封闭的高速传输带里。 注意力计算、专家路由、多卡通信,全部揉进同一个内核。 显卡启动一次,所有计算线程常驻在芯片内部。 中间产生的数据,根本不需要大老远存回高带宽显存,直接在芯片内部极速的片上缓存和寄存器里流转。 更绝的是跨卡通信。 十六张显卡之间倒腾数据的等待时间,直接被塞进计算缝隙里掩盖掉了。 原来把流水线堵死的那些空转等待,被活生生挤了个干净。
新算子合并进 vLLM 之后,结果立刻出来了。 同样是十六张 GB300,同一个 2.8 万亿参数的巨型模型。 单用户的生成速度,直接从原先的基础水平飙升到了每秒 423 个词,整整提速了 1.68 倍。 如果再配合投机解码这种技巧,速度甚至能冲到每秒 800 多个词。
这件事情最耐人寻味的地方,恰恰在这里。 平时大家聊到大模型加速,第一反应永远是买更多更贵的显卡,或者拼命压缩模型参数。 但这次 Benjamin 的这一万两千行代码给出了另一种答案: 哪怕是最顶级的硬件,由于软件调度和数据搬运的内耗,原本的实力也常常被封印了大半。 不用换硬件,把算子切碎重新排兵布阵,就能凭空把巨无霸模型的速度往上推一大截。
同领域解读 / AI & LLMS
导出 / EXPORT
订阅 · SUBSCRIBE
每周一封独立、不受审查的科技与 AI 深度评论,周一发出。带出处,可核查。