跳到主要內容
AI News HubLIVE
站內改寫1 分鐘閱讀

Kernel Forge:用於基於LLM的CUDA內核生成與優化的智能體工具

文章摘要

Kernel Forge是一個開源端到端智能體框架,可接受任何未經修改的PyTorch模型,利用蒙特卡洛樹搜索探索多個優化路徑,並配備圖形用户界面。在NVIDIA DGX Spark上對四個模型進行測試,僅需50次優化迭代,即可在14個內核上超越PyTorch eager模式,最高加速2.83倍。

來源arXiv AI作者: Joshua Brodsky, Dhravid Kumar, Savini Kashmira, Jayanaka Danatanarayana, Jason Mars, Krisztian Flautner, Lingjia Tang
Kernel Forge:用於基於LLM的CUDA內核生成與優化的智能體工具
報告錯誤

更正渠道尚未開通,可先複製下方文章資訊留存。

查看更正說明
直接讀正文

隨着機器學習模型越來越多地嵌入日常軟件,其大部分運行時間都花在少量計算內核上,如矩陣乘法、卷積和歸一化。優化這些內核是減少延遲和成本的最直接方法之一,但傳統上需要專家工程師手動編寫底層GPU代碼。基於大型語言模型(LLM)的智能體系統現在能夠以更少的人力生成和優化內核,但現有工具大多在隨機生成的張量和孤立內核上進行評估,輸出需要開發者手動重新集成的獨立CUDA代碼,主要針對LLM PyTorch模型,並且對檢查與調試結果的支持有限。

針對這些不足,研究團隊提出了Kernel Forge,這是一個開源、端到端的智能體工具,能夠接受任何未經修改的PyTorch模型。與現有工具不同,Kernel Forge支持視覺、擴散和LLM工作負載,採用蒙特卡洛樹搜索(MCTS)探索多條優化路徑,而不是單一的精煉鏈。它還附帶圖形用户界面,用於監控進度、檢查候選內核和調試失敗。

研究團隊在搭載GB10 GPU的NVIDIA DGX Spark上,對涵蓋視覺、擴散和LLM工作負載的四個PyTorch模型進行了評估。每個內核僅進行50次優化迭代,Kernel Forge優化了14個內核,使其性能超越PyTorch eager模式:在ResNet-50上,adaptive_avgpool2d加速1.52倍;在Stable Diffusion 3.5 Medium上,group_norm加速1.70倍;在Gemma 4 E2B上,softmax加速2.83倍;在Qwen 3.5 35B-A3B上,softmax加速1.54倍。這些結果表明,Kernel Forge能夠有效降低內核優化的門檻,為AI部署帶來顯著的性能提升。Kernel Forge的代碼已開源,可供研究者和開發者使用。

展開要點與分析

文章情報

工程師進階

要點

  • Kernel Forge接受任意未經修改的PyTorch模型,支持視覺、擴散和LLM工作負載。
  • 使用蒙特卡洛樹搜索(MCTS)探索多個優化路徑,而非單一線性的精煉鏈。
  • 通過圖形用户界面監控進度、檢查候選內核並調試失敗。
  • 在ResNet-50、Stable Diffusion 3.5 Medium、Gemma 4 E2B和Qwen 3.5 35B-A3B上實現了1.52倍至2.83倍的加速。

要點與分析由自動化流程生成,可能有誤,請結合原始來源核實。