全球首個AI CUDA工程師來了！將PyTorch原生實現提速10-100倍

机器之心發表於2025-02-21

原文網址 : https://www.jiqizhixin.com/articles/2025-02-21-7

用 AI 提高 AI 的效率，它們就能變得像人類大腦一樣高效？

我們的大腦只用了 20 瓦的能量就能完成複雜思考，而現代 AI 系統卻需要成排的高功率 GPU 和驚人的電力消耗。這種差距如何縮小？

日本 AI 初創公司 Sakana AI 團隊提出了一個大膽的願景：利用 AI 本身來最佳化 AI。他們開發的「AI CUDA 工程師」是這一理念的具體實踐。

「AI CUDA 工程師」是第一個用於全自動 CUDA 核心發現和最佳化的綜合智慧體框架。這種方法不僅開創性地將進化計算與大型語言模型相結合，更展示了 AI 自我最佳化的巨大潛力。

CUDA 是一個 low-level 軟體層，可直接訪問 NVIDIA GPU 用於平行計算的硬體指令集。CUDA 核心是用 CUDA 語言編寫的在 GPU 上執行的函式。透過直接在 CUDA 核心層編寫指令，工程師可以為 AI 演算法實現更高的效能。然而，使用 CUDA 需要相當多的 GPU 知識，實際上，大多數機器學習演算法都是在 PyTorch 或 JAX 等更高階別的抽象層中編寫的。

^{「AI CUDA 工程師」生成的高度最佳化 CUDA 核心示例。詳情請參見：https://pub.sakana.ai/ai-cuda-engineer}

「AI CUDA 工程師」是一個利用前沿 LLM 的智慧體框架，旨在自動將標準 PyTorch 程式碼轉換為高度最佳化的 CUDA 核心。透過使用進化最佳化，並利用進化計算中的概念，如「交叉」操作和「創新檔案」來發現有前途的「踏腳石」核心，該團隊提出的框架不僅能夠自動化將 PyTorch 模組轉換為 CUDA 核心的過程，而且高度最佳化的 CUDA 核心通常能夠實現顯著更快的執行時間加速。

該團隊相信這項技術能夠實現加速，從而加快 LLM 或其他生成式 AI 模型等基礎模型的訓練和執行（推理），最終使 AI 模型在 NVIDIA 硬體上執行得更快。

「AI CUDA 工程師」能夠生成比常見 PyTorch 操作加速 10-100 倍的 CUDA 核心。它還能生成比生產環境中常用的現有 CUDA 核心快得多的高度最佳化的 CUDA 核心（加速高達 5 倍）。

^{AI CUDA 工程師智慧體框架的高階概述。}

它的執行流程如下：

第 1 和第 2 階段（轉換和翻譯）：「AI CUDA 工程師」首先將 PyTorch 程式碼翻譯成可執行的 CUDA 核心。即使不明確針對這些目標，也能觀察到初始執行時的改進。

第 3 階段（進化最佳化）：受生物進化的啟發，該框架利用進化最佳化（「適者生存」）來確保只生成最佳的 CUDA 核心。此外，該團隊引入了一種新穎的核心交叉提示策略，以互補的方式組合多個最佳化的核心。

第 4 階段（創新檔案）：正如文化進化如何利用我們祖先幾千年文明的知識來塑造我們的人類智慧一樣，「AI CUDA 工程師」也利用從過去的創新和發現中學到的東西（第 4 階段），從已知高效能 CUDA 核心的家族中建立創新檔案，利用以前的踏腳石來實現進一步的翻譯和效能提升。全球首個AI CUDA工程師來了！將PyTorch原生實現提速10-100倍

該專案釋出後，不少研究者給予了很高的評價，比如英偉達高階 AI 研究科學家 Jim Fan 稱這是他最近見過的最酷的自動程式設計智慧體，認為用當前的計算資源來提高未來計算效率，這是最具回報的投資策略，「AutoML is so back!」

不過，也有人發現了問題。比如 NVIDIA 傑出工程師 Bing Xu 指出「AI CUDA 工程師」的技術報告中存在幾個誤導性部分：

Torch C++ 程式碼並不是 CUDA 核心，它在底層是呼叫 CUDNN 庫。
報告重點強調的 Conv3D GroupNorm 示例中，卷積程式碼根本沒有被生成。如果數值計算結果不正確，聲稱的速度提升就沒有意義。
報告中聲稱 WMMA 可以比 PyTorch（CUBLAS）更快，這絕對是錯誤的。很可能是基準測試出現了問題。

看來，這個「AI CUDA 工程師」的效果還有待驗證。

「AI CUDA 工程師」發現的核心執行時加速

「AI CUDA 工程師」穩健地發現了用於常見機器學習操作的 CUDA 核心，其速度比 PyTorch 中的原生和編譯核心快 10-100 倍。該團隊的方法還能將整個機器學習架構轉換為最佳化的 CUDA 核心。下面是幾個完全自主發現的顯著加速：

^{這些最佳化 CUDA 核心的更多詳情可在互動式網站的排行榜上檢視：https://pub.sakana.ai/ai-cuda-engineer/leaderboard}

該團隊的方法為矩陣乘法、常見的深度學習等操作找到了更高效的 CUDA 核心，截至撰寫本文時，它發現的 CUDA 核心在 KernelBench 上實現了 SOTA 的效能。

技術報告和資料集

Sakana AI 釋出了 AI CUDA 的技術報告，整個技術報告有 80 多頁。

技術報告：https://pub.sakana.ai/static/paper.pdf

報告內容如下：

介紹了一個端到端的智慧體工作流，能夠將 PyTorch 程式碼翻譯成可工作的 CUDA 核心，最佳化 CUDA 執行時效能，並自動融合多個核心。
構建了各種技術來增強 pipeline 的一致性和效能，包括 LLM 整合、迭代分析反饋迴圈、本地核心程式碼編輯和交叉核心最佳化。
報告顯示，「AI CUDA 工程師」穩健地翻譯了被考慮在內的 250 個 torch 操作中的 230 多個，並且對大多數核心實現了強大的執行時效能改進。此外，該團隊的方法能夠有效地融合各種核心操作，並且可以超越幾種現有的加速操作。
釋出了一個包含超過 17,000 個經驗證核心的資料集，這些核心涵蓋了廣泛的 PyTorch 操作。

報告還給出了一些發現的 CUDA 核心的顯著例子，這些核心在 AI 模型的關鍵計算操作上實現了顯著的加速。

AI CUDA Engineer 發現的優質核心

利用新的 LLM 驅動的進化核心最佳化程式，研究團隊穩健地獲得了各種考慮因素的加速。更具體地說，在考慮的 229 個任務中，81% 的效能優於 PyTorch 原生執行時。此外，在所有已發現的 CUDA 核心中，有 20% 的核心速度至少是 PyTorch 實現的兩倍。

^{「AI CUDA 工程師」穩健地發現了優於 PyTorch 實現的 CUDA 核心。}

下面展示了一部分核心。它們突顯了「AI CUDA 工程師」可以成功部署的不同操作的多樣性。這包括 normalization 方法、損失函式、特殊矩陣乘法，甚至整個神經網路架構：

^{「AI CUDA 工程師」生成的高度最佳化 CUDA 核心示例。詳情請參見：https://pub.sakana.ai/ai-cuda-engineer}

「AI CUDA 工程師檔案」

17,000 多個經驗證的 CUDA 核心資料集

^{「AI CUDA 工程師檔案」的文字嵌入視覺化顯示，發現的核心可以按任務（例如 MatMul、Pooling、Convolution）和實現策略（展開、融合、向量化）分組。該檔案可公開訪問，可用於 LLM 的下游微調。}

隨論文一起釋出的還有「AI CUDA 工程師檔案」，這是一個由「AI CUDA 工程師」生成的超過 30,000 個 CUDA 核心組成的資料集。它在 CC-By-4.0 許可下發布，可透過 HuggingFace 訪問：https://huggingface.co/datasets/SakanaAI/AI-CUDA-Engineer-Archive。

該資料集包括 torch 參考實現、torch、NCU 和 Clang-tidy 分析資料、每個任務的多個核心、錯誤訊息以及針對 torch 本地和編譯執行時的加速分數。

^{「AI CUDA 工程師檔案」的摘要統計資料，包含超過 30,000 個核心和超過 17,000 個正確驗證的實現。大約 50% 的所有核心都優於 torch 原生執行時。}

研究團隊設想此資料集可以使開源模型的後訓練執行更好的 CUDA 啟用模組。這包括離線強化學習、偏好最佳化和標準監督微調。

在「AI CUDA 工程師檔案」中探索 17,000 多個核心

該團隊還發布了一個互動式網站，用於互動式檢查超過 17,000 個經驗證的核心及其配置檔案，包括 torch、NCU 和 Clang-Tidy 資料：https://pub.sakana.ai/ai-cuda-engineer。

該網站允許探索 230 個任務中的各種高效能核心。它帶有一個自定義排行榜，可用於檢查跨實驗和 LLM 的相關核心。

^{「AI CUDA 工程師」發現的核心排行榜：https://pub.sakana.ai/ai-cuda-engineer/leaderboard}

此外，你還可以視覺化核心，檢索相關核心，下載程式碼以驗證實現和加速，以及檢視獲得的分析資料。最後，可以深入瞭解最佳化實驗。

^{最佳化的例項 Normalization 核心的詳細檢視，包括分析資料、評估指令碼的下載、相關核心和發現實驗詳細資訊。}

侷限性和有趣的意外發現

雖然將進化最佳化與 LLM 結合非常強大，但這種組合有時也會找到意想不到的方法繞過驗證系統。比如，Twitter 使用者 @main_horse 幫助測試 CUDA 核心時就發現了一個有趣的情況：「AI CUDA 工程師」竟然找到了一種「投機取巧」的方法。這個 AI 系統在評估程式碼中發現了一個記憶體漏洞，在一小部分情況下成功避開了正確性檢查：

發現這個問題後，該團隊立即加強了評估框架的安全性，堵住了這類漏洞，並更新了實驗結果。

有趣的是，這並非該團隊第一次遇到類似情況。在之前的「AI 科學家」專案中，AI 也曾找到方法修改並執行自己的評估指令碼。它沒有努力讓程式碼執行得更快，而是直接嘗試修改程式碼來延長超時時間！研究文獻中已經記錄了這種現象：AI 系統常常能找到創造性的解決方案，這些方案往往出人意料，令開發者感到驚訝。

此外，該團隊還發現前沿 LLM 在使用 TensorCore WMMA 方面存在明顯侷限。雖然 LLM 能夠生成基礎的 CUDA 程式碼，但在實現現代 GPU 架構提供的特殊矩陣乘法加速功能時卻常常力不從心。這可能表明 LLM 的訓練資料中缺乏這方面的資訊，或者模型對這類高階硬體最佳化的理解還不夠深入。

隨著前沿 LLM（特別是那些具有強大程式碼推理能力的模型）變得越來越強大，該團隊預計像他們這樣的程式碼最佳化系統將繼續面臨這些挑戰。他們設想未來的發展方向是：人類工程師與程式碼最佳化 AI 系統協同工作，共同創造最佳且最可靠的結果。

「AI CUDA 工程師」的未來影響

AI 革命才剛剛開始，現在只是處於轉型週期的最初階段。該團隊認為，今天的 LLM 是這一代的「大型主機計算機」。現在仍處於 AI 的早期階段，由於市場競爭和全球創新（尤其是那些在資源限制下進行創新的國家 / 地區），這項技術的效率將提高一百萬倍，這是不可避免的。

目前，AI 系統消耗大量資源，如果技術繼續擴充套件而不考慮效率和能源消耗，結果將不可持續。沒有根本原因說明為什麼 AI 系統不能像人類智慧一樣高效（甚至更高效）。該團隊相信，實現這種更高效率的最佳途徑是利用 AI 使 AI 更加高效。

這是 Sakana AI 正在追求的方向，這個專案是使 AI 快一百萬倍的重要一步。就像早期笨重的大型主機計算機向現代計算發展一樣，人類使用 AI 的方式在幾年內將發生重大變化。

^{參考連結：https://sakana.ai/ai-cuda-engineer/}

全球首個AI CUDA工程師來了！將PyTorch原生實現提速10-100倍

相關文章