跳到主要内容
AI News HubLIVE
站内改写2 分钟阅读

Cohere 的 North Mini Code 解码大内核服务引擎

文章摘要

Cohere 发布了围绕解码大内核(decode megakernel)构建的 North Mini Code 服务引擎。在单张 H100 上使用 BF16,端到端吞吐比 vLLM 快 1.25–1.41 倍;批大小为 1 时达到 62% 的内存带宽理论峰值,且支持生产级特性与工具调用,代码已开源。

Cohere 的 North Mini Code 解码大内核服务引擎
报告错误

纠错通道尚未开通,可先复制下方文章信息留存。

查看更正说明
直接读正文

Cohere 今日发布了一个面向 North Mini Code 模型、以解码大内核(decode megakernel)为核心的服务引擎。官方给出的测试结果显示:在单张 H100 上采用 BF16 精度,其端到端吞吐比 vLLM 快 1.25~1.41 倍;并且这一套引擎的代码已开源在 GitHub 上,方便开发者直接查看或复用。

目前大多数 LLM 服务框架仍把每一次前向计算拆成一连串独立内核:先启动 QKV 投影,等待;再启动注意力,等待;接着启动 MoE,再等待。单独看每一个内核都足够高效,问题出在内核之间的等待上。在小批量解码时,GPU 有相当大比例的时钟周期并不是在计算,而是在等上一个内核结束、等下一个网格被调度。自回归解码,尤其是小批量场景,从根本上说是内存受限的:每一步都要从 HBM 读取大量参数,但真正完成的计算量却不多。

North Mini Code 是一个 30B 模型,每个 token 激活 3.3B 参数;在 BF16 下,每解码一步大约要从 HBM 搬运 6.6 GB 权重,外加 8K 上下文时约 0.5 GB 的 KV cache。H100 的 HBM 带宽是 3.35 TB/s,因此理论上的光速(Speed of Light)约为 470 token/s;而 vLLM 只能跑到 185 token/s,仅为光速的 39%。

大内核(megakernel)正是为缩小这个差距而设计的:不再一次启动几十个小内核,而是将整个 forward pass 放进同一个常驻内核中。Cohere 的做法是,在每个 SM 上只常驻一个线程块,由它从一个由 host 准备好的任务列表中读取待执行任务;任务间的依赖关系不再用内核边界表示,而是通过全局内存中的计数器来显式表达。于是调度的粒度从“一个算子”缩小到“一个算子的一个 tile”,同步的粒度也从整个 GPU 缩小到真正依赖的那些生产者。

速度提升主要来自四个方面。第一是启动和同步开销显著减少;第二是减轻 wave quantization:当 GEMM 的 tile 数量不是 SM 数的整数倍时,传统内核会产生闲置波次,而大内核可以把空闲 SM 立即用后续任务“回填”。North Mini Code 采用并行 transformer 层,注意力与 MoE 可并行计算,因此这种回填还能做得更激进。第三是消除虚假依赖:内核边界意味着全网格同步,所有 SM 都要等最慢的那个;而细粒度计数屏障允许某一路 O-proj 在对应注意力输出完成后立刻开始,不必等所有注意力组都完成。第四是对不依赖当前激活的权重做预取,例如 router 和 QKV 的权重可以在上一层的 O-proj 尚未结束时就开始搬运,利用本来会闲置的带宽。

效果方面,在 batch size 为 1 时,大内核达到 292 token/s,约为光速的 62%,是 vLLM 的 1.58 倍;这一优势在不同 batch size 以及最长 256K 上下文范围内都成立,且没有可测量的精度损失。更重要的是,Cohere 强调这是一个完整的服务系统:支持 continuous batching、paged attention、变长序列,并通过 OpenAI 兼容接口提供 tool calling,因此可以直接把 OpenCode 指向这个端点来进行编码。与此同时,实现本身并不复杂:整个内核就是单个 CUDA 文件,既不需要编译器,也没有引入新的编程范式或抽象,甚至附带了把已有内核改造成大内核的“配方”。

在技术渊源上,Cohere 明确借鉴了 Hazy Research 的“Look Ma, No Bubbles!”开创性工作,沿用了一些关键技术:SM 上的任务解释器模式、基于计数器的同步、以及任务边界处的重叠。不同之处在于,Cohere 的实现更依赖 tensor core 指令,即使 batch size 为 1 也使用 wgmma;另外,他们没有采用共享内存分页来做权重预取,因为这会引入复杂簿记和开销。作为替代,每个 opcode 有一套 warp-specialized 流水线,共享内存布局在编译期静态确定;同时利用同类型相邻 GEMM 任务之间的流水线阶段重叠,以及在 GEMM 内部的权重预取策略,让等待输入的任务仍然在搬运权重字节。这些做法让大内核在一个可部署服务中真正落地。

展开要点与分析

文章情报

工程师进阶

要点

  • 将整个解码前向过程融合为单个常驻大内核,减少内核启动与全网格同步开销。
  • 批大小 1 时吞吐 292 tok/s,是 vLLM 的 1.58 倍;优势在不同批大小和最长 256K 上下文下保持,且无明显精度损失。
  • 支持连续批处理、分页注意力、变长序列以及 OpenAI 兼容接口和工具调用。
  • 实现为单一 CUDA 文件,无需编译器或新编程抽象,并提供移植现有内核的指南。

要点与分析由自动化流程生成,可能有误,请结合原始来源核实。