Cohere 开源 Megakernel 推理系统

Cohere 开源 Megakernel，比 vLLM 快 1.58 倍 Cohere 为 North Mini Code 模型构建了一个围绕 “decode megakernel” 的推理服务系统...

Cohere 开源 Megakernel，比 vLLM 快 1.58 倍 Cohere 为 North Mini Code 模型构建了一个围绕 “decode megakernel” 的推理服务系统：整个 decode 阶段的前向计算由一个常驻 GPU 的持久 kernel 完成，不再逐算子启动 kernel。 cohere.com/blog/megakerne… 在单张 H100、BF16 条件下，batch size 1 时达到 292 tok/s；H100 显存带宽理论上限（speed-of-light）的 62%；比 vLLM（185 tok/s）快 1.58 倍；端到端服务（batch size 8）在真实基准上快 1.25–1.41 倍。 为什么值得做：解码的瓶颈在等待，不在计算 自回归解码在低 batch size 下是显存带宽受限的：每生成一个 token，都要把模型权重完整读一遍。North Mini Code 是 30B 参数模型，每个 token 激活 3.3B 参数，BF16 下每步解码要流过约 6.6 GB 权重，外加 8K 上下文时约 0.5 GB 的 KV cache。H100 的 HBM 带宽约 3.35 TB/s，理论极限约 470 tok/s。 而 vLLM 只跑到 185 tok/s，只利用了理论值的 39%。差距去哪了？传统引擎按算子切分——QKV 投影、attention、MoE、RMSNorm 各自是一个独立 kernel，kernel 之间隔着全 GPU 的隐式同步和启动开销。问题不在 kernel 本身，在 kernel 之间的等待：GPU 在算子边界处闲着，显存带宽在被白白浪费。对 memory-bound 的工作负载来说，SM 空转一拍，就是带宽利用率的直接损失。 Megakernel 是什么？ GPU 里有 100 多个独立的流处理器（H100 有 132 个 SM）。Megakernel 的做法是：每个 SM 启动一个 thread block，且这个 block 常驻整个解码步骤。它不从 driver 接收工作，而是去读一份由 host 预先写进全局内存的任务列表——每个任务是一个 tile 粒度的小工作项。数据依赖不再靠 kernel 边界（全 GPU 屏障）表达，而是靠全局内存中的计数器：任务完成时原子递增计数器，等待输入的任务自旋轮询这些计数器，依赖一满足就开工。 调度粒度从“一个完整算子”缩小到“某个算子的某个 tile”；同步范围从“整块 GPU”缩小到“我真正依赖的那几个生产者”。这是整个设计的核心抽象。 这条路线源于斯坦福 Hazy Research 的 "Look Ma, No Bubbles!"——他们曾把 Llama-3.2-1B 融进单个 kernel，batch 1 时达到 H100 带宽的 78%。Cohere 在此基础上做出了三点关键差异： 1. 即使 batch size 1 也用 tensor core 的 wgmma 指令做矩阵乘（比用 CUDA core 更快、寄存器压力更低）； 2. 不用共享内存分页做预取（Hazy 的方案 bug 多、代价高），而是给每种 opcode 一套静态布局的 warp 专用流水线； 3. 重叠方式改为两类：同一 GEMM 的相邻 tile 之间做多级流水；流水线内部让权重加载先于跨 SM 的激活等待发出。 快在哪里：三个超越“启动开销”的收益 区分“省 kernel 启动开销”和三个更深层的机制： 1. 填补波次量化留下的空洞。 132 个 SM 跑 200 个 tile 的 kernel 需要两个“波次”，第二波只有 68 个 SM 有活，其余 64 个空转。kernel 越小，这种取整损失越频繁。在 megakernel 里，任何输入就绪的 tile 可以在任何空闲 SM 上启动。North Mini Code 的并行 Transformer 层结构（attention 和 MoE 都从同一份归一化输入出发，最后由一个融合的残差相加 + RMSNorm 汇合）让 attention 分支和 MoE 分支天然独立，空闲 SM 可以被另一分支的 tile “回填”。 2. 消除伪依赖。 Kernel 边界是全网格屏障：即使 O-proj 只依赖某一个 KV group 的 attention 输出，传统上也要等整个 attention kernel 结束。细粒度计数器让依赖精确到“我等的那份数据好了没有”。 3. 权重预取。 权重是不可变的，不需要等激活值就绪就能提前从 HBM 拉取。系统在前一层 O-proj 的收尾阶段就开始激进预取下一层的 router 和 QKV 权重。 Cohere @cohere Introducing the next evolution in LLM text generation: the first fully-fledged serving system built around a decode megakernel. Delivering up to 1.58x faster performance than vLLM. Built for North Mini Code, completely open-source. 🔗 View Quoted Tweet 💬 2 🔄 0 ❤️ 1 👀 376 📊 2 ⚡