# 通过warp decode提升MoE模型推理效率

- 来源：Cursor Blog
- 作者：Less Wright, Federico Cassano & Zhiyuan Zhang
- 发布时间：2026-04-06 20:00
- AIHOT 分数：66
- AIHOT 标记：精选
- AIHOT 链接：https://aihot.virxact.com/items/cmosbbgu40075slojjfiqhvmm
- 原文链接：https://cursor.com/blog/warp-decode

## 精选理由

Cursor 把 MoE 推理的并行轴从专家翻转到输出神经元，Blackwell 上吞吐涨 1.84 倍还顺带提精度，这种同时赢性能和精度的内核优化极其罕见，做推理引擎的值得逐行读。

## AI 摘要

针对Blackwell GPU上的小批量解码，研究提出了一种名为“warp decode”的新方法。该方法颠覆了传统以专家为中心的计算路径，改为让每个GPU warp负责计算一个输出神经元。这一根本性改变消除了原有流程中五个纯数据管理的“簿记”步骤，将整个MoE计算层压缩为仅两个内核。其优势在于避免了填充、分散和中间缓冲区的读写，并通过warp独立性实现了更好的调度。在Blackwell GPU上，该方法实现了1.84倍的吞吐量提升，同时输出精度更高，与全FP32参考值的差距缩小了1.4倍，有效加速了模型研发流程。

## 正文

视频 · 前往原文观看

大多数 MoE 推理系统都围绕专家来组织 token 生成路径。这与路由的工作方式一致，也是大规模场景下的标准做法。然而，在 Blackwell GPU 上进行小批量解码时，我们发现围绕输出而非专家来组织内核效果更好。我们将这种方法称为“warp decode”。

我们通过思考 Blackwell 上 MoE 解码实际能达到的最大内存带宽，得出了 warp decode 方案。这促使我们完全翻转了并行化的轴。我们不再将 warp 分配给专家，而是将每个 warp 分配给单个输出值（神经元）。

能够同时提升性能和准确率的内核十分罕见，而 warp decode 正是其中之一。在 Blackwell 上，它带来了 1.84 倍的吞吐量提升，同时还将输出与完整 FP32 参考值的接近程度提升了 1.4 倍，从而提高了准确率。这加速了 Composer 的研究和训练流程，使我们能够更快地改进模型，并更频繁地发布新版本。

传统的 MoE 路径

现代 MoE 模型会将每个 token 路由到一组专门的专家网络子集，例如，在某一层从 128 个专家中选择 8 个。标准的实现方式会围绕这些专家组织所有计算：收集每个专家所需的 token，执行数学运算，然后重新组装结果。

这种方法在预填充和大批量场景下效果很好，因为每个专家承担的共享工作足以分摊组织数据的开销。但在自回归解码步骤中，我们一次只生成一个 token，没有足够的共享工作来证明这种做法的合理性。传统路径中的八个阶段里，有五个纯粹是为了管理以专家为中心的数据布局而存在，并不执行任何实际计算。

我们做出的改变

Warp decode 通过将并行化重新组织为围绕输出而非专家，消除了那五个“记账”步骤。

现代 GPU 以称为 warp 的 32 条并行处理通道为一组来执行指令。在我们的新方法中，每个 warp 被分配恰好一个输出值来计算。该 warp 直接从内存中流式读取所需的权重数据，将所有八个路由专家的结果累加到一个运行总计中，然后写入一个结果。

这种 warp 独立性使得 warp 解码无需任何暂存、交接、跨 warp 同步点或中间缓冲区即可运行。整个 MoE 计算层被压缩为两个内核：`moe_gate_up_3d_batched` 和 `moe_down_3d_batched`。

这两个内核的工作原理

在门控/上行内核中，每个协作线程数组（CTA）包含八个 warp，每个 warp 拥有一个中间神经元，对应每个 token 与每个被路由专家的配对。该 warp 加载被路由的专家 ID，读取该神经元的门控和上行权重行，并流式处理输入激活向量。MXFP8 权重在运行时被转换为 FP32，两个点积都在私有寄存器中累加。

由于这两个内核被融合为一次传递，激活向量只读取一次，并立即在两个投影中复用，无需任何共享内存暂存。在 warp 级归约之后，warp 应用 SiLU(门控) × 上行，并写入一个中间值。

在下行内核中，每个 warp 拥有一个 token 的一个输出维度。它遍历所有 top-k 被路由专家，加载相应的下行投影权重行，并流式处理中间激活值，同时将每个专家的路由权重折叠到一个单一的运行 FP32 累加器中。

处理完所有专家后，我们使用 `__shfl_xor_sync` 进行 warp 级蝶形归约，将 32 个通道局部的部分和进行归约。这直接编译为 PTX `shfl.sync.bfly` 指令，这是一个单一的硬件原语，可在 warp 内的通道之间交换寄存器，完全绕过共享内存。

这样做的好处是，我们不需要 L1 往返、存储体冲突或显式屏障，因为同步通过通道掩码内置于指令中。最终的加权 top-k 组合不再是单独的后处理步骤，而是成为投影本身的一部分。

warp 解码中的每个 warp 都是独立的，并在其整个生命周期中获得一个单一、稳定的任务：生成一个输出标量。正是这种 warp 独立性消除了传统路径所需的共享内存暂存、跨 warp 同步和中间缓冲区。

流水线简化与加速

Warp decode 通过两种主要机制实现性能提升：一是移除传统路径所需的阶段和缓冲区，二是创建 warp 独立性，从而实现更好的调度和延迟隐藏。

阶段消除

阶段消除带来了大部分吞吐量增益。我们消除了填充、分散和合并步骤。移除这些阶段需要从根本上重新组织并行方式，而不仅仅是融合传统流水线的各个阶段。

消除填充

传统路径：将每个专家的 token 列表填充到 2 的幂次方或 128 字节边界，以满足分组内核的要求。在解码阶段，对于单个 token 来说，这是无法分摊的开销。

Warp decode 路径：通过从不形成按专家划分的批次，完全避免了这一开销。

消除分散和合并

传统路径：每个专家完成后，会向 GPU 内存写入八个中间结果，然后运行一个独立的规约步骤来合并它们。

Warp decode 路径：每个专家的路由权重被折叠到 warp 内部的运行累加器中。这八个中间结果永远不会在内存中具体化，从而节省了后续规约过程的写入和读取成本。

缓冲区消除

这种重组还移除了传统路径因其以专家为中心的布局而需要的两个中间内存缓冲区。

第一个是激活值收集缓冲区，即输入激活向量被复制并重新排列成专家主序布局。在批次大小为 1 时，这相当于对已存在的数据进行完整复制。第二个是按专家划分的输出缓冲区。对于八个专家和隐藏维度 2048，每个 token 在 BF16 格式下需要 8 × 2048 × 2 字节 = 32 KB 的缓冲区，经过分配、写入、立即读取一次，然后丢弃。

Warp decode 通过将八个专家的贡献折叠到跨 32 个 warp 通道的寄存器累加器中，消除了这两个缓冲区，在此过程中，直到最终单标量写入之前，没有任何数据到达全局内存。每个 token 移除 32+ KB 的中间缓冲区流量，为真正决定性能的权重行释放了 L2 缓存容量。

Warp 独立性

此次重组还使保留的计算变得更快，因为该内核在设计上是“极度并行”的：每个线程束完全独立于其他线程束。由于每个线程束恰好拥有一个输出标量，并且只读取它需要的权重行，因此线程束之间不存在共享的可变状态。

在单个线程束的层面上，这种独立性是彻底的。输入激活值是只读的，累加器位于私有寄存器中，输出写入则指向一个唯一的地址。从硬件调度器的角度来看，整个输出维度就是一个由独立工作项组成的扁平化池。

GPU的线程束调度器可以在任何时间、以任何顺序发出任何线程束，且不受正确性约束。当一个线程束因等待内存加载而停滞时，调度器会立即切换到另一个线程束。由于在B200的148个流式多处理器上同时有数千个线程束在运行，内存延迟几乎完全被来自其他线程束的有效计算所掩盖。

该内核还能线性扩展，因此将输出维度加倍会使独立线程束的数量加倍，且无需增加任何同步操作。这一点在token批次维度上同样成立，因此调度器看到的是一个所有工作节点之间无边连接的扁平化命名空间。这与传统路径形成鲜明对比，在传统路径中，专家级别的GEMM内核需要块内协调。

结果

大规模端到端解码吞吐量

在我们使用NVIDIA B200 GPU运行Qwen-3风格模型的内部推理系统上进行测试，产生了持续的吞吐量提升。该吞吐量提升在所有上下文长度区间内保持一致，证实了这是一个纯粹的生成阶段改进，不依赖于提示词长度。

精度提升

移除中间激活量化步骤会带来可衡量的质量影响。将BF16激活值转换为MXFP8再转换回来，会引入一个舍入误差基底，该误差会在模型各层中累积。Warp解码在整个过程中将激活值保持在BF16格式，累加器保持在FP32格式，因此归约操作永远不会基于降级后的输入进行。结果是，与经典路径相比，Warp解码生成的输出与完整的32位真实值的接近程度提升了1.4倍。

硬件效率

我们开发 warp decode 的起点是探究：我们能将性能逼近硬件最大吞吐量的多少？B200 在连续内存读取场景下的实测峰值为 6.8 TB/s（通过拷贝内核测得）。Warp decode 在 B=32 时可持续达到 3.95 TB/s，即峰值的 58%。剩余差距很可能反映了专家路由产生的随机访问模式所带来的内存延迟开销，因为每个 token 可能路由到不相邻的专家，例如 5、8、14、19 等。

相比之下，峰值吞吐量是通过连续（0,1,2,3）内存读取测得的。在所有批次大小下，与参考实现的一致性都非常严格：最小余弦相似度 > 0.999996，最大绝对差值为 0.001953。

Warp decode 与 Composer 训练

Warp decode 并非专家中心化执行的通用替代方案。像预填充和大批量推理这类高负载任务，仍然受益于专家中心化的打包方式，因为许多 token 共享同一个专家，且组织这些 token 的开销可以通过足够的实际计算来分摊，从而变得值得。

当每个专家共享的工作量不足以证明上述开销的合理性时（这在 MoE 解码场景中经常出现），warp decode 便成为更优选择。这使得它成为我们持续改进 Composer 的重要组成部分。虽然预训练数据和强化学习的投入决定了模型输出的质量，但像 warp decode 这样的推理投入则决定了这些输出能以多快的速度和多高的准确性到达开发者手中。

推出 Grok 4.5

通过自动安装引导 Composer

Shomil、Joshua 与 Andrew

保持 Cursor 应用的稳定性

Andrew 与 Kevin