上月底，創業公司Groq的產品一夜爆火。

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

圖片

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

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

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

Groq LPU的神秘面紗

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

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

圖片

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

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

這兩篇文章解釋了Groq的張量流處理器（TSP）的設計和實現，以及他們如何使用TSP構建分布式推理引擎。

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

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

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

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

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

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

舉個例子，現代CPU一般具有如下設計：

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

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

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

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

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

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

而GPU還有其他一些非確定性因素，包括緩存、共享和全局內存、動態資源分區等。

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

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

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

TSP架構

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

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

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

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

TSP具有以下功能切片：

MXM：用于執行矩陣運算 

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

MEM：內存讀/寫運算 

VXM：向量上的算術運算 

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

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

TSP中的指令執行

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

這些切片以生產者——消費者的方式進行交互。

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

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

圖片

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

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

TSP的編譯器和ISA

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

編譯器有權訪問TSP硬件的以下狀態：

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

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

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

- 220M全局共享SRAM。

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

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

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

硬件中的確定性數據路徑 

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

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

從TSP擴展到LPU

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

設計多TSP系統

與TSP一樣，分布式多TSP系統的設計目標也圍繞著確定性數據流和指令執行，以及節點之間的低延遲通信。

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

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

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

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

其中144個端口在機架內本地使用，以便在機架內快速傳輸數據，其余144個端口用于連接到其他機架。

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

在基于TSP的分布式系統中實現確定性

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

使用硬件對齊計數器同步TSP的時鐘

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

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

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

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

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

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

初始計劃調整

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

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

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

- 要擴展多跳系統，可以在生成樹的每個hop上重復執行以上方案。

運行時重新同步

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

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

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

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

編譯器在軟件計劃網絡中的作用

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

編譯器不是動態管理數據流，而是在編譯時靜態解析所有內容。

已知流量模式

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

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

計劃的數據流

在傳統的網絡系統中，通過網絡的數據包流由硬件管理，硬件在感應到網絡中的負載時會優化路由。數據流中的這種被動調整會增加延遲，并在數據流中引入非確定性。

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

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

確定性負載均衡

在編譯時調度數據流的另一個優點是，它允許編譯器有效地跨可用鏈接對流進行負載均衡。在傳統網絡中，硬件根據路由器中可用的擁塞指標，按數據包執行路由決策。

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