揭開Groq LPU神秘面紗：世界最快硬體加速器的底層架構設計！

【導讀】 Groq一夜爆火的背後，是自研的語言處理單元硬體LPU，近日，Substack的專欄作家Abhinav Upadhyay為我們一步步揭示了LPU底層架構的奧秘

上月底，新創公司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系統的情況下，編譯器會根據資料量以最佳方式執行調度，並選擇要分散流量的連結數量。這樣可以有效地利用系統中的可用頻寬，並減少整體延遲。（新智元）

參考資料：

https://codeconfessions.substack.com/p/groq-lpu-design