用 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 研究科學家 Jim Fan 稱這是他最近見過的最酷的自動編程智能體，認為用當前的計算資源來提高未來計算效率，這是最具回報的投資策略 ，「AutoML is so back!」

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

1. Torch C++ 代碼并不是 CUDA 內核，它在底層是調用 CUDNN 庫。
2. 報告重點強調的 Conv3D GroupNorm 示例中，卷積代碼根本沒有被生成。如果數值計算結果不正確，聲稱的速度提升就沒有意義。
3. 報告中聲稱 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://mp.weixin.qq.com/s/uWoySPIqrOWEYAm64w18Kw??