脈動陣列

Systolic Array

TPU核心的矩陣運算架構

詳細解釋

脈動陣列(Systolic Array)是專為矩陣計算優化的硬體架構,數據像血液在心脏脈動一樣流經處理單元陣列,實現高效的並行計算。

概念起源:

  • 類比心臟脈動:數據流動的節奏性
  • 1970年代提出,用於信號處理
  • 深度學習復興:矩陣乘法需求激增

架構特點:

  • 規則陣列:2D網格排列的處理單元(PE)
  • 數據流動:數據從邊緣流入,流經PE陣列
  • 本地計算:每個PE簡單計算,結果傳遞給鄰居
  • 最小存儲:數據流動中完成計算,減少記憶體訪問

矩陣乘法實現:

  • 權重矩陣:預加載到PE陣列
  • 激活向量:從左側流入
  • 部分和:從頂部累積到底部
  • 結果:從底部或右側輸出

優勢:

  • 高並行性:多個PE同時工作
  • 數據重用:數據流經多個PE被重用
  • 能效:減少數據移動,提高能效
  • 簡單控制:規則數據流,控制邏輯簡單

應用:

  • Google TPU:核心計算單元是脈動陣列
  • NVIDIA Tensor Core:類似概念優化矩陣乘
  • 專用AI芯片:幾乎所有AI加速器的核心
  • CNN加速器:卷積的矩陣乘法映射

與其他架構的對比:

  • CPU:通用但效率低
  • GPU:靈活並行,但控制和記憶體開銷大
  • 脈動陣列:專用高效,但僅適合規則計算

設計考量:

  • 陣列大小:權衡面積和效率
  • 數據精度:INT8/FP16/BF16選擇
  • 帶寬匹配:IO帶寬與計算速率匹配
  • 靈活性:支援不同矩陣大小和形狀

現代變體:

  • 2D脈動陣列:傳統設計
  • 3D堆疊:與HBM整合
  • 稀疏支持:處理稀疏矩陣

脈動陣列是AI加速器的核心計算架構。

探索更多AI詞彙

查看所有分類,繼續學習AI知識