专栏名称: 大模型智能
机器学习算法、深度学习算法、自然语言处理等干货知识集中营
目录
相关文章推荐
新华每日电讯  ·  在“世界工厂”东莞,看见“素人写作” ·  昨天  
为你读诗  ·  绝美奥斯汀刷边版特装版,值得收藏 ·  昨天  
51好读  ›  专栏  ›  大模型智能

AndrejKarpathy盛赞!斯坦福团队新作,让Llama-1B实现毫秒级推理

大模型智能  · 公众号  ·  · 2025-06-02 10:00

正文

请到「今天看啥」查看全文


先说实验结果,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,大幅减少了指令之间等待的浪费,使得整个内核执行尽可能地接近理论并发。

图片

除此之外,研究团队还对 CUDA 异步屏障(asynchronous barrier)的性能进行了测量,发现即便在 barrier 已“pass”的状态下,每次仍需 60ns,同步操作成本不可忽视。而在实际执行中,尤其在 matrix-vector(矩阵乘向量)这类关键操作中,他们发现:在 Hopper 架构(如 H100)上,使用常规 CUDA 核心(非 Tensor Core)可以更有效,不过在 Blackwell 架构上,Tensor Core 性能占优。

这也说明在硬件不同世代中,Megakernel 的最佳实现路径也应适配微架构差异,而非一套方案通用所有平台。







请到「今天看啥」查看全文