斯坦福讓“GPU高速運轉”的新工具火了，比FlashAttention2更快

AIGC動態歡迎閱讀

原標題：斯坦福讓“GPU高速運轉”的新工具火了，比FlashAttention2更快
關鍵字：張量,內存,研究人員,寄存器,硬件
文章來源：量子位
內容字數：0字

內容摘要：

西風 發自 凹非寺量子位 | 公眾號 QbitAIAI算力資源越發緊張的當下，斯坦福新研究將GPU運行效率再提升一波——
內核只有100行代碼，讓H100比使用FlashAttention-2，性能還要提升30%。
怎么做到的？
研究人員從“硬件實際需要什么？如何滿足這些需求？”這兩個問題出發，設計了 一個嵌入式CUDA DSL工具，名為ThunderKittens（暫且譯為雷貓）。
雷貓可簡化AI內核的編寫，同時充分利用底層硬件能力。
具體來說，雷貓的主要抽象是寄存器和共享內存中的小型張量塊（tile），和目前GPU中對小矩陣乘法的優化相匹配。
通過操作這些tile，開發者可相對簡單地編寫代碼，充分利用張量核心、異步數據傳輸和共享內存等硬件特性。
使用雷貓實現的注意力機制內核，代碼量少且能實現很高的硬件利用率，性能超過直接使用底層庫（如Cutlass）。
詳細討論過程以及雷貓是怎么設計出的，研究人員以“GPUs Go Brrr”為題，發在了斯坦福Hazy Research的Blog網站上。
網友們對此討論也十分熱烈。
有網友表示讀這篇Blog時，讓他想起了初次了解超標量CPU架構時的

原文鏈接：斯坦福讓“GPU高速運轉”的新工具火了，比FlashAttention2更快

聯系作者

文章來源：量子位
作者微信：QbitAI
作者簡介：追蹤人工智能新趨勢，關注科技行業新突破