IBM Spyre加速器的PyTorch原生支持構建之路

IBM Spyre加速器的PyTorch原生支持構建之路

我們最近發布了在PyTorch生態系統中啟用IBM Spyre加速器的2026年上半年路線圖 。 在配套的技術說明中 , 我們介紹了該硬件的特性——32個活躍AI核心、混合精度SIMD-脈動陣列以及可編程數據流設計 。

在這里 , 我們想詳細介紹路線圖的關鍵主題 , 以及我們如何在一個主要由GPU塑造的生態系統中為數據流加速器構建一流的PyTorch支持 。 我們的理念是生態優先——我們利用上游機制 , 最小化自定義代碼 , 并回饋那些讓下一個加速器更容易遵循相同路徑的組件 。
擴展Inductor以支持數據流抽象
我們正在樹外擴展inductor , 以處理數據流加速器所需的抽象 。 三個擴展最為重要:首先 , 我們引入了基于塊的張量布局 , 使編譯器能夠推理Spyre環形連接核心所期望的塊結構數據移動 。 其次 , 我們添加了多核工作劃分過程 , 在編譯期間而非運行時將塊分配到Spyre的32個核心上 。 第三 , 我們添加了暫存器優化——Spyre核心使用顯式管理的片上內存而非硬件緩存 , inductor在調度數據時需要考慮這一點 。 這些擴展將使torch.compile能夠為我們2026年上半年范圍內的每個優先模型生成高效的Spyre代碼 , 從Llama 3.1 8B到Granite 4 Hybrid 30B 。
后端編譯器中間表示的兩階段方法
我們對位于inductor高級圖和Spyre機器碼之間的后端編譯器中間表示采取兩階段方法 。
在第一階段 , SuperDSC(SDSC)將作為后端編譯器中間表示——所有操作降級和代碼生成的單一入口點 。 我們優先模型所需的每個torch操作都可以在SDSC中表達 , 在PyTorch集成層和硬件特定優化之間提供清晰的分離 。
在第二階段 , 我們將過渡到KernelTile IR(KTIR) , 這是一個更符合社區的規范 , 與TileIR等新興計劃更加一致 。 KTIR將泛化塊級表示 , 使其他數據流加速器——不僅僅是Spyre——可以將其用于更低級別的調度和代碼生成 。 我們計劃在今年上半年發布完整的KTIR規范 , 并且我們正在設計的開源調度算法可以適應我們自己硬件之外的應用 。
設備集成與分布式通信
Spyre將完全通過樹外擴展注冊為PyTorch設備:設備生命周期、內存管理、數據傳輸和調度 。 我們的目標是以這種方式處理100%的注冊 , 與直接設備訪問相比開銷低于5% 。 我們計劃將構建的通用原語回饋到PyTorch核心的OpenReg測試基礎設施中 。
對于多卡推理 , 我們通過torch.inductor編譯功能性集合操作(all-reduce、all-gather) , 這將使我們在2026年上半年實現所有優先模型的分布式推理 。 長期來看 , 我們計劃遷移到torch.distributed , 并最終隨著社區通信層的穩定遷移到torch.comms 。
生產推理與vLLM集成
生產推理將通過vLLM運行 。 我們將Spyre作為vLLM平臺插件啟用 , 采用上游模型實現而非維護我們自己的分支 。 我們的優先模型將通過Spyre上的vLLM進行端到端服務 。
兩項優化將推動實際可用性 。 新的Spyre注意力后端將消除同質序列長度約束 , 直接減少Token間延遲 。 上游vLLM中改進的torch.compile工件緩存將把啟動時間降至幾秒鐘 。 我們正在與vLLM社區合作 , 以穩定平臺插件接口 。
測試與持續集成
我們正在構建一個分層測試金字塔 , 將驗證整個堆棧:操作級正確性、inductor編譯和降級、模塊級測試(包括注意力、歸一化和激活)、頂層模型質量和性能 , 以及端到端vLLM推理 。 所有測試都將限定在優先模型范圍內 , 并每晚運行 , 回歸失敗將在數小時內標記 。
我們正在將此持續集成基礎設施構建為對PyTorch生態系統的樹外貢獻 , 建立其他加速器團隊可以采用的模式 。 我們的目標是每晚運行的通過率超過95% , 完整流程在三小時內完成 。
回饋社區
生態優先意味著回饋 , 而不僅僅是在上面構建 。 本半年有三項突出貢獻:我們計劃上游OpenReg原語 , 使樹外設備測試成為PyTorch的一流能力 。 我們正在努力將KTIR泛化為社區規范 , 使數據流加速器共享通用的塊級中間表示 , 而不是各自發明自己的 。 我們將記錄樹外持續集成模式 , 使下一個加速器團隊不必從頭解決基礎設施問題 。
設計文檔和RFC存放在我們的公共存儲庫中 。 我們歡迎參與——無論是審查KTIR規范、試用Spyre跟蹤分析器 , 還是參與關于PyTorch原生加速器支持應該是什么樣子的討論 。
Q&A
Q1:IBM Spyre加速器有什么硬件特點?
A:IBM Spyre加速器擁有32個活躍AI核心、混合精度SIMD-脈動陣列以及可編程數據流設計 。 其核心使用顯式管理的片上內存而非硬件緩存 , 核心之間采用環形連接架構 , 適合塊結構數據移動 。
Q2:PyTorch如何支持Spyre這樣的數據流加速器?
A:通過樹外擴展inductor來處理數據流加速器所需的抽象 , 包括引入基于塊的張量布局、添加多核工作劃分過程以及暫存器優化 。 Spyre將完全通過樹外擴展注冊為PyTorch設備 , 處理設備生命周期、內存管理、數據傳輸和調度 , 目標是與直接設備訪問相比開銷低于5% 。
Q3:KTIR在Spyre的PyTorch支持中扮演什么角色?
【IBM Spyre加速器的PyTorch原生支持構建之路】A:KTIR(KernelTile IR)是第二階段的后端編譯器中間表示 , 它是一個更符合社區的規范 , 與TileIR等新興計劃一致 。 KTIR將泛化塊級表示 , 使其他數據流加速器也可以將其用于更低級別的調度和代碼生成 , 而不僅僅服務于Spyre硬件 。

    推薦閱讀