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 步：幾乎無法區分。
更多技術細節請閱讀原博客。
