機(jī)器之心報(bào)道

編輯：+0

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Hopper 與 Blackwell 上的相對(duì)反量化成本。

被忽視的「量化稅」

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

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

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

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

Cursor 如何從零重寫MoE 層？

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

優(yōu)化策略

• 擁抱原生硬件指令

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

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

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

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

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

• 針對(duì) MoE 工作負(fù)載進(jìn)行優(yōu)化

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

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

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

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

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

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

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

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

更多技術(shù)細(xì)節(jié)請(qǐng)閱讀原博客。