第 7 讲 · 高级 CUDA：Mega-Kernel · Advanced CUDA Programming: Mega Kernel

对应 Maps to：CMU 15-779（Fall 2025，Zhihao Jia），Lecture 7 · Fri 9/19 · Advanced CUDA Programming: Mega Kernel。 必读材料：B. Spector, J. Juravsky 等（HazyResearch），Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B（博客，2025-05-27）；Zhihao Jia 等，Compiling LLMs into a MegaKernel: A Path to Low-Latency Inference（博客，2025 年 6 月，介绍 Mirage Persistent Kernel，MPK）。本讲还引用：X. Cheng, Z. Zhang, …, T. Chen, Z. Jia, MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs（arXiv 2512.22219）。 依据说明：本讲以原课 slides 07-mega-kernel.pdf 的主线组织：内核是什么；逐算子内核的局限（启动开销、层间无法流水、计算与通信无法重叠）；mega-kernel 的优点与挑战；MPK 的任务图（与 CUDA Graph 的区别）、编译器、运行时（GPU 内核里的任务并行运行时）、降低任务启动开销的技巧；把 LLM 推理时延推向硬件极限；开放问题。HazyResearch 的数字引自其博客；MPK 的数字引自 MPK 的 arXiv 论文和 GitHub README（MPK 博客原文在编写时无法访问，见代码 README）。执行模型 megasim.py、其中的参数和全部实验都是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture07/（3 个脚本 + 共享模块 megasim.py；只用 Python 标准库和 numpy；每个脚本 < 1 秒） 预计阅读时间：3 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标

学完这一讲，你应该能够：

1. 说清楚”内核”在 GPU 执行模型里的含义，以及内核边界隐含了什么（启动开销、全局同步、资源重新分配）。
2. 给批 1 解码算一笔账：读多少字节、理想多快；把逐算子内核的额外时间分解成启动间隙、延迟、尾部效应三项。
3. 解释 mega-kernel 的核心想法（一个常驻内核 + GPU 内部的任务调度），以及它为什么能消除这三项：特别是”权重和激活值无关，可以在等依赖时提前读”。
4. 说出 mega-kernel 的主要挑战：GPU 内部的调度与同步、资源分配、代码生成、调度开销。
5. 描述 MPK 的设计：以 SM 为粒度的任务图（任务 + 事件）、从读写区域推导依赖、worker 与调度器、AOT 与 JIT 两种任务启动方式；并和 CUDA Graph 比较。
6. 读懂 HazyResearch 与 MPK 的实验数字，并能判断一个简单模型的结论在哪些地方过于乐观。

0.2 先修与本讲范围

0.3 本讲地图

§1  内核是什么，"一个算子一个内核"是怎么来的
§2  小批量解码的账：读权重就是一切                       ← 代码 01（演示 1）
§3  内核边界的三笔代价：启动间隙、延迟、尾部              ← 代码 01（演示 2–5）
§4  mega-kernel：一个常驻内核 + GPU 里的调度器；挑战        ← 代码 02（演示 1–2）
§5  HazyResearch 的 Llama-1B mega-kernel                   ← 代码 02（演示 3–5）
§6  MPK：任务图、编译器、运行时                            ← 代码 03
§7  把推理时延推向硬件极限：哪些开销还剩下？
§8  开放问题
§9  论文精读：HazyResearch 博客；MPK
§10 小结  §11 常见误区  §12 练习  §13 延伸阅读  §14 下一讲预告

0.4 关于模型的数字

megasim.py 把一次解码步拆成任务（一个 SM 上的一小块工作），在 132 个 SM 上用两种方式执行。参数：

这是模型：它对两种执行方式都偏乐观。讲义会把模型结果和博客、论文的实测并排放。

1. 内核是什么 · What is a kernel

第 3 讲：内核（kernel）是从 CPU 上启动、在 GPU 上由一个线程块网格执行的函数。一次内核启动包括：CPU 把参数和网格配置交给驱动 → 驱动把命令放进 GPU 的队列 → GPU 的调度器把线程块分给 SM → 所有线程块执行完 → 内核结束，下一个内核才能开始（同一个流里）。

内核边界隐含了三件事：

1. 启动开销：每次启动都要经过 CPU、驱动、GPU 前端，即使内核什么都不做也要几微秒；
2. 全局同步：同一个流里，下一个内核要等上一个内核的所有线程块结束——这是 CUDA 里最便宜的”全局屏障”（第 4 讲的归约就靠它分解）；
3. 资源重新分配：每个内核有自己的寄存器数、共享内存大小，内核结束后共享内存里的数据全部丢失。

深度学习框架很自然地采用一个算子一个内核：PyTorch 的每个算子（矩阵乘法、RMSNorm、SiLU……）调用一个预先写好的库内核（cuBLAS、cuDNN 或自己的 CUDA 代码）。第 2 讲的算子融合把几个算子合并成一个内核，第 5 讲的 FlashAttention 把注意力的三个内核合成一个。mega-kernel 把这个思路推到极致：整个模型只有一个内核。

2. 小批量解码的账 · The cost of small-batch decoding

第 5 讲 §8：解码批 1 时每个权重只用一次，算术强度约为 1，完全是访存受限。所以一次解码步的理想时间就是”把所有权重读一遍”的时间。

代码 01 演示 1 以 Llama-3.2-1B 为例（HazyResearch 博客的例子）：

每层切成的任务（每个任务是一个 SM 上的一小块工作，例如”QKV 投影的 16 行”）：

一次前向一共约 4 万个任务。逐算子执行时是 114 个内核（每层 7 个：RMSNorm、QKV、注意力、输出投影、RMSNorm、gate/up、down，再加最后的 RMSNorm 和 LM head）——HazyResearch 博客说主流引擎把一次前向拆成”约一百个”内核，量级一致。

3. 内核边界的三笔代价 · Three costs of kernel boundaries

3.1 一层里的每个内核（代码 01 演示 2）

这一层理想 36.9 µs，实际 56.5 µs（65%）。三笔代价：

(1) 启动间隙。 每个内核之间约 1.3 µs（CUDA Graph）到 2.1 µs（普通 CUDA stream）。HazyResearch 在 H100 上测量空内核：CUDA stream 上约 2.1 µs，用 CUDA Graph 约 1.3 µs。MPK 论文的测量：Qwen3-8B 每个 token 有 293 次内核启动，直接启动每次 3.8 µs（合计 1.11 ms），用 CUDA Graph 每次 0.8 µs（合计 0.23 ms）。

(2) 延迟（开头和结尾）。 一个内核开始后才能开始读它的权重——要先付一次访存延迟才有数据；结束时要把结果写出去，下一个内核才能看到。权重和激活值无关，本可以在上一个内核还在算的时候就开始读，但内核边界不允许。每个内核约 1.1 µs，114 个就是 125 µs。

(3) 尾部效应。 QKV 有 192 个任务、132 个 SM：第二轮只有 60 个 SM 有活干，其余 72 个空等。理想 3.76 µs 的读取用了 6.26 µs。下一个内核必须等这个内核最慢的线程块结束。这是第 3 讲”波次量化”的又一次出现。

小内核（RMSNorm、解码时的注意力）更糟：几乎全是固定开销。

3.2 整个解码步（代码 01 演示 3）

CUDA Graph 和融合都有效，但只要还有内核边界，三项开销就不会消失。

这个模型对逐算子内核相当宽容：每个内核都是理想的 GEMV，没有 CPU 开销。HazyResearch 博客报告，vLLM 和 SGLang 在这个场景下最多只用到 H100 一半的带宽。

3.3 硬件越快，边界越贵（代码 01 演示 4）

换成 B200（148 个 SM、8 TB/s），理想时间从 748 µs 降到 313 µs，但每个内核的固定开销（启动、延迟）不变：融合后的逐算子内核只用到 58% 的带宽（H100 上 76%）。这和 HazyResearch 的实测趋势一致：他们的 mega-kernel 在 B200 上比 vLLM 快 3.5 倍以上，在 H100 上约 2.5 倍。硬件带宽每提升一代，固定开销的占比就更大。

3.4 第四个局限：计算和通信无法重叠

多 GPU 推理（张量并行，第 14 讲）时，每层之后要做一次 all-reduce。逐算子执行中，all-reduce 是一个独立的内核（或一次 NCCL 调用），必须等前面的矩阵乘法全部结束才能开始；而后面的计算又要等 all-reduce 全部结束。实际上，矩阵乘法每算完一块输出，那一块就可以开始传输了。MPK 把通信也拆成任务，用同样的事件机制和计算交错（§6）。

3.5 已有的部分补救

• CUDA Graph：把一串内核的启动”录”下来一次性重放，降低启动开销（但仍有间隙，而且图基本是静态的：形状、参数变化时要重新录制或准备多份）。
• 算子融合：减少内核数（RMSNorm 融合进矩阵乘法），但融合的范围受限于手写或编译器能生成的内核。
• PDL（Programmatic Dependent Launch，Hopper）：允许下一个内核在上一个内核结束前就启动、做准备工作（例如读权重），在读上一个内核的输出之前再同步。HazyResearch 指出它有帮助，但同步粒度太粗：例如注意力仍要等全部 Q、K、V 都算完。

4. Mega-kernel · The idea

4.1 想法：一个常驻内核 + GPU 里的调度器

把整个前向编译成一个内核：启动时每个 SM 上放一个常驻的 worker（第 6 讲的 persistent kernel 推广到整个模型），worker 在一个循环里不断取任务、执行任务。任务之间的依赖用全局内存里的计数器表示：任务做完时给计数器加 1，依赖它的任务等计数器到达目标值再开始。

while (还有任务) {
    task = 我的队列.next();
    预取 task 的权重到共享内存（不需要等依赖！）;
    等待 task 依赖的计数器到位;
    读激活值、计算、写出结果;
    atomicAdd(task 触发的计数器, 1);
}

4.2 它怎样消除三笔代价

• 没有启动间隙：只启动一次。
• 延迟被藏住：权重和激活值无关。worker 在等上一个算子的结果时，已经在读下一个任务的权重（受共享内存容量限制，例如 13 页 × 16 KB、一个任务占 4 页 → 最多提前读 3 个任务）。
• 尾部被填平：一个 SM 先做完自己那份 QKV，不需要等别的 SM，直接开始读注意力或输出投影的权重。不同的 SM 在同一时刻可能在处理不同的算子——跨算子的流水线。

代码 02 演示 1–2 的结果（同一个模型、同样的参数）：

第 0 层前 60 µs 的时间线（每个字母表示这个 SM 正在读哪个算子的数据，Q = QKV、A = 注意力、O = 输出投影、G = gate/up、D = down，· = 没在读）：

逐算子内核
    SM0   |··QQQQQQ···A···OOO···GGGGGGGGGGGGGGGGGGGGGGGGG···DDDDDDDDDDDD···QQQQQQ··|
    SM60  |··QQQ······A···OOO···GGGGGGGGGGGGGGGGGGGGGGGGG···DDDDDDDDDDDD···QQQ·····|
    SM131 |··QQQ················GGGGGGGGGGGGGGGGGGGGGG······DDDDDDDDD······QQQ·····|
mega-kernel
    SM0   |QQQQQQAOOOGGGGGGGGGGGGGGGGGGGGGGGGGDDDDDDDDDDDDQQQAOOOGGGGGGGGGGGGGGGGGG|
    SM60  |QQQAOOOGGGGGGGGGGGGGGGGGGGGGGGGGDDDDDDDDDDDDQQQQQQOOOGGGGGGGGGGGGGGGGGGG|
    SM131 |QQQAOOOGGGGGGGGGGGGGGGGGGGGGGGGGDDDDDDDDDDDDQQQAOOOGGGGGGGGGGGGGGGGGGGGG|

逐算子内核在每个算子之间都有一段所有 SM 都空着的缝；QKV 的第二轮只有一部分 SM 在读（SM60、SM131 只读了一半时间）。mega-kernel 的每一行几乎都是连续的字母。

4.3 挑战

听起来很美，但 mega-kernel 很难写：

1. GPU 内部的调度：谁决定哪个任务在哪个 SM 上、什么时候执行？调度本身的开销必须远小于一个任务（§5.3、§6.3）。
2. 同步：依赖靠全局内存里的计数器和原子操作实现，要处理内存一致性（写出的数据必须在计数器更新之前对其他 SM 可见）。
3. 资源分配：一个内核只有一套寄存器数和共享内存配置，必须满足最苛刻的那个任务（例如注意力）；其他任务只能共用这套配置。HazyResearch 用”共享内存分页”让不同的任务按需申请。
4. 代码生成：每种任务都要写成能在同一个内核里运行的设备函数；模型一变就要重写——所以需要编译器（MPK）。
5. 调试与性能分析：一个内核里发生了几万个任务，传统的按内核分析的工具不再适用。

5. HazyResearch 的 Llama-1B mega-kernel · Look Ma, No Bubbles

5.1 设计

HazyResearch 手写了一个 Llama-1B 的 mega-kernel，设计要点（根据博客）：

• 解释器 + 指令集：每个 SM 执行一串”指令”，所有指令用同一个 CUDA 模板实现。调度表提前在 Python 里排好，之后反复用于几百次前向。一共 7 种指令：融合的 RMSNorm + QKV + RoPE；注意力计算；注意力归约（长序列时）；输出投影 + 残差；融合的 RMSNorm + up/gate + SiLU；down 投影 + 残差；RMSNorm + LM head。
• 共享内存分页：H100 上前 213 KB 的共享内存分成 13 页、每页 16 KiB；指令按需申请和释放页，解释器把释放的页立即交给下一条指令，让它尽早开始加载。
• 计数器同步：全局内存里一组初值为 0 的计数器；每条指令完成时加 1，依赖它的指令等到目标值。MLP 的中间结果分成 4 块、各有一个计数器，down 投影只等它需要的那一块。

5.2 结果（博客）

• H100：一次前向不到 1 ms，用到 78% 的显存带宽；比 vLLM 快近 2.5 倍、比 SGLang 快 1.5 倍以上（两者都开了 CUDA Graph 和 torch.compile）。博客称这是他们所知第一次在 GPU 上用不到 1 ms 跑完 16 位、10 亿参数以上语言模型的前向。
• B200：一次前向不到 680 µs，比 vLLM 快 3.5 倍以上、比 SGLang 快 1.5 倍以上。
• 博客自己的估算：H100 上 3.35 TB/s ÷ 2.48 GB ≈ 每秒 1,350 次前向；如果每层 7 个内核、每个内核乐观地只停顿 5 µs，就只剩约 770 次。

5.3 哪个设计最关键？（代码 02 演示 3–4）

在这个模型里，关键是”提前读权重”：只能提前读 1 个任务时，等依赖的时间又暴露出来，退化到和逐算子内核差不多。依赖粒度（细粒度 vs 粗粒度）只差几微秒——提前读得够深，等待就已经被藏住了。这是代码 README”计划 vs 实际”第 2 条记录的一个意外：原以为细粒度依赖是主要收益。

调度开销：一个 64 KB 任务的权重读取约 2.6 µs。代码 02 演示 4 把每个任务的调度开销从 0.1 µs 增加到 5 µs：2 µs 以内结果不变（750 µs），3 µs 时 951 µs，5 µs 时 1,553 µs——比逐算子内核还慢。任务越细，调度必须越便宜，这就是 HazyResearch 提前在 Python 里排好调度表、MPK 用 AOT 队列和预取任务描述的原因。

5.4 模型漏掉了什么（代码 02 演示 5）

模型说 99.7%，实测 78%。HazyResearch 博客给出了 B200 上一次前向（约 600 µs）的时间分解：

占大头的是”写出结果、检查计数器、读激活值”——每个任务都要做，模型里只算了一次延迟、而且假设能被藏住。博客还提到，CUDA 的异步屏障即使已经处于”通过”状态，每次也要约 60 ns。模型回答的是”去掉内核边界最多能省多少”；真实系统的工程难点恰恰在模型忽略的细节里。

6. MPK：从计算图自动生成 mega-kernel · Mirage Persistent Kernel

HazyResearch 的 mega-kernel 是为一个模型手写的。MPK（Mirage Persistent Kernel，Zhihao Jia 组与合作者）的目标是自动把一个模型（甚至多 GPU 推理）编译成 mega-kernel。

6.1 任务图：以 SM 为粒度

MPK 的核心表示是一张 SM 级的任务图：

• 任务（task）：一个 SM 上的一段计算或通信；
• 事件（event）：任务之间的同步点。任务和事件在图里交替出现：任务只连向它触发的事件，只从它依赖的事件接收边。
• 一个任务在它依赖的事件被激活后才能运行；完成时通知它触发的事件；一个事件收到所有关联任务的通知后被激活。

和 CUDA Graph 的区别：CUDA Graph 的节点是内核，边表示”整个内核 A 结束后整个内核 B 才能开始”；MPK 的节点是任务（一个 SM 上的一块工作），依赖细到”任务 B 只等它真正需要的那几个任务”。

6.2 编译器：从读写区域推导依赖

编译器把每个算子按输出切成任务（每个任务一个 SM），然后：对共享一个张量的两个算子，如果任务 A 的输出区域和任务 B 的输入区域有重叠，就建立依赖，并用事件表示。

代码 03 演示 1 在一个小网络上实现了这个分析（GEMV → 按块 tanh → GEMV → split-K GEMV）：28 个任务、52 条逐对的依赖边，按”依赖集合相同”合并后只有 9 个事件。例如 tanh 的每个任务只依赖覆盖它那 24 个元素的 3 个 GEMV 任务，split-K 的每一段只依赖写它那 16 个元素的 2 个任务——这些细粒度的依赖是内核边界表达不了的。

对 Llama-1B 的一层（megasim 的任务划分）：

一层 1,984 个任务，逐对的边约 38 万条，合并成 15 个事件后只要 3,968 条边。同一层用 CUDA Graph 表示是 5 个内核节点。事件是”逐对依赖”和”内核级依赖”之间的折中：既保留了细粒度（注意力按 KV 头、down 按 K 段），又不至于让图爆炸。

6.3 运行时：GPU 内部的任务并行运行时

MPK 论文描述的运行时：

• worker 与调度器：每个 worker 占一个 SM，有自己的任务队列，按先进先出执行；调度器以 warp 为粒度运行，每个 SM 上 4 个调度器 warp。论文的配置：H100 的 132 个 SM 里 128 个做 worker，4 个 SM 留给 16 个调度器（A100：104 + 16；B200：144 + 16）。
• 队列和计数器：事件队列和任务队列都是设备内存里的环形缓冲区，用 atomicAdd 入队和出队；每个事件记录激活前需要的触发次数。
• 两种任务启动方式：JIT 任务在它依赖的事件激活后才由调度器分配，适合工作量取决于数据的任务（例如注意力，长短取决于上下文长度），能做负载均衡；AOT 任务在执行前就按轮转预先放进 worker 的队列，省掉调度的往返，适合形状固定的任务。worker 优先执行 JIT 任务，AOT 任务在依赖满足时执行。
• 降低开销的其他技巧：每个任务描述（352 字节）预取到共享内存；分页的共享内存让下一个任务的数据在当前任务计算时就开始加载（与 HazyResearch 的思路相同）。论文的测量：Qwen3-8B 上 GPU 内部调度器只占运行时间的 0.28%。
• 计算与通信重叠：通信也是任务。每个 all-reduce 被拆成 GPU 之间的传输任务（用 NVSHMEM 的 signal-wait 实现）和本地的归约任务；后面的计算任务只要它需要的那部分通信完成就可以开始。

6.4 AOT 还是 JIT？（代码 03 演示 3）

128 个 worker：

两个教训：

• 任务一样大时，AOT 的轮转已经完美均衡，调度开销又低，AOT 更好。
• 任务长短随数据变化时，AOT 在编译时不知道这一步谁长谁短，比下限多 15%。但 JIT 如果只是”谁空闲谁取”并不一定更好；JIT 的真正优势是运行时知道每个任务多长，先分最长的，就能做到离下限 5%。这是代码 README”计划 vs 实际”第 4 条记录的结果。

6.5 正确性：依赖必须推导，不能手写（代码 03 演示 2）

代码 03 写了一个真的会算的 mega-kernel 解释器：6 个 worker、AOT 队列、用计数器判断依赖，每一步随机挑一个能前进的 worker（模拟 SM 之间不确定的执行速度）。

• 依赖正确时，20 种随机交错全部算对（与直接计算差 ≤ 7.8 × 10⁻¹⁶）；但逐位比较有 10 种不同的结果——split-K 的 4 个部分和用原子累加，顺序随交错方式变化。需要逐位可复现时，就要固定累加顺序。
• 只删掉一条依赖边（GEMV3[7] → GEMV4[3]），20 种交错里有 6 种算错：GEMV4[3] 读到了还没写好的数据。少一条依赖不会每次都出错——这正是数据竞争难以靠测试发现的原因。所以依赖必须由编译器从读写区域推导出来。

6.6 MPK 的结果（论文）

• Qwen3-8B 在 A100 上：每个 token 的解码时延从 vLLM、SGLang 的 14.5 ms 降到 12.5 ms；论文估计的硬件下限约 10 ms（16 GB 参数 ÷ 1.6 TB/s）。
• 单 GPU 批 1 服务：比 SGLang、vLLM 快 1.0–1.7 倍（摘要称时延最多降低 1.7 倍）；比 PyTorch 快 10 倍以上。
• 多 GPU：Qwen3-1.7B 张量并行扩展到 8 块 H100，比 SGLang、vLLM 快 1.1–1.4 倍。
• 消融：跨任务流水线让 Qwen3-8B 在 B200 上最后一个线性层快 1.2–1.3 倍；计算与通信重叠（Qwen3-1.7B，4 块 H100）带来 1.1 倍。
• GitHub README 称 MPK 能把 LLM 推理时延降低 1.2–6.7 倍（不同的对比基线和设置）。

7. 把推理时延推向硬件极限 · Towards the hardware limit

把本讲的数字放在一起看批 1 解码的”效率阶梯”（H100，Llama-3.2-1B）：

模型对两者都乐观，所以它们之间的差距（76% vs 99.7%）和实测的差距（≤ 50% vs 78%）方向一致、量级不同。剩下的 22% 去哪了？§5.4 的分解说：主要是每个任务写出结果、检查计数器、读激活值的开销，以及同步和初始化。继续逼近极限的方向包括：更粗的任务（减少每任务的固定开销，但会加重尾部）、把激活值留在片上（例如同一个 SM 连续执行有依赖的任务）、更便宜的同步原语。

另一个视角：批 1 解码的下限就是读一遍权重。要突破这个下限，必须改变”读一遍权重只产生一个 token”这件事——这是批处理（第 16 讲）、推测解码（第 17 讲）和量化（附加讲 A2）的出发点。

8. 开放问题 · Open problems

1. 动态性：服务系统里每一步的批大小、每个请求的上下文长度都在变（连续批处理）。mega-kernel 的任务图怎样适应？MPK 用”步数”和”token”两个元张量记录运行时状态，并用 JIT 任务处理依赖数据的工作量；更一般的动态形状仍然困难。
2. 资源分配：一个内核只有一套寄存器 / 共享内存配置，最苛刻的任务决定了所有任务的资源。怎样让不同任务各取所需？
3. 调度策略：哪些任务该 AOT、哪些该 JIT？任务切多细？这和第 9 讲的自动调优是同一类问题。
4. 可移植性与可调试性：mega-kernel 和具体 GPU 架构强绑定；一个内核里几万个任务，性能分析和调试工具需要重新设计（MPK 提供了一个按任务可视化时间线的分析器）。
5. 和服务系统的结合：mega-kernel 优化的是单步时延，服务系统还要考虑吞吐、多租户、KV 缓存管理。
6. 多 GPU 与多节点：通信任务化以后，怎样在更大的规模上调度？

9. 论文精读 · Paper reading

9.1 Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B（HazyResearch 博客，2025）

1. 问题：批 1 解码是访存受限的，理论上只受显存带宽限制；但主流引擎在 H100 上最多只用到一半带宽。原因是一次前向被拆成约一百个内核，每个内核的启动、收尾都让 GPU 空闲。

2. 想法：把整个 Llama-1B 前向融合成一个 mega-kernel：解释器执行预先排好的指令序列；共享内存分页，释放的页立刻交给下一条指令提前加载；全局内存计数器表示依赖，按块细分（MLP 中间结果 4 块）。

3. 验证：H100 上不到 1 ms（78% 带宽），比 vLLM 快近 2.5 倍、比 SGLang 快 1.5 倍以上；B200 上不到 680 µs。博客给出了详细的时间分解，坦率地指出剩余开销在哪里。

4. 从系统的角度看：这篇博客最有价值的是诊断——用一个简单的估算（“每层 7 个内核、每个停顿 5 µs”）说明问题的量级，再用一个手写的实现证明可以做到多好。它把”内核边界”本身识别为一种需要消除的开销，而不是天经地义的执行单位。

5. 局限：手写、只针对一个模型和批 1；调度表静态排好，难以处理动态的批处理和变长请求；工程复杂度高。这些正是 MPK 想用编译器解决的问题。

9.2 Compiling LLMs into a MegaKernel（MPK 博客）与 MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs（arXiv 2512.22219）

1. 问题：逐算子执行的三个局限——内核屏障让相关算子串行、无法跨算子流水；算子级依赖无法让计算和通信细粒度重叠；每次迭代成百上千次内核启动（CUDA Graph 降低了开销，但图基本是静态的）。手写 mega-kernel 又太难，不能推广到每个模型。

2. 想法：一个编译器 + 一个 GPU 内部的运行时。编译器把张量程序降低成 SM 级任务图（任务 + 事件），依赖由读写区域推导；为每个任务生成 CUDA 代码。运行时在一个常驻内核里执行任务：worker 执行任务、调度器分配 JIT 任务、AOT 任务预先排队、事件计数器激活后继任务。通信（all-reduce）也被拆成任务，和计算重叠。

3. 验证：单 GPU 批 1 服务比 SGLang、vLLM 快 1.0–1.7 倍（Qwen3-8B 在 A100 上 14.5 → 12.5 ms，下限约 10 ms）；多 GPU（8 块 H100）快 1.1–1.4 倍；GPU 内部调度器只占 0.28% 的运行时间；消融实验分别量化了跨任务流水线和计算通信重叠的收益。

4. 从系统的角度看：MPK 把”mega-kernel”从一次性的工程壮举变成了一个编译器问题：表示（任务图）、分析（依赖推导）、代码生成、运行时调度。这和本课程的主线一致——把手工优化变成自动化系统（第 8–11 讲的编译器、第 15 讲的自动并行都是同一个思路）。它也和第 11 讲的 Mirage 超优化器出自同一个项目：Mirage 优化单个内核，MPK 把整个模型变成一个内核。

5. 局限与开放问题：§8 列出的动态性、资源分配、调度策略、可调试性；此外（本讲义的推断，见练习 7），批大时计算更重、内核边界的占比下降，mega-kernel 相对成熟推理引擎的收益也会变小。

9.3 两篇材料之间的联系

HazyResearch 证明了”可以做到多好”，MPK 回答”怎样自动做到”。两者的核心机制几乎相同：常驻内核、共享内存分页让加载提前开始、计数器表示依赖。差别在于谁来决定任务和依赖：HazyResearch 由人设计 7 种指令和调度表；MPK 由编译器从计算图推导。本讲的模型（代码 02 演示 3）提示了一个共同的关键：让权重的加载不受依赖的约束——这是两者都用分页共享内存的原因。

10. 小结 · Summary

1. 内核边界隐含启动开销、全局同步、资源重新分配。深度学习框架”一个算子一个内核”，批 1 解码时一次前向有约一百个内核。
2. 批 1 解码的理想时间就是读一遍权重（Llama-3.2-1B 在 H100 上约 748 µs）。逐算子内核的额外时间 = 启动间隙 + 延迟 + 尾部效应（模型：CUDA Graph、不融合时 148 + 125 + 45 µs）。
3. CUDA Graph、融合、PDL 都只是部分补救；硬件越快，固定开销占比越大（B200 上 58%）。
4. mega-kernel：一个常驻内核，每个 SM 一个 worker，计数器表示依赖。关键是”权重和激活值无关、可以在等依赖时提前读”——它同时消除了间隙、延迟和尾部（模型：99.7%）。
5. HazyResearch 的 Llama-1B mega-kernel：H100 上不到 1 ms、78% 带宽，比 vLLM 快近 2.5 倍；剩下的开销主要是每个任务写出、检查计数器、读激活值。
6. MPK：SM 级任务图（任务 + 事件）、从读写区域推导依赖、worker + 调度器、AOT（形状固定的任务）+ JIT（依赖数据的任务）；Qwen3-8B 在 A100 上 14.5 → 12.5 ms（下限约 10 ms）。
7. 依赖必须推导：少一条依赖不会每次都出错（代码 03：6/20），这正是数据竞争难以测试的原因。

11. 常见误区 · Common misconceptions

• “mega-kernel 快是因为省掉了内核启动。” 启动间隙只是三笔代价之一。代码 01：CUDA Graph 已经把间隙降到 1.3 µs，延迟和尾部仍有 170 µs；代码 02 的消融说明，mega-kernel 最关键的是”提前读权重”。
• “CUDA Graph 就是 mega-kernel。” CUDA Graph 把一串内核的启动录下来重放，节点仍是内核，内核之间仍有全局同步；mega-kernel 只有一个内核，依赖细到任务。
• “任务切得越细越好。” 每个任务有固定开销（调度、写出、检查计数器、读激活值）。代码 02 演示 4：调度开销超过一个任务的读取时间后，mega-kernel 比逐算子内核还慢。
• “JIT 调度总比 AOT 好，因为它是动态的。” 任务一样大时 AOT 更好（开销低、已经均衡）；JIT 的价值在于利用运行时信息（代码 03）。
• “mega-kernel 对所有场景都有大收益。” 收益集中在批小、时延敏感的解码。批大时计算更重，内核边界的占比下降。
• “模型说 99.7%，所以 mega-kernel 能跑到 99.7%。” 实测 78%。模型帮助识别开销的来源，不能预测绝对性能。

12. 练习 · Exercises（含参考答案）

练习 1　Llama-3-8B（约 80 亿参数，BF16）批 1 解码，在 H100 上理想的每 token 时延是多少？如果每个 token 有 290 次内核启动、每次 CUDA Graph 间隙 0.8 µs，启动间隙占理想时间的多少？

答案：约 16 GB ÷ 3.35 TB/s ≈ 4.8 ms。290 × 0.8 µs = 0.23 ms，约 5%。（MPK 论文里 Qwen3-8B 在 A100 上的下限约 10 ms，因为 A100 的带宽只有约 1.6 TB/s。）

练习 2　QKV 有 192 个任务、132 个 SM。如果把每个任务改成 24 行（任务数 128），尾部效应会怎样？代价是什么？

答案：3072 / 24 = 128 个任务，一轮就完成，只有 4 个 SM 空闲，尾部效应几乎消失（理想 3.76 µs、每个 SM 读 96 KB → 约 3.8 µs）。代价：任务变大，每个任务的共享内存需求增加（96 KB = 6 页），mega-kernel 里能提前读的任务数从 3 降到 2；而且”让任务数恰好是 SM 数的倍数”在不同 GPU（132、148、108 个 SM）上要分别调整。

练习 3　用代码 01 的方法估算：如果某个推理引擎每层有 12 个内核（RoPE、残差加法都是单独的内核），每个内核固定开销约 3 µs（间隙 + 延迟），Llama-3.2-1B 的一次前向在 H100 上带宽利用率大约是多少？

答案：16 × 12 + 2 = 194 个内核 × 3 µs ≈ 580 µs 固定开销。总时间 ≈ 748 + 580 + 尾部（约 45） ≈ 1,370 µs，利用率约 55%——接近博客说的”最多 50%”。

练习 4　解释为什么 mega-kernel 可以在等上一个算子的结果时开始读下一个任务的权重，而逐算子内核不行。PDL 能做到多少？

答案：权重是常量，不依赖任何激活值；只有”读激活值、计算”必须等依赖。逐算子内核里，下一个内核要等上一个内核结束才开始，根本没有机会提前读。PDL 允许下一个内核提前启动并读权重，但在读输入前要和上一个内核整体同步，所以它能消除”读权重的延迟”，消除不了尾部效应（仍要等最慢的线程块）和粗粒度依赖的等待。

练习 5　代码 03 里删掉一条依赖边后 20 种交错中有 6 种算错。为什么不是全部算错？如果把 worker 数从 6 改成 1，会怎样？

答案：只有当 GEMV4[3] 恰好在 GEMV3[7] 写完之前读取 y3 时才出错；很多交错下 GEMV3[7] 已经先完成了。1 个 worker 时所有任务按队列（拓扑序）依次执行，GEMV3[7] 一定在 GEMV4[3] 之前完成，永远不会出错——数据竞争在串行测试里完全看不出来。

练习 6　mega-kernel 的寄存器和共享内存配置必须满足最苛刻的任务。假设注意力任务需要每线程 168 个寄存器、每块 200 KB 共享内存，而 GEMV 任务只需要 64 个寄存器、64 KB。这对 GEMV 任务有什么影响？HazyResearch 的分页设计怎样缓解？

答案：整个内核按 168 个寄存器、200 KB 配置，每个 SM 只能放一个 CTA，GEMV 任务也只能用这么多线程，占用率受限。分页让共享内存不是”每个任务固定一大块”，而是按页申请：GEMV 任务只用 4 页，其余页可以留给后面的任务提前加载——把”多余的”共享内存变成了预取深度。

练习 7　为什么批大小增大时 mega-kernel 相对逐算子内核的收益会下降？

答案：批大时每个权重被多个 token 复用，计算量随批大小增长，每个内核的运行时间变长；而启动间隙、延迟这些固定开销不变，占比下降。到预填充或大批量时，矩阵乘法本身是计算受限的，内核边界的开销几乎可以忽略。

练习 8　（开放）设计一个实验，用 megasim.py 回答：上下文长度从 1K 增加到 128K 时，mega-kernel 相对逐算子内核的收益怎样变化？提示：build_graph(context=...)。

参考思路：上下文变长时注意力读的 KV 缓存线性增长（128K 时每层约 256 MiB，超过权重本身），注意力任务从”小内核”变成”大内核”，固定开销占比下降，两者都更接近理想；但注意力任务的长短更不均匀（批里请求长度不同时），JIT 调度的价值上升。

13. 延伸阅读 · Further reading

• HazyResearch, Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B（2025）。
• X. Cheng, Z. Zhang, …, T. Chen, Z. Jia, MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs（arXiv 2512.22219）；项目主页 github.com/mirage-project/mirage。
• NVIDIA CUDA 文档：CUDA Graphs；Programmatic Dependent Launch（PDL）。
• K. Gupta, J. A. Stuart, J. D. Owens, A Study of Persistent Threads Style GPU Programming for GPGPU Workloads（InPar 2012）——“常驻线程”编程风格的早期研究。
• 第 6 讲的 persistent GEMM 与 warp 特化；第 11 讲的 Mirage 超优化器。

14. 下一讲预告 · Next lecture

第 3–7 讲都在手写（或手工设计）GPU 内核：分块、合并访存、共享内存、warp 特化、mega-kernel。每一步都需要对硬件的深入了解，代码也越来越复杂。第 8 讲《ML 编译器（一）：基于分块的 DSL》换一个角度：能不能让程序员只描述”每个块做什么”，把线程、共享内存、流水线交给编译器？我们会学习 Triton（按块编程、指针运算、为了 L2 复用的分组排序、编译栈）和 TVM（张量表达式与调度原语），并在 CPU 上实现一个极小的”类 Triton”分块 DSL 和一个 L2 缓存模拟器。