Cursor为Blackwell从零构建MXFP8内核，MoE层提速3.5倍，端到端训练提速1.5倍

在构建更强大的 AI 模型的这场竞赛中，传统路径很简单：升级到最新最强大的硬件。但 Cursor 发现释放下一代 GPU 的真正潜力远非即插即用那么简单。

在从 NVIDIA 的 Hopper H100s 升级到新旗舰 Blackwell B200s 后，该团队遇到了一个「升级陷阱」：硬件性能翻倍，但实际训练速度却被 MoE 层的效率拖慢，新架构的设计反而放大了数据搬运和量化的开销。

这就像给一辆赛车换上了动力翻倍的新引擎，却发现原有的轮胎完全无法承载这股力量，导致速度反而下降。

他们的解决方案是回归基础，自己定制「赛车胎」：在 GPU 内核级别从零开始重写整个混合专家（MoE）训练层。

Cursor 不仅解决了瓶颈问题，还彻底释放了 Blackwell 架构的潜能。通过抛弃对现有 CUDA 库的依赖，他们能够：

• 直接针对 TMEM 的新特性设计数据流管线，避免无谓的寄存器搬运开销；

• 将量化与反量化逻辑融入内核计算流程，大幅压缩了内存带宽占用；

• 优化 MXFP8 的 microscaling 实现，在保证训练收敛质量的同时，把性能推到极限。

最终效果是：MoE 层在前向和反向传播中都实现了 3.5 倍提速，端到端训练速度在 Blackwell 上快了 1.5 倍，相比最初的 Hopper GPU 方案实现了 2 倍的加速。

与 BF16 相比，MXFP8 MoE 的相对加速（归一化为 1.0）。

Cursor 团队在博客中详细介绍了相关技术细节，并分享了他们的工程经验和性能数据。

• 博客地址：https://cursor.com/en/blog/kernels

为什么现有 MoE 内核在 Blackwell 上失效？

为了降低计算成本，模型训练普遍采用低精度数据格式（如 FP8）。但简单地将高精度数字（如 0.0001）转换为 FP8 会导致其被四舍五入为零，丢失信息。

微缩放（MX）通过将张量（Tensor）分割成许多小数据块（例如每 32 个元素一块），并为每个块计算一个独立的缩放因子（scale factor）来解决这个问题。

MXFP8 量化示例：每个 1x32 块共享一个缩放因子。

这样，每个块内的数据都能被有效缩放到 FP8 的可表示范围内，从而在保留精度的同时享受低精度计算带来的性能优势。Cursor 使用的 MXFP8 就是这样一种格式。

张量内存（TMEM）瓶颈

在 Hopper (H100) 架构上，张量核心的计算结果直接累积在寄存器中，后续的「反量化」等操作可以流畅地进行。

然而，Blackwell (B200) 引入了新的张量内存（TMEM）来存储累加结果。这意味着任何自定义的算术操作都必须经历一次低效的数据往返：TMEM → 寄存器 → CUDA 核心处理 → TMEM。

这种异步数据传输会在张量核心的计算管线中产生「气泡」，大幅降低执行效率。更关键的是，尽管 Blackwell 的 FP8 张量核心吞吐量翻倍，其 CUDA 核心性能仅提升了约 33%，导致反量化速度严重滞后于计算速度。

该甘特图截取自我们定制的 Blackwell 注意力核。第一行显示了张量核心（QKT）的活动情况；第二行显示了 CUDA 核心的活动情况（数据从 TMEM 加载至寄存器，然后执行 softmax）。从 TMEM 到寄存器的加载延迟，导致了张量核心出现流水线气泡。

数据显示，在特定配置下，Blackwell 上的反量化耗时是矩阵乘法本身的 1.76 倍，远高于 Hopper 上的 1.03 倍。

Hopper 与 Blackwell 上的相对反量化成本。

被忽视的「量化税」

除了 TMEM 瓶颈，数据「量化」过程本身也成了性能杀手。

以一个典型的 MoE 矩阵乘法为例，计算本身可能仅需 1.16 毫秒，但将输入矩阵量化为 MXFP8 格式并写回内存就需要搬运近 2.9 GB 的数据，耗时约 0.44 毫秒，占到计算时间的近 40%。

在反向传播中，这个开销因需要转置-量化而翻倍，达到 0.88 毫秒，占比高达 76%。这意味着，如果优化不当，MXFP8 带来的性能提升可能被完全抵消。

此外，现有的开源量化内核不仅带宽利用率低，其生成的缩放因子（scale factor）布局还与 Blackwell 的硬件指令不兼容，需要额外的、拖慢性能的重塑操作。

Cursor 如何从零重写MoE 层？

面对这些挑战，并发现现有的开源库（如 NVIDIA 的 TransformerEngine）并非最佳选择，Cursor 团队选择放弃高层依赖，使用纯 CUDA 和 PTX 汇编语言亲自编写 MoE 层的 GPU 代码。

优化策略

• 拥抱原生硬件指令

他们没有与 TMEM 架构对抗，而是围绕原生的 tcgen05.mma 指令构建内核。这使得 GPU 硬件自身能够处理 MXFP8 所需的缩放，完全消除了 TMEM 和 CUDA 核心之间低效的数据移动。

• 设计高效的数据流水线

他们实现了一个复杂的流水线，采用了诸如「Warp 专精」（将特定任务分配给不同的线程组）和 2-CTA（协同线程阵列）模式等技术。

Warp 专精将特定的任务分配给不同的线程组（Warp）。例如，Warp 0 负责从主内存加载数据到共享内存，Warp 1 负责加载缩放因子，Warp 2 负责将缩放因子从共享内存移至 TMEM，而 Warp 3 则专门负责启动矩阵乘法计算。这使得各个环节可以高度并行。

2-CTA 模式允许两个 GPU 流式多处理器（SM）协同完成单个矩阵乘法，通过共享 B 矩阵来减少内存流量，带来了 15-20% 的性能提升。

• 针对 MoE 工作负载进行优化

对于 MoE 训练中特有的分组矩阵乘法，他们应用了一种名为「专家级超分组」的 L2 缓存优化启发式算法。这确保了内存访问模式保持高效，将标准矩阵乘法与分组矩阵乘法之间的性能下降限制在仅 4%。

「秘密武器」：量化内核与低精度配方

该团队开发了一个自定义的 MXFP8 量化内核，他们称这是目前用于 MoE 训练的最快内核。微基准测试显示，其内核持续的内存带宽超过 6.2 TB/s，相比他们从现有开源工具测得的约 4.5 TB/s 有了显著提升。

至关重要的是，他们的内核输出的数据内存布局与 tcgen05.mma 指令所要求的完全一致，避免了其他工具所必需的、耗时的额外「重塑」步骤。

基于内存带宽利用率的 MXFP8 量化内核比较（E4M3，32 块大小的缩放）。

团队还确定了一种特定的低精度「配方」，能够在不影响训练质量的情况下提供最高速度。通过使用元素类型为 FP8E4M3、块大小为 32 的 MXFP8 格式，他们能够使训练损失的收敛情况与速度慢得多的 BF16 格式几乎完全匹配。

团队公布的训练损失曲线显示，两种方法几乎没有区别，证明了性能的提升并未以牺牲准确性为代价。

BF16 与 MXFP8 训练损失超过 10k 步：几乎无法区分。

文章来自于微信公众号“机器之心”。