在構(gòu)建更強大的 AI 模型的這場競賽中，傳統(tǒng)路徑很簡單：升級到最新最強大的硬件。但 Cursor 發(fā)現(xiàn)釋放下一代 GPU 的真正潛力遠非即插即用那么簡單。

在從 NVIDIA 的 Hopper H100s 升級到新旗艦 Blackwell B200s 后，該團隊遇到了一個「升級陷阱」：硬件性能翻倍，但實際訓(xùn)練速度卻被 MoE 層的效率拖慢，新架構(gòu)的設(shè)計反而放大了數(shù)據(jù)搬運和量化的開銷。

這就像給一輛賽車換上了動力翻倍的新引擎，卻發(fā)現(xiàn)原有的輪胎完全無法承載這股力量，導(dǎo)致速度反而下降。

他們的解決方案是回歸基礎(chǔ)，自己定制「賽車胎」：在 GPU 內(nèi)核級別從零開始重寫整個混合專家（MoE）訓(xùn)練層。

Cursor 不僅解決了瓶頸問題，還徹底釋放了 Blackwell 架構(gòu)的潛能。通過拋棄對現(xiàn)有 CUDA 庫的依賴，他們能夠：

• 直接針對 TMEM 的新特性設(shè)計數(shù)據(jù)流管線，避免無謂的寄存器搬運開銷；
• 將量化與反量化邏輯融入內(nèi)核計算流程，大幅壓縮了內(nèi)存帶寬占用；
• 優(yōu)化 MXFP8 的 microscaling 實現(xiàn)，在保證訓(xùn)練收斂質(zhì)量的同時，把性能推到極限。

最終效果是：MoE 層在前向和反向傳播中都實現(xiàn)了 3.5 倍提速，端到端訓(xùn)練速度在 Blackwell 上快了 1.5 倍，相比最初的 Hopper GPU 方案實現(xiàn)了 2 倍的加速。

與 BF16 相比，MXFP8 MoE 的相對加速（歸一化為 1.0）。

Cursor 團隊在博客中詳細介紹了相關(guān)技術(shù)細節(jié)，并分享了他們的工程經(jīng)驗和性能數(shù)據(jù)。

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

為什么現(xiàn)有 MoE 內(nèi)核在 Blackwell 上失效？

為了降低計算成本，模型訓(xùn)練普遍采用低精度數(shù)據(jù)格式（如 FP8）。但簡單地將高精度數(shù)字（如 0.0001）轉(zhuǎn)換為 FP8 會導(dǎo)致其被四舍五入為零，丟失信息。

微縮放（MX）通過將張量（Tensor）分割成許多小數(shù)據(jù)塊（例如每 32 個元素一塊），并為每個塊計算一個獨立的縮放因子（scale factor）來解決這個問題。

MXFP8 量化示例：每個 1x32 塊共享一個縮放因子。

這樣，每個塊內(nèi)的數(shù)據(jù)都能被有效縮放到 FP8 的可表示范圍內(nèi)，從而在保留精度的同時享受低精度計算帶來的性能優(yōu)勢。Cursor 使用的 MXFP8 就是這樣一種格式。

張量內(nèi)存（TMEM）瓶頸

在 Hopper (H100) 架構(gòu)上，張量核心的計算結(jié)果直接累積在寄存器中，后續(xù)的「反量化」等操作可以流暢地進行。

然而，Blackwell (B200) 引入了新的張量內(nèi)存（TMEM）來存儲累加結(jié)果。這意味著任何自定義的算術(shù)操作都必須經(jīng)歷一次低效的數(shù)據(jù)往返：TMEM → 寄存器 → CUDA 核心處理 → TMEM。

這種異步數(shù)據(jù)傳輸會在張量核心的計算管線中產(chǎn)生「氣泡」，大幅降低執(zhí)行效率。更關(guān)鍵的是，盡管 Blackwell 的 FP8 張量核心吞吐量翻倍，其 CUDA 核心性能僅提升了約 33%，導(dǎo)致反量化速度嚴重滯后于計算速度。

該甘特圖截取自我們定制的 Blackwell 注意力核。第一行顯示了張量核心（QKT）的活動情況；第二行顯示了 CUDA 核心的活動情況（數(shù)據(jù)從 TMEM 加載至寄存器，然后執(zhí)行 softmax）。從 TMEM 到寄存器的加載延遲，導(dǎo)致了張量核心出現(xiàn)流水線氣泡。

數(shù)據(jù)顯示，在特定配置下，Blackwell 上的反量化耗時是矩陣乘法本身的 1.76 倍，遠高于 Hopper 上的 1.03 倍。

Hopper 與 Blackwell 上的相對反量化成本。

被忽視的「量化稅」

除了 TMEM 瓶頸，數(shù)據(jù)「量化」過程本身也成了性能殺手。

以一個典型的 MoE 矩陣乘法為例，計算本身可能僅需 1.16 毫秒，但將輸入矩陣量化為 MXFP8 格式并寫回內(nèi)存就需要搬運近 2.9 GB 的數(shù)據(jù)，耗時約 0.44 毫秒，占到計算時間的近 40%。

在反向傳播中，這個開銷因需要轉(zhuǎn)置-量化而翻倍，達到 0.88 毫秒，占比高達 76%。這意味著，如果優(yōu)化不當，MXFP8 帶來的性能提升可能被完全抵消。

此外，現(xiàn)有的開源量化內(nèi)核不僅帶寬利用率低，其生成的縮放因子（scale factor）布局還與 Blackwell 的硬件指令不兼容，需要額外的、拖慢性能的重塑操作。

Cursor 如何從零重寫MoE 層？

面對這些挑戰(zhàn)，并發(fā)現(xiàn)現(xiàn)有的開源庫（如 NVIDIA 的 TransformerEngine）并非最佳選擇，Cursor 團隊選擇放棄高層依賴，使用純 CUDA 和 PTX 匯編語言親自編寫 MoE 層的 GPU 代碼。

優(yōu)化策略

• 擁抱原生硬件指令

他們沒有與 TMEM 架構(gòu)對抗，而是圍繞原生的 tcgen05.mma 指令構(gòu)建內(nèi)核。這使得 GPU 硬件自身能夠處理 MXFP8 所需的縮放，完全消除了 TMEM 和 CUDA 核心之間低效的數(shù)據(jù)移動。

• 設(shè)計高效的數(shù)據(jù)流水線

他們實現(xiàn)了一個復(fù)雜的流水線，采用了諸如「Warp 專精」（將特定任務(wù)分配給不同的線程組）和 2-CTA（協(xié)同線程陣列）模式等技術(shù)。

Warp 專精將特定的任務(wù)分配給不同的線程組（Warp）。例如，Warp 0 負責(zé)從主內(nèi)存加載數(shù)據(jù)到共享內(nèi)存，Warp 1 負責(zé)加載縮放因子，Warp 2 負責(zé)將縮放因子從共享內(nèi)存移至 TMEM，而 Warp 3 則專門負責(zé)啟動矩陣乘法計算。這使得各個環(huán)節(jié)可以高度并行。

2-CTA 模式允許兩個 GPU 流式多處理器（SM）協(xié)同完成單個矩陣乘法，通過共享 B 矩陣來減少內(nèi)存流量，帶來了 15-20% 的性能提升。

• 針對 MoE 工作負載進行優(yōu)化

對于 MoE 訓(xùn)練中特有的分組矩陣乘法，他們應(yīng)用了一種名為「專家級超分組」的 L2 緩存優(yōu)化啟發(fā)式算法。這確保了內(nèi)存訪問模式保持高效，將標準矩陣乘法與分組矩陣乘法之間的性能下降限制在僅 4%。

「秘密武器」：量化內(nèi)核與低精度配方

該團隊開發(fā)了一個自定義的 MXFP8 量化內(nèi)核，他們稱這是目前用于 MoE 訓(xùn)練的最快內(nèi)核。微基準測試顯示，其內(nèi)核持續(xù)的內(nèi)存帶寬超過 6.2 TB/s，相比他們從現(xiàn)有開源工具測得的約 4.5 TB/s 有了顯著提升。

至關(guān)重要的是，他們的內(nèi)核輸出的數(shù)據(jù)內(nèi)存布局與 tcgen05.mma 指令所要求的完全一致，避免了其他工具所必需的、耗時的額外「重塑」步驟。

基于內(nèi)存帶寬利用率的 MXFP8 量化內(nèi)核比較（E4M3，32 塊大小的縮放）。

團隊還確定了一種特定的低精度「配方」，能夠在不影響訓(xùn)練質(zhì)量的情況下提供最高速度。通過使用元素類型為 FP8E4M3、塊大小為 32 的 MXFP8 格式，他們能夠使訓(xùn)練損失的收斂情況與速度慢得多的 BF16 格式幾乎完全匹配。

團隊公布的訓(xùn)練損失曲線顯示，兩種方法幾乎沒有區(qū)別，證明了性能的提升并未以犧牲準確性為代價。

BF16 與 MXFP8 訓(xùn)練損失超過 10k 步：幾乎無法區(qū)分。

更多技術(shù)細節(jié)請閱讀原博客。