深入North Mini Code的Megakernel服务引擎（22分钟阅读）

# Cohere 的 North Mini Code 超级内核服务引擎 | Cohere 来源：https://cohere.com/blog/megakernels 今日，Cohere 发布了一款围绕解码超级内核构建的 North Mini Code (https://cohere.com/north-mini-code) 服务引擎：在单张 H100 上使用 BF16 精度，端到端速度比 vLLM 快 1.25 至 1.41 倍。您可以在 GitHub (https://github.com/cohere-ai/cohere-megakernel) 上探索该服务引擎背后的代码。 大多数大型语言模型（LLM）服务栈仍视每次前向传播为一系列内核的序列：启动 QKV，等待；启动注意力，等待；启动混合专家（MoE），等待。单个内核的启动本身没问题，问题在于中间的等待。在小批量大小下，GPU 在每个解码步骤中花费了惊人比例的时间在等待而非计算。 自回归解码，尤其是在较低批量大小下，本质上是内存受限的。在每个解码步骤中，我们从高带宽内存（HBM）中移动了大比例的内存，而相对计算较少。这意味着正确的问题是：我们能多有效地利用内存带宽，而非浮点运算能力。以 North Mini Code 为例，这是一个 30B 参数的模型，每个 token 激活 3.3B 参数，使用 BF16 精度意味着在每个解码步骤中需要流式传输 6.6 GB 的权重，加上约 8K 上下文长度下 0.5 GB 的键值（KV）缓存。H100 通过 HBM 提供 3.35 TB/s 的带宽，理论极限速度（SoL）约为 470 token/s。vLLM 服务该模型的速度为 185 token/s，仅为理论极限的 39%。 超级内核近来备受关注，被认为是缩小这一差距的方法：与其运行上百个微小内核，不如将整个前向传播作为一个持久化内核运行。从 Hazy Research 的开创性工作 “Look Ma, No Bubbles!” (https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles)（我们将在下文回顾其设计）开始，众多后续工作相继发布，实现了不同程度的加速。现有工作主要沿两个方向发展：自动生成超级内核的编译器，以及在批量大小为 1 时测量解码速度的独立演示。 我们更进一步。本文介绍了我们所认为的首个完全成熟的、围绕解码超级内核构建的服务系统。它支持真实服务器所需的一切功能：连续批处理、分页注意力以及不规则序列长度，所有这些都通过一个兼容 OpenAI 并支持工具调用的端点提供。将 OpenCode 指向它，您就可以用它进行编码。 在批量大小为 1 时，我们的超级内核达到 292 token/s，或理论极限的 62% —— 比 vLLM 快 1.58 倍。这一优势在不同批量大小和长达 256K 的上下文长度下均得以保持，且精度没有可测量的损失。 图 1：不同上下文长度下批量大小为 1 的解码吞吐量。超级内核始终优于 vLLM。 我们还发现超级内核比其声誉所暗示的要容易编写得多，因此我们提供了一份将已有内核移植为超级内核的方案。我们的内核是一个单独的 CUDA 文件：无需编译器，无需新的编程范式，无需奇异的抽象 —— 只是普通的平铺 GEMM（通用矩阵乘法）和普通的分页注意力，被重构以适应单一的调用约定。 ## **什么是超级内核？** GPU 大约有 100–150 个独立的处理器，称为 SM（流式多处理器），它们都运行相同的程序 —— 一个*内核* —— 处理不同的数据。超级内核是一个单一的持久化内核，它运行整个前向传播：我们为每个 SM 启动一个线程块，并且它在整个解码步骤中保持常驻。每个块不从驱动接收工作，而是读取一个*任务列表* —— 一份由主机准备并存放于全局内存中的、它应该执行的小工作片段列表。依赖关系不是通过内核边界编码的，而是通过全局内存中的显式计数器来表达，任务在完成时递增计数器，在需要输入时自旋等待。 其结果是调度的单位从整个操作缩小到一个操作的一个图块，同步的单位从整个 GPU 缩小到任务所依赖的特定生产者。 图 2：一个解码步骤，从主机到设备。主机不是为每个操作启动一个内核，而是将步骤分解为任务 —— 每个任务是一个操作的一个图块 —— 并以轮询方式分配给各个 SM，这样每个 SM 在全局内存中都有自己的任务列表。大多数任务的顺序由主机上的调度器决定（见调度器部分）。全注意力和混合专家（MoE）是例外：它们的任务数量取决于活动的序列长度和路由，因此它们进入共享工作队列，任何 SM 都可以从中获取。我们将在后续章节中描述详细信息。 ## **速度提升来源** 推理引擎通常为每个操作（RMSNorm、QKV、注意力、MoE 等）启动一个内核，大部分优化工作都放在使每个内核尽可能快上。这种方法在训练和预填充阶段效果很好，因为工作负载是计算密集型的 GEMM，并且每个内核都有足够的工作来占满 SM。 解码则相反：它受限于内存带宽和延迟。解码步骤主要是低算术强度的 GEMV（通用矩阵向量乘法），因此其速度取决于权重从 HBM 流式传输到共享内存的速度。任何阻碍权重移动的东西都是时间损失，而每个操作一个内核的方法有多个会导致权重停止移动的地方。这些停顿加起来占用了典型推理引擎未使用的 61% 带宽中的大部分。 最简单的好处来自**减少启动和同步开销**。在两个连续的内核之间，所有 SM 必须完成，然后才能启动下一个内核，并且驱动程序必须调度下一个网格。对于一个由数十个小内核组成的解码步骤，这些间隙会累积起来。超级内核每步只支付一次成本，而不是每个操作支付一次。我们按影响程度大致列出对本模型更重要的另外三个好处： ### **1. 减少波次量化** 假设一个内核有 200 个图块的工作要做，而 GPU 有 132 个 SM。前 132 个图块并行运行；剩余的 68 个在第二个波次运行，同时有 64 个 SM 空闲。该内核花费了两个波次的时间来完成 1.5 个波次的工作量，内核越小，这种舍入造成的效率损失就越严重。这不是我们通过更均匀地分配工作就能解决的问题。GEMM 图块的形状受到矩阵维度和内核设计的限制，因此总图块数很少能正好是 SM 数量的整数倍。 图 3：波次量化降低了 GPU 利用率。 在超级内核中，没有需要向上舍入的边界：一个输入准备就绪的图块会在任何一个空闲的 SM 上启动。North Mini Code 从这种方式中获益比大多数架构更多，因为它使用了*并行 Transformer 层*：注意力和 MoE 前馈网络基于相同的归一化输入计算，并且仅在层末尾通过一个融合的残差相加 + RMSNorm 重新汇合，因此注意力和 MoE 都不需要彼此的输出。 图 4：North Mini Code 使用的并行 Transformer 层。 使用超级内核，我们可以用准备就绪的工作来“回填”空闲的 SM。并行 Transformer 层允许我们以更积极的方式进行回填：只要有可能，我们就会*确定性地*将可能准备运行的任务放置在空闲的 SM 上。任务放置的细节在任务调度器部分描述。在下图中，我们比较了由传统服务栈和超级内核运行的 MoE 解码层。 图 5：一个 MoE 解码层，相同的操作按相同顺序调度，两种方式。上：每个操作一个内核。由于波次量化和硬件抖动，注意力未能同时完成。内核边界处的屏障使每个 SM 都等待最慢的一个（阴影部分）。相同的模式在每个内核边界处重复。下：超级内核。随着每个 SM 完成其注意力工作，它开始一个 MoE 图块，因此相同的间隔承载着工作而非等待。使用 16 个 SM 和每个操作几十个图块绘制，以使任务可见；实际内核有 132 个 SM 和数千个任务。时长为说明性，时间轴为相对值。 ### **2. 消除虚假依赖** 即使工作负载相同，SM 也并非总是在同一时间完成工作。内核边界是一个全网格屏障，因此最慢的 SM 决定了所有 SM 的节奏。例如，如果注意力被分割到 4 个键/值组，其中一个组提前完成，该 SM 将闲置直到其他三个组赶上，尽管其下一个操作实际需要的数据已经位于内存中。细粒度的屏障消除了这种虚假依赖：对于给定的 KV 组，O-proj（输出投影）在*该组的*注意力输出就绪时即可开始。同样，MoE 下投影任务可以在对应专家的上投影任务完成后立即开始，无需等待所有专家的上投影完成。 ### **3. 权重预取** 权重是不可变的 —— 它们根本不依赖于当前步骤的激活。因此，一个任务可以在其激活依赖项被满足*之前*就开始将其权重图块从 HBM 流式传输到共享内存中，而内核边界是禁止这种做法的。我们在路由器和 QKV 投影中对此进行了最积极的应用，它们在前一层 O-proj 的尾部，甚至在 RMSNorm 运行之前，就开始预取权重，以利用否则会被浪费的带宽。 ## **我们的起点** 我们的设计借鉴了前面提到的 Hazy Research 的开创性文章中的大量洞察，该文章将 Llama-3.2-1B 的前向传播融合到一个内核中，在批量大小为 1 时达到了 H100 内存带宽的 78%，而 vLLM 和 SGLang 大约只有一半。他们的三个想法对我们至关重要： - **GPU 上的“任务解释器”模式。** 每个 SM 遍历在主机上准备并跨前向传播重用的任务描述符列表。一个控制线程束读取描述符并将任务分派到各种设备函数，每个函数实现一种操作。我们在图 2 和图 6 中描述了我们对该想法的实现。 - **基于计数器的同步。** 依赖屏障是全局内存中的普通整数，在每个步骤开始前置零。一个任务完成时递增一个计数器，开始前自旋等待一个计数器。我们在屏障部分和图 7 中详细描述了我们的实现。 - **跨任务边界的重叠。** 一个任务可以在同一 SM 上的前一个任务仍在存储其结果时，开始加载其权重。 ### **我们与众不同的实现** 我们的实现与他们的超级内核在几个方面有所不同。我们的 GEMM 实现严重依赖张量核心指令，即使在批量大小为 1 时也是如此，因为我们发现 wgmma 在我们的场景中比 CUDA 核心稍快，并且减少了寄存器压力。 另一个区别是我们不使用共享内存分页来实现权重预取。我们最初尝试了共享内存分页，它允许在前一个任务释放缓冲区之前开始内存加载。然而在实践中，簿记很复杂，是稳定的 bug 来源，并且高开销超过了收益。相反，每个操作码都有自己的线程束专用流水线，其共享内存布局在编译时静态确定，我们从两个更廉价的地方获得重叠。 **在相同类型的连续 GEMM 任务之间。** 一个 MoE GEMM 任务遍历图块列表，流水线在其整个列表中保持其阶段，而不是在每个图块边界处清空和重新填充流水线。在 MoE GEMM 内部，下一个图块的权重已经在传输中，而当前图块仍在张量核心或尾声处理上。因为它是具有相同共享内存布局的相同操作，共享内存分页简化为一个简单的多阶段流水线。这种预取与持久化分组 GEMM 内核具有相同的精神。 **在 GEMM 流水线内部。** 权重是不可变的，因此生产者线程束在跨 SM 等待激活*之前*就发出其权重图块加载；在本文后面展示的伪代码（列表 1）中，`prefetch_weight_tiles` 位于 `wait_input_bars` 之上。MoE 下投影就是一个这样的例子：在“上/门”仍在计算下投影将要消耗的隐藏状态时，其专家权重就开始从 HBM 流式传输。因此，一个阻塞在输入上的任务仍在移动数据。另一个例子是在 RMSNorm 完成之前进行 QKV 和路由器预取。 我们还在调度方面投入了大量精力。大多数任务遵循主机构建的静态调度，我们可以精确调整；注意力和 MoE 使用本地工作窃取来平衡来自连续批处理和路由的运行时依赖性工作。我们将在下文讨论这一点。 ### **每个操作一个 ABI** 整个超级内核可以看作是各种较小内核通过一个共同的调用约定拼接在一起 —— 每个较小的内核必须用恰好 3 个线程束组（每个线程束组有 4 个线程束）实现，包含 8 个消费者线程束、1 个控制器线程束、1 个生产者和 1 个存储者线程束，其中每个线程束都有自己的寄存器要求。此外，每个较小的内核必须从固定大小的任务描述符中读取其“参数”。我们觉得这种约定类似于编译器和操作系统中的应用二进制接口（ABI）。在整篇文章中，我们将使用**“ABI”**一词来描述该约定。 让我们以一个 GEMM 为例来理解 ABI 是如何工作的。 图 6：一个 GEMM 图块如何成为 SM 上的工作。主机将输出分割成图块，并将每个图块写成一个 32-int32 描述符。在设备上，一个控制器线程束读取下一个描述符并根据其操作码（QKV、注意力、O-proj 或其他十四种中的任何一种）进行切换，因此同一个 SM 可以连续运行不同的操作。2×4 网格是说明性的；真正的批量为 1 的 QKV 有 80 个图块。 使其可手工编写的关键在于，每个操作，而不仅仅是 GEMM，都遵循该 ABI：任务是什么，它以什么形状运行，以及如何发出完成信号。一旦所有操作都遵循相同的 ABI，将它们组装成超级内核就变得非常容易。本节的其余部分将详细说明这三件事。 ### **任务** 超级内核不发明新操作。它采用常规的解码计算图 —— QKV、注意力、O-proj、路由器、MoE、残差/RMSNorm、语言模型头 —— 并将其降级为一组小型的平铺*任务*。十六种操作码覆盖整个解码步骤： 密集 GEMM - QKV_PROJ, O_PROJ, FFN_UPGATE_ACT, FFN_DOWN, LM_HEAD 一个矩阵乘法的一个输出图块，可选一个 split-K 归约的切片。QKV_PROJ 还在其尾声中在寄存器内应用 RoPE；FFN_UPGATE_ACT 在其尾声中融合了 SiLU 及乘法操作。 注意力 - ATTN_DECODE, ATTN_COMBINE, ATTN_DRAIN ATTN_DECODE 在部分 KV 切片上计算注意力；ATTN_COMBINE 将这些部分合并为真正的输出。ATTN_DRAIN 仅用于全注意力层。它是对动态生成的工作队列的声明器（参见后面的调度部分）。 MoE 路由 - ROUTER_GEMM, ROUTER_TOPK, ROUTE_FINALIZE, MOE_GATHER 为 128 个专家评分，计算每个 token 的 top-k 专家，并生成 MoE GEMM 工作队列。 MoE GEMM - MOE_UPGATE_ACT_DRAIN, MOE_DOWN_DRAIN, MOE_COMBINE 专家 FFN，作为对 ROUTE_FINALIZE 生成的工作队列的声明器。MOE_COMBINE 将每个 token 的 8 个专家输出按其路由器权重求和。 **所有 GEMM 操作码共享一个流水线。** O-proj、路由器、密集 FFN、语言模型头以及两个 MoE 专家 GEMM 运行相同的线程束专用体：生产者将权重和激活图块流式传输到共享内存阶段环，消费者在张量核心上累积，存储者写入输出图块并到达屏障。每个操作码不同的是一个简短的每操作细节列表 —— 哪些张量、等待哪个屏障以及发出哪个屏障的信号、是否