研究

利用多代理系統將 GPU kernel 加速 38%

我們的多代理系統在短短 3 週內自主最佳化了 235 個 NVIDIA Blackwell 200 GPU 的 CUDA kernel,相較基準達成 38% 的幾何平均加速。

Wilson Lin – Cursor;Sahil Modi、Yuan Zhang 與 Edward Lin – NVIDIA閱讀時間 4 分鐘

過去幾個月來,我們一直在開發一套能夠自主打造、維護並部署複雜軟體的多代理系統。作為這項工作的一部分,我們持續在各種領域測試這套系統,包括讓它從零打造一個瀏覽器,以及在 First Proof benchmark 上解出研究等級的數學問題。

最近,我們開始與 NVIDIA 合作,迎接一項新的挑戰:將多代理工具鏈套用於最佳化 CUDA kernel。這些都是艱鉅的技術問題,而且在現實世界中影響重大:CUDA kernel 是支撐 NVIDIA GPU 上 AI (人工智慧) 模型訓練與推論的核心軟體。更快的 kernel 代表更高的 GPU 利用率、更低的能耗、更低的延遲,以及更低的每 token 成本——讓供應商能同時為更多使用者提供更大、能力更強的模型。

我們的多代理工具鏈在 235 個問題上自主運作了三週。這套系統從零打造並最佳化 Blackwell GPU kernel,一路深入到組合語言層級,達成 38% 的幾何平均加速。

這種幅度的效能提升,通常只有非常資深的 kernel 工程師投入數月甚至數年的時間才做得到。而這套多代理系統在短短幾週內就完成了,解決了一長串以現有方法難以實際處理的 kernel 問題。

將 Kernel 最佳化作為測試代理系統能力的方法

評估長時間執行的多代理系統,最佳方式之一就是給它們開放式的最佳化問題,而連我們自己也不知道正確答案是什麼。Kernel 最佳化問題正符合這項條件:它們提供可衡量的目標,讓系統能反覆迭代並持續最佳化,而不是只針對已知且簡單的差異進行調整。

如今,工程師會將模型拆解成個別的數學運算,並分別調校每一項來最佳化 Kernel。這讓問題更容易處理,但也留下了尚未發揮的效能空間,因為零碎式的最佳化會錯失跨整個系統同步最佳化所能帶來的潛在收益。迄今為止,GPU 效能一直受限於我們無法跳脫這些人工簡化,去探索更完整的解空間。

透過這項實驗,我們想看看多代理系統是否能突破這些限制,探索更廣泛的解空間,產出更快的 Kernel。

用於問題生成與基準測試的 SOL-ExecBench

NVIDIA 使用 SOL-ExecBench,從超過 124 個正式環境中的開源模型 (例如 Deepseek、Qwen、Gemma、Kimi 和 Stable Diffusion) 生成了 235 個最佳化問題。不同於合成資料或玩具級 kernel,每個問題都是真實世界中訓練或推論工作負載所面臨的限制條件,涵蓋多種模型架構:LLM、擴散、視覺、音訊、影片,以及多模態混合模型。

SOL-ExecBench 在 LLM、擴散、視覺、音訊、影片和多模態架構中生成了 235 個最佳化問題SOL-ExecBench 在 LLM、擴散、視覺、音訊、影片和多模態架構中生成了 235 個最佳化問題

我們也使用 SOL-ExecBench,在 27 張 NVIDIA Blackwell 200 GPU 上對多代理 kernel 解法進行基準測試。SOL-ExecBench 是一種有效的評測工具,會將 kernel 效能與現有軟體基準,以及理論上的硬體效能極限進行比較。如果代理採用像快取這類作弊手法,交出超出 B200 所能支援的效能,流程就會判定該結果無效。

我們如何執行這項實驗

這個多代理系統透過部署一個規劃者代理,根據效能指標在自主工作代理之間分派並重新平衡工作,因此在單次執行中解決了全部 235 個 GPU kernel 最佳化問題。

整個協調協定都寫在單一的 Markdown 檔案中,明確規定了輸出格式、規則與測試。這個多代理系統也在執行過程中自行學會呼叫基準測試流程,形成一個迴圈,讓系統在沒有任何開發者介入的情況下,持續測試、除錯並最佳化 kernel。

多代理系統的自主除錯、測試與最佳化迴圈多代理系統的自主除錯、測試與最佳化迴圈

為了更準確評估這個多代理系統的能力,我們要求它在兩次彼此獨立的執行中,以兩種語言撰寫解法,分別代表 GPU 抽象層級光譜的兩端:

  • CUDA C with inline PTX,讓代理能直接存取暫存器與 ISA 層級指令,用來測試系統是否能在最低層級理解硬體。
  • CuTe DSL,提供高階且可組合的抽象,但在公開訓練資料中幾乎不存在,用來測試系統是否能僅根據提供的文件學會全新的 API。

加速 38%,其中 19% 的最佳化結果提升超過 2 倍

我們以兩種方式呈現多代理系統的效能:

  • 幾何平均加速:相較於由單一代理最佳化的 PyTorch 程式碼基準,其幾何平均加速比。
  • Speed-of-Light (SOL) scores:以對數曲線表示某個解決方案相較於理論硬體極限的表現。0.5 分代表最佳化後的 PyTorch 基準,1.0 則代表效能極限。
Speed-of-Light 分數刻度:0.5 代表最佳化後的 PyTorch 基準,1.0 則是理論硬體極限Speed-of-Light 分數刻度:0.5 代表最佳化後的 PyTorch 基準,1.0 則是理論硬體極限

我們的多代理系統在 235 個問題中有 149 個的表現優於基準 (63%) ,幾何平均比率為 1.38 倍 (幾何平均加速 38%) 。

多代理系統在 235 個 kernel 最佳化問題中達成 1.38 倍的幾何平均加速多代理系統在 235 個 kernel 最佳化問題中達成 1.38 倍的幾何平均加速

在 235 個問題中,有 45 個 (19%) 的最佳化幅度相較於基準超過 2 倍。您可以在這個公開儲存庫中查看我們系統開發出的最終解決方案。

各問題加速比的分佈,其中 19% 的提升超過 2 倍各問題加速比的分佈,其中 19% 的提升超過 2 倍

針對不同問題的不同最佳化策略

為了展現系統面對不同類型問題時的適應性,我們著重介紹三個案例;在這些案例中,系統自然而然地發展出各不相同的最佳化策略。

BF16 分組查詢注意力與分頁式 Prefill

帶有分頁式 prefill 的分組查詢注意力,是現代 LLM 推論堆疊中常見的提示階段操作。經過充分最佳化的實作,能在相同的 GPU VRAM 上支援更長的上下文、更高的並行度,以及更好的吞吐量。

該代理使用 CUDA C++,針對從 SGLang 中 Llama 3.1 8B 的推論流程擷取出的這個注意力問題進行最佳化。隨著代理持續迭代這個 kernel,它成功運用了特定的硬體指令來處理記憶體載入與數學運算,透過持久化 kernel 導入改進的排程,並針對輸入大小做了高度最佳化。

我們將多代理系統的自訂 kernel,與 FlashInfer 函式庫中人工最佳化的基準實作進行比較。我們發現,該系統產出的解決方案已接近硬體極限,SOL 分數達到 0.9722,相較於基準實作,幾何平均加速了 84%。

我們也替換了 SGLang 中現有的 kernel,並觀察到在 Llama 3.1 8B 上,首個 token 時間 (TTFT) 加速了 3%。鑑於這個注意力問題依服務組態不同,會占 prefill 流程 2–5% 的時間,我們認為這是不可忽視的端對端加速。

BF16 Grouped Query Attention 的最佳化軌跡,達到 0.9722 的 SOL 分數與 84% 的幾何平均加速BF16 Grouped Query Attention 的最佳化軌跡,達到 0.9722 的 SOL 分數與 84% 的幾何平均加速

具 Gating 的 NVFP4 MoE Linear

這個問題呈現了 Qwen3 這類 Mixture-of-Experts 模型中常見的雙 kernel 模式,但需注意的是,輸入張量和中間乘積輸出都被量化為 NVFP4 (4 位元浮點) 。

代理正確找出了量化部分是主要瓶頸,並因此將 scale 計算與捨入融合在一起。它沒有在量化時先縮放再捨入,而是使用預先計算好的 threshold buckets,直接將 FP32 值對應到 FP4 編碼;這之所以可行,是因為 NVFP4 只有 16 種可能值。接著,它將這些最佳化套用到更大的測試案例上。

代理最終超越了經過最佳化的 PyTorch 基準,達成 39% 的幾何平均加速和 0.58 的 SOL 分數。

NVFP4 MoE Linear 最佳化軌跡,達成 39% 的幾何平均加速和 0.58 的 SOL 分數NVFP4 MoE Linear 最佳化軌跡,達成 39% 的幾何平均加速和 0.58 的 SOL 分數

BF16 矩陣乘法

矩陣乘法向來是最難最佳化的問題之一,因為它需要對各種硬體單元及其排程有深刻的理解。要寫出效能完全發揮的矩陣乘法 kernel (GEMM) ,需要用到 inline PTX (類似組合語言) 、管線化,以及在 kernel 內進行 staging。因此,撰寫高速 GEMM 歷來幾乎都是少數經驗非常豐富的 kernel 專家的專業領域。

Cursor 多代理系統從零產生了一個特化的 CUDA C++ GEMM kernel,效能已相當接近 NVIDIA cuBLAS 函式庫中經過精細調校的人工基準實作 (86%) 。系統之所以能達到這個結果,是因為它能獨立學會使用 Blackwell 專屬指令、針對硬體最佳化記憶體讀寫,並進一步針對特定形狀做極致最佳化。

此外,在小 M 的測試案例上——這對 LLM 推論的 decode 尤其重要——這個多代理系統 kernel 的表現最高比函式庫快 9%。這項結果顯示,多代理系統很快就可能連在最困難的 kernel 問題上都超越領域專家。

BF16 矩陣乘法結果:達到人工最佳化 cuBLAS 效能的 86%,並在小 M 案例中最高超越 9%BF16 矩陣乘法結果:達到人工最佳化 cuBLAS 效能的 86%,並在小 M 案例中最高超越 9%

用於打造軟體的多代理系統

雖然多代理工具鏈相較於基準帶來了 38% 的幾何平均加速,但 SOL 分數的中位數仍只有 0.56,顯示仍有相當大的最佳化空間。我們認為,只要有更多算力,多代理解決方案還能大幅改進,因為當時我們只有 27 張 GPU,卻要同時執行數百個問題和代理。這限制了我們充分發揮多代理系統潛力的能力。如果有更多 GPU,系統就能探索更深、更具新意的解法。

軟體開發中最具挑戰性的任務往往是開放式的,沒有明確解法。單一代理系統在這類情況下往往力有未逮,因為模型最擅長的是範圍明確、且已在訓練過程中見過的任務。我們將 kernel 最佳化實驗視為進一步的驗證:多代理架構將迅速成為打造軟體的預設方式,因為它們能處理遠遠超出訓練資料分布的全新問題。

我們正在研究的這些技術,很快就會融入 Cursor 的核心產品。如果你有興趣投入多代理協作中的棘手問題,Cursor 團隊很樂意收到你的來信:hiring@cursor.com

分類於: 研究

作者s: Wilson Lin, Sahil Modi, Yuan Zhang & Edward Lin