斯坦福 Hazy Research 团队刚刚发布了一项重量级优化作用 :他们将开源模型 Llama-3.2-1B 的盛赞斯坦前向推理整合成了一个“Megakernel”,并将低推迟推理才干面向了极限 。福团
在某些实时性极高的队新运用中 ,例如对话式 AI 和人类参加的作让交互式工作流中,大言语模型的盛赞斯坦呼应速度不只重要 ,乃至能够决议用户体会的福团胜败。
团队以为约束 LLM 推理速度的队新瓶颈其实是在内存加载的问题上,他们经过研讨发现,作让现有的盛赞斯坦开源推理引擎(如 vLLM 、SGLang) ,福团在极低推迟的队新单序列生成使命下,即使在尖端 GPU(如 H100)上,作让也只能运用不到 50% 的盛赞斯坦内存带宽。
这首要是福团由于每层 Transformer 模块被拆解成几十到上百个 CUDA kernel ,每个 kernel 履行十分小的队新操作(比方 RMS norm 、注意力、MLP、Rotary Position Embedding 等) ,它们之间存在许多上下文切换与等候 。
更严峻的是,这些 kernel 发动与收尾的本钱加起来,并不会被 CUDA Graph 或 PDL(Programmatic Dependent Launch)等机制充沛躲藏,反而在短时使命中被扩大。
换句话说,GPU 花了许多时刻“等着干活” ,而不是“在干活”。Hazy 团队的研讨也正是围绕着这个问题打开。
Megakernel:从零规划的交融思路 。
先说试验成果,Megakernel在 H100 上的推理推迟紧缩至缺乏 1 毫秒,显存带宽运用率高达 78% ,相较于 vLLM 提高了 2.5 倍 、相较 SGLang 提高 1.5 倍;而在更先进的 B200 渠道上,推迟进一步下降至 600~680 微秒,迫临理论极限 。
从一次完好推理的时刻散布来看 ,250 微秒用于存储激活、等候一致性与数据加载 ,200 微秒用于 RMSNorm 与 matvec(其间 matvec 占比达 95%),权重加载仅需 30 微秒,流水机制体现安稳 。warp 间同步与 barrier 带来 40 微秒的推迟 ,其他如 setup 、参数传递与页状况符号等杂项开支算计约 80 微秒 。
全体来看 ,在精心调度下 ,Hazy 团队的 Megakernel 简直已将当时硬件功能压榨至极限 。
而能够得到以上作用,其实都归功于 Hazy 团队提出的一个急进但有用的规划思路:将整个前向传达进程整合为一个单一 CUDA kernel ,也便是所谓的 Megakernel。
试验中 ,他们根据已有 ThunderMLA 架构 ,开发了一个 GPU 上运转的轻量“指令解说器”体系。该体系为每个 Streaming Multiprocessor(SM)预先分配一段“履行计划” ,其间包括多条按顺序排列的指令,每条指令代表 Transformer 模型中的一个结构单元。
这些指令包括 :
交融 RMSNorm 、QKV projection 、RoPE 的复合指令;
attention 矩阵乘与减缩核算(支撑长序列 GQA);
O-projection 与 residual 相加;
MLP 的 RMSNorm、gate 激活(SiLU)与上投影;
down projection 和终究 residual;
最终一层 RMSNorm + language modeling head。
每个指令都根据一致的 CUDA 模板构建 ,完结对 load 、store 、compute 的标准化封装 。指令间依靠由解说器在运转前静态排布,每个 SM 能够重复复用同一个 schedule 以处理多个 token。
此外 ,为确保高效的数据途径 ,解说器会将这些履行计划按模型结构静态编列 ,避免调度时动态分支,提高吞吐与并发履行才干 。
一同为了完结流水化核算并避免 shared memory 抵触,团队还对 GPU 的同享内存进行了分页办理 ,例如 :
将前 213KB 的 shared memory 分为 13 个 16KiB 页面;
剩下部分用于存储指令参数 、页分配信息等;
每条指令在加载前显现恳求页,完毕后归还给解说器调度器;
当页被开释时,解说器会当即将其分配给下一条等候中的指令。
这种机制确保了下一个核算阶段能够尽早开端预加载权重,然后最大化带宽运用率并消除“气泡” 。
不过 Megakernel 结构无法依靠传统的 kernel 间隐式同步,因而 Hazy 团队还运用了一个计数器体系:他们在 global memory 中保护一组整数 ,每条指令完结后会对对应计数器 +1 ,若某条指令依靠从前进程的成果,它会等候计数器到达特定值才履行。
例如:在 MLP 下投影阶段,团队将中心态拆成 4 个 chunk ,每个 chunk 在写入后当即触发后续核算 ,然后完结并行流 。此外,团队经过准确设置依靠图,避免了大局 barrier ,大幅减少了指令之间等候的糟蹋 ,使得整个内核履行尽或许地挨近理论并发。