上月底，創(chuàng)業(yè)公司Groq的產(chǎn)品一夜爆火。

憑借自研的硬件加速器LPU，達成了500個token/s的神級推理速度，當場秒殺了ChatGPT。

圖片

Groq提供的響應速度刷新了人們的認知，而這要歸功于背后的語言處理單元硬件LPU（language processing unit hardware）。

Groq的研發(fā)團隊在LPU上應用了創(chuàng)新的硬件架構設計，并配套了強大的編譯器。

下面讓我們跟隨Substack的專欄作家Abhinav Upadhyay一起，一步步揭開Groq LPU底層架構的神秘面紗。

Groq LPU的神秘面紗

到目前為止，Groq并沒有給出任何關于LPU本身的論文，但在過去幾年中，他們發(fā)表了下面兩篇論文：

論文地址：https://dl.acm.org/doi/10.1109/ISCA45697.2020.00023

圖片

論文地址：https://dl.acm.org/doi/abs/10.1145/3470496.3527405

兩篇工作分別在2020年和2022年發(fā)表在計算機體系結構頂會ISCA上，后一篇還是獲獎論文。

這兩篇文章解釋了Groq的張量流處理器（TSP）的設計和實現(xiàn)，以及他們?nèi)绾问褂肨SP構建分布式推理引擎。

盡管沒有正式聲明，但LPU很可能是基于這個分布式系統(tǒng)來進行設計和擴展的。

那么，我們就首先詳細分解一下TSP及其編譯器的架構，然后以此為基礎來分析Groq如何使用這些TSP，構建可靠且高吞吐量的分布式AI推理引擎。

TSP的架構與傳統(tǒng)的CPU或GPU芯片有很大不同，主要目的是為了讓TSP硬件更具確定性。

這里就先要提一嘴CPU或GPU的不確定性。

CPU和GPU微架構中的非確定性

基于微架構的設計，在CPU和GPU上執(zhí)行指令是不確定的，——即無法保證特定指令何時執(zhí)行、完成需要多長時間以及何時提供結果。

舉個例子，現(xiàn)代CPU一般具有如下設計：

- 超標量（Super scalar architecture）：每個周期能夠發(fā)出多條指令；

- 亂序執(zhí)行（Out-of-order execution）：以任意順序執(zhí)行指令；

- 預測執(zhí)行（Speculative execution）：對于分支，它會猜測分支條件是真是假，并提前預測執(zhí)行該分支以提高吞吐量（當然如果猜錯了，就需要放棄并返回另一條分支）； 

- 指令流水線（Instruction pipelining）：將指令分為多個階段，以流水線的方式執(zhí)行，再次提高了指令吞吐量；

- 多級緩存（Multiple levels of caches）：CPU有2到3級緩存，可以減少從內(nèi)存加載數(shù)據(jù)帶來的延遲。

所有這些都使得CPU中指令執(zhí)行的順序和時間不確定且難以推理。

而GPU還有其他一些非確定性因素，包括緩存、共享和全局內(nèi)存、動態(tài)資源分區(qū)等。

非確定性帶來的問題是，我們很難推理程序的性能，也很難保證最壞情況下的性能限制。

因此，Groq為TSP提出了一個全新的設計，高度并行，且沒有不確定行為。這消除了硬件的復雜性，使編譯器能夠獲得更大的權力，精確調度和控制指令的執(zhí)行，保證對程序性能的限制。

下面，讓我們從內(nèi)部了解TSP的架構是什么樣子的。

TSP架構

TSP的硬件設計與CPU或GPU的設計形成鮮明對比。傳統(tǒng)的多核芯片采用平鋪架構，在下圖（a）中，每個小方塊（tile）代表一個處理核心。

核心由一組功能單元組成，負責執(zhí)行不同類型的計算（算術運算、內(nèi)存運算、邏輯運算、指令控制等）。

而TSP的設計師將這種傳統(tǒng)設計徹底顛覆了。他們將功能單元移到核心之外，以2d網(wǎng)格方式排列。

網(wǎng)格的每一列只包含特定類型的功能單元，稱為切片（slice）。下圖顯示了傳統(tǒng)多核芯片和TSP在設計上的區(qū)別。

TSP具有以下功能切片：

MXM：用于執(zhí)行矩陣運算 

SXM：用于對矢量進行移位和旋轉操作 

MEM：內(nèi)存讀/寫運算 

VXM：向量上的算術運算 

ICU：指令控制單元，這個有點特殊，就是上圖（b）底部那一條水平的藍色條，它負責獲取和調度指令并在其他切片上執(zhí)行。

在了解了TSP的架構之后，讓我們將注意力轉移到它的核心：指令執(zhí)行。

TSP中的指令執(zhí)行

TSP以SIMD（單指令多數(shù)據(jù)）方式執(zhí)行指令。每個功能切片由20個tile組成，每個tile能夠處理16個數(shù)。因此，一個完整的切片可以處理并生成最大320個元素的向量。

這些切片以生產(chǎn)者——消費者的方式進行交互。

當從內(nèi)存中讀取向量時，會為其分配流ID（介于0到31之間）和流向（東或西）。每個切片都可以自由處理流并生成新的結果流，也可以讓流按原樣流向下一個相鄰切片。

為了有效地處理完整的向量，指令以流水線方式執(zhí)行，如下圖所示：

圖片

TSP中指令的流水線執(zhí)行會導致流在切片之間交錯移動。上圖的黑色塊描繪了流在切片中不同時間戳的移動。

當然了，想要愉快地執(zhí)行指令，必然少不了編譯器和指令集（ISA）設計。

TSP的編譯器和ISA

TSP的設計人員簡化了硬件，所以壓力就給到了編譯器這邊。編譯器需要精確地調度指令和數(shù)據(jù)流，以正確執(zhí)行給定的程序，并以最有效的方式執(zhí)行。

編譯器有權訪問TSP硬件的以下狀態(tài)：

- 320個通道的編程抽象：TSP芯片中的每個tile都能夠以SIMD方式在矢量的16個單元（16個通道）上運行。垂直切片由20個這樣的tile組成，因此總共有320個SIMD通道可供執(zhí)行；

- 144個獨立指令隊列：芯片上有144個指令隊列，每個周期能夠發(fā)出一條或多條指令。編譯器可以完全控制每個隊列中的程序順序；

- 每個通道64個邏輯流：每個通道可以訪問64個邏輯流，可用于移動操作數(shù)或結果，其中32個可用于向東移動數(shù)據(jù)，而另外32個用于向西移動數(shù)據(jù)；

- 220M全局共享SRAM。

由于TSP硬件中沒有非確定性行為，因此編譯器可以準確了解每條指令的延遲，以及程序中的數(shù)據(jù)流（DNN的計算圖等）。

編譯器識別計算任務之間的依賴關系，并分配到TSP的可用功能單元上并行執(zhí)行。

TSP編程模型依賴于兩個關鍵要素：

硬件中的確定性數(shù)據(jù)路徑 

通過ISA獲得的有關指令延遲的信息

編譯器的后端可以跟蹤片上任何流的位置和使用時間，稱為軟件定義硬件。

從TSP擴展到LPU

TSP是LPU的基礎單元。許多TSP以機架的形式組合在一起，形成一個能夠提供大量吞吐量的分布式系統(tǒng)。

設計多TSP系統(tǒng)

與TSP一樣，分布式多TSP系統(tǒng)的設計目標也圍繞著確定性數(shù)據(jù)流和指令執(zhí)行，以及節(jié)點之間的低延遲通信。

分布式TSP系統(tǒng)的設計從節(jié)點開始。節(jié)點由機箱內(nèi)8個TSP設備組成。這些設備中的每一個都由11個引腳組成，其中7個引腳用于將每個TSP設備連接到節(jié)點中的其他7個TSP設備，其余4個引腳用于形成全局鏈接。

節(jié)點中的每個設備都有4個全局鏈路，總共有32個全局鏈路，共同構成了一個32個虛擬端口的高基數(shù)路由器（high-radix router）。

高基數(shù)路由器支持大量連接、高帶寬和高性能，這正是高性能分布式系統(tǒng)所需要的。

將9個這樣的TSP節(jié)點和8個TSP組合成一個機架。機架中的每個節(jié)點都有32個端口，因此機架總共有288個全局端口。

其中144個端口在機架內(nèi)本地使用，以便在機架內(nèi)快速傳輸數(shù)據(jù)，其余144個端口用于連接到其他機架。

最大配置的系統(tǒng)可以支持145個相互互連的機架，包括10440個TSP，系統(tǒng)中任何兩個TSP之間最多有5個hops。

在基于TSP的分布式系統(tǒng)中實現(xiàn)確定性

在這種擴展的分布式系統(tǒng)制度中，單個TSP的功能單元充當大規(guī)模并行處理器的單個處理核心。TSP的計算模型基于確定性硬件，所以整個分布式系統(tǒng)也應具有同樣的確定性。

使用硬件對齊計數(shù)器同步TSP的時鐘

每個TSP設備都包含一個稱為硬件對齊計數(shù)器（HAC）的硬件計數(shù)器，溢出周期為256。TSP通過以下步驟使用它來相互同步：

- 當兩個TSP互連時，其中一個TSP將其HAC值傳輸給對方。然后，對方將該值返回發(fā)送方。發(fā)送方觀察當前HAC值與返回值之間的差值。

- 這個差值就代表了兩個設備之間鏈路的延遲。此過程重復多次，得到兩個TSP之間的平均鏈路延遲。

- 之后，兩個設備以父子關系排列。父級定期將當前HAC值發(fā)送給子級。子級將平均鏈路延遲與自己的HAC值相加，并與自己的HAC值進行比較。

- 兩個值之間的差值表示由于連續(xù)時鐘漂移而導致的初始未對準。然后子級調整其HAC值以減小此差異。在多次重復此過程后，兩個TSP的HAC值會收斂在一個小鄰域內(nèi)，表示鏈路延遲的抖動。

- 協(xié)議允許兩個TSP相互同步，并且可以通過在網(wǎng)絡中建立生成樹來擴展TSP多跳網(wǎng)絡。

初始計劃調整

程序在多TSP系統(tǒng)上執(zhí)行之前，需要對齊所有TSP，以正確調度整個系統(tǒng)的數(shù)據(jù)流和指令執(zhí)行。這涉及到以下機制：

- 在單個TSP級別，有幾個獨立的功能單元和144個獨立的指令隊列。為了同步它們，TSP支持SYNC和NOTIFY指令。SYNC指令將所有指令隊列置于停放狀態(tài)，其中一個隊列充當通知程序。當通知器發(fā)出 NOTIFY指令時，該指令被廣播到芯片上的所有隊列，此時它們被同步并恢復操作。 

- 對于多TSP系統(tǒng)，兩個TSP使用HAC相互同步，另外每個TSP都支持DESKEW指令，用于停止處理任何后續(xù)指令，直到TSP的HAC溢出。 

- 要擴展多跳系統(tǒng)，可以在生成樹的每個hop上重復執(zhí)行以上方案。

運行時重新同步

雖然TSP在程序開始時進行一次性同步，但它們也需要在程序執(zhí)行期間重新同步，因為每個TSP都有自己獨立的時鐘源。

為此，TSP使用更輕量級的方案。除了HAC之外，每個TSP都有一個軟件對齊計數(shù)器（SAC），其溢出周期與HAC相同。

但是，SAC在TSP之間不同步，SAC只是計算TSP的時鐘周期。HAC值表示分布式系統(tǒng)的全局時間，而SAC表示本地時間。因此，HAC和SAC值之間的增量決定了累積漂移。

為了重新同步本地和全局時間，TSP執(zhí)行一條RUNTIME_DESKEW指令。系統(tǒng)中的每個TSP同時執(zhí)行該指令，根據(jù)累積的漂移調整全局時間與本地時間。

編譯器在軟件計劃網(wǎng)絡中的作用

到目前為止，編譯器能夠對TSP內(nèi)以及整個網(wǎng)絡中的數(shù)據(jù)移動進行周期準確的了解。編譯器知道在源TSP上注入向量的確切時間以及它到達目標TSP的確切時間，稱為軟件計劃網(wǎng)絡。

編譯器不是動態(tài)管理數(shù)據(jù)流，而是在編譯時靜態(tài)解析所有內(nèi)容。

已知流量模式

對于深度學習模型，編譯器可以根據(jù)模型的靜態(tài)計算圖推斷數(shù)據(jù)流。編譯器還可以在網(wǎng)絡中可用的TSP設備之間自動分配計算任務。

因此，編譯器會計算每個子任務的精確執(zhí)行時間以及各層之間的激活交換。這使得并行分解步驟顯式，并完全由編譯器控制。

計劃的數(shù)據(jù)流

在傳統(tǒng)的網(wǎng)絡系統(tǒng)中，通過網(wǎng)絡的數(shù)據(jù)包流由硬件管理，硬件在感應到網(wǎng)絡中的負載時會優(yōu)化路由。數(shù)據(jù)流中的這種被動調整會增加延遲，并在數(shù)據(jù)流中引入非確定性。

為了避免這種情況，分布式多TSP系統(tǒng)使用編譯器顯式調度通過網(wǎng)絡的數(shù)據(jù)流。編譯器巧妙地路由數(shù)據(jù)，以便在任何時間點都不會在網(wǎng)絡中出現(xiàn)擁塞積聚。

除此之外，編譯器計劃的數(shù)據(jù)流還改善了網(wǎng)絡中的延遲，因為編譯器可以調度數(shù)據(jù)主動推送，而不是必須通過設備請求。

確定性負載均衡

在編譯時調度數(shù)據(jù)流的另一個優(yōu)點是，它允許編譯器有效地跨可用鏈接對流進行負載均衡。在傳統(tǒng)網(wǎng)絡中，硬件根據(jù)路由器中可用的擁塞指標，按數(shù)據(jù)包執(zhí)行路由決策。

但是，在多TSP系統(tǒng)的情況下，編譯器會根據(jù)數(shù)據(jù)量以最佳方式執(zhí)行調度，并選擇要分散流量的鏈路數(shù)量。這樣可以有效地利用系統(tǒng)中的可用帶寬，并減少整體延遲。