IBM Spyre加速器的PyTorch原生支持構(gòu)建之路

我們最近發(fā)布了在PyTorch生態(tài)系統(tǒng)中啟用IBM Spyre加速器的2026年上半年路線圖。在配套的技術(shù)說(shuō)明中，我們介紹了該硬件的特性——32個(gè)活躍AI核心、混合精度SIMD-脈動(dòng)陣列以及可編程數(shù)據(jù)流設(shè)計(jì)。

在這里，我們想詳細(xì)介紹路線圖的關(guān)鍵主題，以及我們?nèi)绾卧谝粋€(gè)主要由GPU塑造的生態(tài)系統(tǒng)中為數(shù)據(jù)流加速器構(gòu)建一流的PyTorch支持。我們的理念是生態(tài)優(yōu)先——我們利用上游機(jī)制，最小化自定義代碼，并回饋那些讓下一個(gè)加速器更容易遵循相同路徑的組件。

擴(kuò)展Inductor以支持?jǐn)?shù)據(jù)流抽象

我們正在樹(shù)外擴(kuò)展inductor，以處理數(shù)據(jù)流加速器所需的抽象。三個(gè)擴(kuò)展最為重要：首先，我們引入了基于塊的張量布局，使編譯器能夠推理Spyre環(huán)形連接核心所期望的塊結(jié)構(gòu)數(shù)據(jù)移動(dòng)。其次，我們添加了多核工作劃分過(guò)程，在編譯期間而非運(yùn)行時(shí)將塊分配到Spyre的32個(gè)核心上。第三，我們添加了暫存器優(yōu)化——Spyre核心使用顯式管理的片上內(nèi)存而非硬件緩存，inductor在調(diào)度數(shù)據(jù)時(shí)需要考慮這一點(diǎn)。這些擴(kuò)展將使torch.compile能夠?yàn)槲覀?026年上半年范圍內(nèi)的每個(gè)優(yōu)先模型生成高效的Spyre代碼，從Llama 3.1 8B到Granite 4 Hybrid 30B。

后端編譯器中間表示的兩階段方法

我們對(duì)位于inductor高級(jí)圖和Spyre機(jī)器碼之間的后端編譯器中間表示采取兩階段方法。

在第一階段，SuperDSC（SDSC）將作為后端編譯器中間表示——所有操作降級(jí)和代碼生成的單一入口點(diǎn)。我們優(yōu)先模型所需的每個(gè)torch操作都可以在SDSC中表達(dá)，在PyTorch集成層和硬件特定優(yōu)化之間提供清晰的分離。

在第二階段，我們將過(guò)渡到KernelTile IR（KTIR），這是一個(gè)更符合社區(qū)的規(guī)范，與TileIR等新興計(jì)劃更加一致。KTIR將泛化塊級(jí)表示，使其他數(shù)據(jù)流加速器——不僅僅是Spyre——可以將其用于更低級(jí)別的調(diào)度和代碼生成。我們計(jì)劃在今年上半年發(fā)布完整的KTIR規(guī)范，并且我們正在設(shè)計(jì)的開(kāi)源調(diào)度算法可以適應(yīng)我們自己硬件之外的應(yīng)用。

設(shè)備集成與分布式通信

Spyre將完全通過(guò)樹(shù)外擴(kuò)展注冊(cè)為PyTorch設(shè)備：設(shè)備生命周期、內(nèi)存管理、數(shù)據(jù)傳輸和調(diào)度。我們的目標(biāo)是以這種方式處理100%的注冊(cè)，與直接設(shè)備訪問(wèn)相比開(kāi)銷低于5%。我們計(jì)劃將構(gòu)建的通用原語(yǔ)回饋到PyTorch核心的OpenReg測(cè)試基礎(chǔ)設(shè)施中。

對(duì)于多卡推理，我們通過(guò)torch.inductor編譯功能性集合操作（all-reduce、all-gather），這將使我們?cè)?026年上半年實(shí)現(xiàn)所有優(yōu)先模型的分布式推理。長(zhǎng)期來(lái)看，我們計(jì)劃遷移到torch.distributed，并最終隨著社區(qū)通信層的穩(wěn)定遷移到torch.comms。

生產(chǎn)推理與vLLM集成

生產(chǎn)推理將通過(guò)vLLM運(yùn)行。我們將Spyre作為vLLM平臺(tái)插件啟用，采用上游模型實(shí)現(xiàn)而非維護(hù)我們自己的分支。我們的優(yōu)先模型將通過(guò)Spyre上的vLLM進(jìn)行端到端服務(wù)。

兩項(xiàng)優(yōu)化將推動(dòng)實(shí)際可用性。新的Spyre注意力后端將消除同質(zhì)序列長(zhǎng)度約束，直接減少Token間延遲。上游vLLM中改進(jìn)的torch.compile工件緩存將把啟動(dòng)時(shí)間降至幾秒鐘。我們正在與vLLM社區(qū)合作，以穩(wěn)定平臺(tái)插件接口。

測(cè)試與持續(xù)集成

我們正在構(gòu)建一個(gè)分層測(cè)試金字塔，將驗(yàn)證整個(gè)堆棧：操作級(jí)正確性、inductor編譯和降級(jí)、模塊級(jí)測(cè)試（包括注意力、歸一化和激活）、頂層模型質(zhì)量和性能，以及端到端vLLM推理。所有測(cè)試都將限定在優(yōu)先模型范圍內(nèi)，并每晚運(yùn)行，回歸失敗將在數(shù)小時(shí)內(nèi)標(biāo)記。

我們正在將此持續(xù)集成基礎(chǔ)設(shè)施構(gòu)建為對(duì)PyTorch生態(tài)系統(tǒng)的樹(shù)外貢獻(xiàn)，建立其他加速器團(tuán)隊(duì)可以采用的模式。我們的目標(biāo)是每晚運(yùn)行的通過(guò)率超過(guò)95%，完整流程在三小時(shí)內(nèi)完成。

回饋社區(qū)

生態(tài)優(yōu)先意味著回饋，而不僅僅是在上面構(gòu)建。本半年有三項(xiàng)突出貢獻(xiàn)：我們計(jì)劃上游OpenReg原語(yǔ)，使樹(shù)外設(shè)備測(cè)試成為PyTorch的一流能力。我們正在努力將KTIR泛化為社區(qū)規(guī)范，使數(shù)據(jù)流加速器共享通用的塊級(jí)中間表示，而不是各自發(fā)明自己的。我們將記錄樹(shù)外持續(xù)集成模式，使下一個(gè)加速器團(tuán)隊(duì)不必從頭解決基礎(chǔ)設(shè)施問(wèn)題。

設(shè)計(jì)文檔和RFC存放在我們的公共存儲(chǔ)庫(kù)中。我們歡迎參與——無(wú)論是審查KTIR規(guī)范、試用Spyre跟蹤分析器，還是參與關(guān)于PyTorch原生加速器支持應(yīng)該是什么樣子的討論。

Q&A

Q1：IBM Spyre加速器有什么硬件特點(diǎn)？

A：IBM Spyre加速器擁有32個(gè)活躍AI核心、混合精度SIMD-脈動(dòng)陣列以及可編程數(shù)據(jù)流設(shè)計(jì)。其核心使用顯式管理的片上內(nèi)存而非硬件緩存，核心之間采用環(huán)形連接架構(gòu)，適合塊結(jié)構(gòu)數(shù)據(jù)移動(dòng)。

Q2：PyTorch如何支持Spyre這樣的數(shù)據(jù)流加速器？

A：通過(guò)樹(shù)外擴(kuò)展inductor來(lái)處理數(shù)據(jù)流加速器所需的抽象，包括引入基于塊的張量布局、添加多核工作劃分過(guò)程以及暫存器優(yōu)化。Spyre將完全通過(guò)樹(shù)外擴(kuò)展注冊(cè)為PyTorch設(shè)備，處理設(shè)備生命周期、內(nèi)存管理、數(shù)據(jù)傳輸和調(diào)度，目標(biāo)是與直接設(shè)備訪問(wèn)相比開(kāi)銷低于5%。

Q3：KTIR在Spyre的PyTorch支持中扮演什么角色？

A：KTIR（KernelTile IR）是第二階段的后端編譯器中間表示，它是一個(gè)更符合社區(qū)的規(guī)范，與TileIR等新興計(jì)劃一致。KTIR將泛化塊級(jí)表示，使其他數(shù)據(jù)流加速器也可以將其用于更低級(jí)別的調(diào)度和代碼生成，而不僅僅服務(wù)于Spyre硬件。