你有沒(méi)有想過(guò)一個(gè)看似簡(jiǎn)單的技術(shù)決策是如何在幾年后徹底改變一個(gè)行業(yè)格局的2015年當(dāng)谷歌的工程師們開(kāi)始為搜索推薦和AlphaGo等AI應(yīng)用的計(jì)算瓶頸而頭疼時(shí)他們面前擺著幾條路繼續(xù)堆疊更貴的GPU嘗試更靈活的FPGA或者走一條當(dāng)時(shí)看起來(lái)風(fēng)險(xiǎn)極高、幾乎無(wú)人涉足的路——從頭設(shè)計(jì)一款專(zhuān)為神經(jīng)網(wǎng)絡(luò)計(jì)算而生的芯片。今天當(dāng)我們談?wù)揂I算力時(shí)TPU張量處理單元已經(jīng)是一個(gè)繞不開(kāi)的名字。但回到故事的開(kāi)端這絕不是一個(gè)“因?yàn)橛绣X(qián)所以造芯片”的簡(jiǎn)單故事。它關(guān)乎一個(gè)核心判斷當(dāng)通用計(jì)算架構(gòu)成為瓶頸時(shí)真正的突破往往來(lái)自于為特定工作負(fù)載進(jìn)行“暴力簡(jiǎn)化”的專(zhuān)用設(shè)計(jì)。這篇文章我們不打算復(fù)述那些輝煌的戰(zhàn)績(jī)而是想拆解TPU誕生背后的工程邏輯與決策路徑看看一個(gè)頂級(jí)技術(shù)團(tuán)隊(duì)是如何從實(shí)際問(wèn)題出發(fā)完成一次從軟件到硬件的協(xié)同創(chuàng)新并最終沉淀出一套可復(fù)用的專(zhuān)用加速器設(shè)計(jì)方法論的。1. 問(wèn)題的起點(diǎn)為什么通用GPU也開(kāi)始不夠用了要理解TPU為什么會(huì)出現(xiàn)首先要回到2013-2015年的谷歌。那時(shí)深度學(xué)習(xí)已經(jīng)開(kāi)始在圖像識(shí)別、語(yǔ)音處理和搜索排名中展現(xiàn)出巨大潛力。工程師們很快發(fā)現(xiàn)基于CPU的傳統(tǒng)服務(wù)架構(gòu)在運(yùn)行這些模型時(shí)延遲和成本都高得難以接受。于是他們自然轉(zhuǎn)向了GPU。1.1 GPU的“甜蜜點(diǎn)”與“阿喀琉斯之踵”GPU確實(shí)是當(dāng)時(shí)的最佳選擇。其強(qiáng)大的并行浮點(diǎn)計(jì)算能力尤其是對(duì)矩陣乘法的加速讓模型訓(xùn)練和推理的速度提升了數(shù)十倍。在訓(xùn)練階段GPU幾乎是無(wú)可替代的。然而當(dāng)模型進(jìn)入生產(chǎn)環(huán)境進(jìn)行大規(guī)模在線(xiàn)推理Inference時(shí)問(wèn)題開(kāi)始浮現(xiàn)。推理任務(wù)與訓(xùn)練有本質(zhì)不同確定性高模型已經(jīng)固定計(jì)算圖是靜態(tài)的。延遲敏感用戶(hù)搜索時(shí)結(jié)果必須在幾十毫秒內(nèi)返回。吞吐量要求巨大全球每秒的搜索請(qǐng)求是天文數(shù)字。能效比至關(guān)重要數(shù)據(jù)中心的空間、電力和冷卻成本是運(yùn)營(yíng)的核心開(kāi)銷(xiāo)。GPU為通用并行計(jì)算設(shè)計(jì)其架構(gòu)在推理場(chǎng)景下顯得“過(guò)于復(fù)雜”緩存與控制邏輯冗余為了支持復(fù)雜的圖形渲染和通用計(jì)算GPU有大量用于分支預(yù)測(cè)、亂序執(zhí)行、復(fù)雜調(diào)度的硬件單元。對(duì)于高度規(guī)整的矩陣運(yùn)算這些單元大部分時(shí)間處于閑置狀態(tài)卻在持續(xù)消耗功耗。訪(fǎng)存瓶頸盡管有高帶寬顯存但神經(jīng)網(wǎng)絡(luò)模型參數(shù)巨大頻繁的數(shù)據(jù)搬運(yùn)仍然是性能瓶頸。GPU的通用內(nèi)存架構(gòu)并非為這種“權(quán)重固定、數(shù)據(jù)流動(dòng)”的模式做極致優(yōu)化。精度過(guò)剩許多推理任務(wù)并不需要GPU擅長(zhǎng)的FP32高精度INT8甚至更低精度就能滿(mǎn)足要求同時(shí)能大幅降低計(jì)算和存儲(chǔ)開(kāi)銷(xiāo)。這時(shí)谷歌的工程師們算了一筆賬如果繼續(xù)沿著“買(mǎi)更多、更快的GPU”這條路走為了滿(mǎn)足未來(lái)幾年AI服務(wù)需求的指數(shù)級(jí)增長(zhǎng)數(shù)據(jù)中心的建設(shè)和電力成本將變得不可持續(xù)。他們需要的不是一個(gè)更快的通用處理器而是一個(gè)為“神經(jīng)網(wǎng)絡(luò)推理”這個(gè)單一任務(wù)量身定制的、極度高效的“計(jì)算鍋爐”。1.2 十字路口的三條路徑CPU、FPGA與ASIC面對(duì)瓶頸技術(shù)團(tuán)隊(duì)通常會(huì)評(píng)估幾種方案方案核心思路優(yōu)勢(shì)劣勢(shì)適用于TPU誕生前的場(chǎng)景優(yōu)化軟件繼續(xù)用CPU改進(jìn)算法、模型壓縮、利用CPU新指令集無(wú)需新硬件利用現(xiàn)有基礎(chǔ)設(shè)施性能提升有天花板無(wú)法解決根本能效比問(wèn)題早期探索或?qū)ρ舆t不敏感的后臺(tái)任務(wù)采用FPGA用可編程門(mén)陣列實(shí)現(xiàn)定制化數(shù)據(jù)流靈活性高可快速迭代邏輯設(shè)計(jì)峰值性能、能效比通常低于專(zhuān)用芯片開(kāi)發(fā)難度大成本較高原型驗(yàn)證或算法尚未固化、需要頻繁變更的領(lǐng)域設(shè)計(jì)專(zhuān)用ASIC從頭設(shè)計(jì)一款只干神經(jīng)網(wǎng)絡(luò)推理的芯片性能、能效比可達(dá)到理論極限量產(chǎn)成本低研發(fā)周期長(zhǎng)通常2-3年投入巨大一旦設(shè)計(jì)完成幾乎無(wú)法修改需求明確、規(guī)模巨大、算法相對(duì)穩(wěn)定的場(chǎng)景谷歌的團(tuán)隊(duì)最終選擇了最艱難但最具潛力的ASIC之路。這個(gè)決策背后有幾個(gè)關(guān)鍵判斷需求足夠明確且龐大搜索引擎的神經(jīng)網(wǎng)絡(luò)推薦模型已經(jīng)是核心服務(wù)規(guī)模效應(yīng)足以攤平ASIC的巨額研發(fā)成本。算法相對(duì)穩(wěn)定盡管深度學(xué)習(xí)在發(fā)展但底層核心操作矩陣乘加、激活函數(shù)、池化是穩(wěn)定的。為這些操作設(shè)計(jì)硬件是安全的。時(shí)間窗口他們預(yù)見(jiàn)到AI推理需求即將爆發(fā)必須提前布局。等需求來(lái)了再行動(dòng)就晚了。這個(gè)選擇回答了“為什么是ASIC”的問(wèn)題但更大的挑戰(zhàn)在于“如何設(shè)計(jì)這個(gè)ASIC”。這不僅僅是硬件工程師的任務(wù)而是一場(chǎng)從軟件模型、編譯器到硬件架構(gòu)的深度協(xié)同。2. TPU v1的設(shè)計(jì)哲學(xué)為“推理”做極致的減法TPU的第一代設(shè)計(jì)目標(biāo)非常純粹在保證嚴(yán)格延遲限制的前提下最大化神經(jīng)網(wǎng)絡(luò)推理的吞吐量和能效比。它不是一顆通用的AI芯片而是一個(gè)針對(duì)已訓(xùn)練模型進(jìn)行前向傳播的專(zhuān)用加速器。2.1 架構(gòu)核心脈動(dòng)陣列與片上緩存TPU v1最標(biāo)志性的設(shè)計(jì)是其脈動(dòng)陣列。這是一個(gè)二維網(wǎng)格狀的處理單元陣列每個(gè)單元都能執(zhí)行一次乘加運(yùn)算。數(shù)據(jù)和權(quán)重從陣列的邊緣流入像波浪一樣在陣列中“脈動(dòng)”傳遞并在傳遞過(guò)程中完成計(jì)算。這種設(shè)計(jì)的精妙之處在于最大化數(shù)據(jù)復(fù)用每個(gè)權(quán)重?cái)?shù)據(jù)一旦被加載進(jìn)陣列會(huì)在多個(gè)計(jì)算周期內(nèi)與流經(jīng)的不同輸入數(shù)據(jù)相乘極大地減少了從外部?jī)?nèi)存讀取權(quán)重的次數(shù)。這是應(yīng)對(duì)訪(fǎng)存瓶頸的關(guān)鍵。簡(jiǎn)化控制流整個(gè)陣列由統(tǒng)一的時(shí)鐘驅(qū)動(dòng)數(shù)據(jù)流是規(guī)整同步的省去了通用處理器中復(fù)雜的指令分發(fā)、亂序執(zhí)行等控制邏輯功耗大幅降低。高計(jì)算密度在給定的芯片面積上脈動(dòng)陣列能部署遠(yuǎn)超通用核心數(shù)量的計(jì)算單元。與脈動(dòng)陣列配套的是巨大的片上統(tǒng)一緩沖區(qū)。這個(gè)緩沖區(qū)充當(dāng)了數(shù)據(jù)和權(quán)重的“中轉(zhuǎn)站”其容量經(jīng)過(guò)精心設(shè)計(jì)能夠容納神經(jīng)網(wǎng)絡(luò)關(guān)鍵層的全部參數(shù)使得計(jì)算核心可以長(zhǎng)時(shí)間、高速地從片上獲取數(shù)據(jù)而不是頻繁訪(fǎng)問(wèn)速度更慢、功耗更高的片外DRAM。注意這里體現(xiàn)了一個(gè)重要的硬件設(shè)計(jì)原則——“面向數(shù)據(jù)流設(shè)計(jì)”。TPU的硬件架構(gòu)本質(zhì)上是對(duì)神經(jīng)網(wǎng)絡(luò)計(jì)算數(shù)據(jù)流圖的直接硬件映射。計(jì)算在哪里發(fā)生數(shù)據(jù)如何流動(dòng)都在設(shè)計(jì)階段被固化下來(lái)以換取極致的效率。2.2 軟件棧的協(xié)同設(shè)計(jì)從TensorFlow到編譯器再優(yōu)秀的硬件如果沒(méi)有友好的軟件接口也無(wú)法被廣泛應(yīng)用。TPU的成功另一半功勞要?dú)w于其軟件棧的提前布局。谷歌在研發(fā)TPU硬件的同時(shí)也在內(nèi)部大力推廣TensorFlow。TensorFlow使用數(shù)據(jù)流圖來(lái)描述計(jì)算這與TPU的硬件數(shù)據(jù)流思想天然契合。更重要的是TensorFlow為T(mén)PU提供了一個(gè)完美的抽象層用戶(hù)層面研究人員和工程師用TensorFlow API定義模型無(wú)需關(guān)心底層是CPU、GPU還是TPU。編譯器層面XLA等編譯器將TensorFlow計(jì)算圖進(jìn)行優(yōu)化、分區(qū)并編譯成能在TPU上高效執(zhí)行的指令流。這種軟硬協(xié)同意味著降低使用門(mén)檻應(yīng)用團(tuán)隊(duì)不需要學(xué)習(xí)新的硬件編程模型。發(fā)揮硬件極限編譯器可以進(jìn)行圖層融合、內(nèi)存布局優(yōu)化等高級(jí)操作將硬件的潛力榨干?？焖俚布F(tuán)隊(duì)和軟件框架團(tuán)隊(duì)可以緊密合作根據(jù)實(shí)際模型的需求調(diào)整硬件微架構(gòu)或編譯器策略。2.3 與CPU/GPU的定位區(qū)分不是替代是補(bǔ)充一個(gè)常見(jiàn)的誤解是TPU旨在取代CPU和GPU。實(shí)際上谷歌從一開(kāi)始就明確了異構(gòu)計(jì)算的定位CPU負(fù)責(zé)通用邏輯、控制流、數(shù)據(jù)預(yù)處理和后處理、任務(wù)調(diào)度等。GPU負(fù)責(zé)模型訓(xùn)練和某些對(duì)靈活性要求較高的推理任務(wù)。TPU負(fù)責(zé)大規(guī)模、低延遲、高吞吐的定型模型推理。在服務(wù)器中TPU以PCIe加速卡的形式存在。主機(jī)CPU負(fù)責(zé)將模型加載到TPU的內(nèi)存中準(zhǔn)備好輸入數(shù)據(jù)然后啟動(dòng)TPU執(zhí)行。TPU完成計(jì)算后CPU再取回結(jié)果。這是一個(gè)典型的主從協(xié)作模式。3. 從v1到后續(xù)演進(jìn)專(zhuān)用化的邊界與擴(kuò)展TPU v1在推理任務(wù)上取得了巨大成功但其設(shè)計(jì)邊界也非常清晰僅支持推理且僅支持特定的數(shù)值精度和算子。這既是其成功的秘訣也帶來(lái)了挑戰(zhàn)。隨著AI的發(fā)展谷歌的TPU架構(gòu)也在不斷演進(jìn)。3.1 挑戰(zhàn)訓(xùn)練需求與模型復(fù)雜化模型訓(xùn)練對(duì)硬件的要求比推理更高需要反向傳播涉及梯度計(jì)算需要支持更復(fù)雜的計(jì)算圖。需要高數(shù)值精度訓(xùn)練過(guò)程對(duì)數(shù)值穩(wěn)定性敏感通常需要FP32甚至混合精度。需要更大的內(nèi)存不僅要存儲(chǔ)模型參數(shù)還要存儲(chǔ)中間激活值、梯度等。需要靈活性研究中的模型結(jié)構(gòu)變化快。為了支持訓(xùn)練后續(xù)的TPU版本如TPU v2/v3在架構(gòu)上進(jìn)行了重大升級(jí)支持浮點(diǎn)計(jì)算加入了FP32和BF16浮點(diǎn)計(jì)算單元。大幅增加HBM高帶寬內(nèi)存以容納大型模型和中間狀態(tài)。設(shè)計(jì)互聯(lián)網(wǎng)絡(luò)通過(guò)高速互聯(lián)將多個(gè)TPU芯片組成Pod支持大規(guī)模分布式訓(xùn)練。增強(qiáng)軟件棧編譯器需要支持自動(dòng)微分、分布式策略等復(fù)雜功能。3.2 核心矛盾的平衡效率 vs. 靈活性TPU的演進(jìn)史就是一部在“專(zhuān)用效率”和“通用靈活性”之間尋找新平衡點(diǎn)的歷史。v1極端偏向效率為推理定制。v2/v3/v4逐步增加靈活性以支持訓(xùn)練但通過(guò)保留脈動(dòng)陣列核心、定制互聯(lián)等方式在特定領(lǐng)域尤其是大規(guī)模矩陣運(yùn)算仍保持遠(yuǎn)超通用硬件的效率。軟件定義通過(guò)不斷強(qiáng)大的編譯器XLA和框架TensorFlow/JAX用軟件來(lái)彌補(bǔ)硬件靈活性的不足將用戶(hù)多變的計(jì)算圖“翻譯”成硬件高效執(zhí)行的模式。這個(gè)過(guò)程中一個(gè)關(guān)鍵的設(shè)計(jì)方法論得以鞏固“通過(guò)軟件定義硬件通過(guò)硬件加速軟件”。硬件提供一個(gè)高效但有一定約束的計(jì)算范式軟件則負(fù)責(zé)將廣闊的應(yīng)用需求優(yōu)雅地映射到這個(gè)范式之上。4. 給工程師的啟示從TPU案例中學(xué)到什么TPU的故事遠(yuǎn)不止于一款芯片的成功。它為我們提供了一個(gè)教科書(shū)級(jí)別的案例展示了如何從真實(shí)的業(yè)務(wù)痛點(diǎn)出發(fā)進(jìn)行跨棧的深度優(yōu)化。即使我們不設(shè)計(jì)芯片也能從中汲取對(duì)日常開(kāi)發(fā)極具價(jià)值的經(jīng)驗(yàn)。4.1 當(dāng)性能遇到瓶頸時(shí)向上堆資源不是唯一解我們習(xí)慣的做法是應(yīng)用慢了加機(jī)器數(shù)據(jù)庫(kù)慢了升級(jí)配置。這本質(zhì)是在通用解決方案上做線(xiàn)性投入。TPU的啟示在于要敢于審視整個(gè)工作負(fù)載的特性。如果你的業(yè)務(wù)中有某種計(jì)算模式反復(fù)出現(xiàn)、消耗了絕大部分資源并且模式相對(duì)穩(wěn)定那么為其設(shè)計(jì)一個(gè)“專(zhuān)用”的解決方案可能會(huì)帶來(lái)數(shù)量級(jí)的效率提升。這種“專(zhuān)用化”可以發(fā)生在不同層面算法層面用近似算法替代精確算法。數(shù)據(jù)結(jié)構(gòu)層面設(shè)計(jì)針對(duì)性的數(shù)據(jù)格式如稀疏矩陣存儲(chǔ)。服務(wù)層面將熱點(diǎn)功能拆分為獨(dú)立的微服務(wù)并針對(duì)性?xún)?yōu)化?；A(chǔ)設(shè)施層面使用特定的數(shù)據(jù)庫(kù)、緩存或計(jì)算引擎。關(guān)鍵在于識(shí)別出那個(gè)“穩(wěn)定且昂貴”的核心模式。4.2 軟硬協(xié)同設(shè)計(jì)是系統(tǒng)級(jí)優(yōu)化的終極形態(tài)大多數(shù)軟件工程師視硬件為黑盒硬件工程師則不太關(guān)心上層應(yīng)用的具體邏輯。TPU的成功打破了這堵墻。它告訴我們最大的優(yōu)化空間往往存在于棧與棧之間的接口處。在實(shí)際工作中這意味著理解下層原理開(kāi)發(fā)高性能應(yīng)用的工程師需要了解CPU緩存行、內(nèi)存帶寬、NUMA架構(gòu)、GPU SM等硬件知識(shí)。為上層設(shè)計(jì)接口設(shè)計(jì)中間件或基礎(chǔ)庫(kù)時(shí)要考慮如何讓上層應(yīng)用能更自然地表達(dá)其計(jì)算意圖從而方便底層優(yōu)化。數(shù)據(jù)流思維像TPU設(shè)計(jì)一樣梳理你系統(tǒng)中的關(guān)鍵數(shù)據(jù)流消除不必要的拷貝和轉(zhuǎn)換讓數(shù)據(jù)盡可能“流”起來(lái)而不是被反復(fù)“搬運(yùn)”。4.3 專(zhuān)用化之前必須明確邊界和代價(jià)TPU v1的成功建立在“只做推理”這個(gè)清晰的邊界之上。盲目追求通用性往往會(huì)犧牲掉專(zhuān)用化的極致優(yōu)勢(shì)。在決定為一個(gè)系統(tǒng)或模塊做深度定制優(yōu)化前必須問(wèn)自己幾個(gè)問(wèn)題需求是否穩(wěn)定如果業(yè)務(wù)邏輯每月一變專(zhuān)用化可能血本無(wú)歸。規(guī)模是否足夠大優(yōu)化帶來(lái)的收益能否覆蓋研發(fā)和維護(hù)的額外成本是否構(gòu)成了關(guān)鍵路徑上的主要瓶頸優(yōu)化一個(gè)只占1%時(shí)間的模塊意義不大。有沒(méi)有更輕量的替代方案例如能否先用FPGA或高度優(yōu)化的軟件庫(kù)如oneDNN、TensorRT驗(yàn)證效果4.4 迭代路徑從驗(yàn)證核心價(jià)值開(kāi)始谷歌沒(méi)有一開(kāi)始就設(shè)計(jì)能訓(xùn)練所有模型的萬(wàn)能TPU。他們的路徑非常清晰聚焦最痛點(diǎn)先解決規(guī)模最大、成本最高的在線(xiàn)推理問(wèn)題。打造最小可行產(chǎn)品TPU v1功能單一但在其目標(biāo)場(chǎng)景下效果驚人。建立生態(tài)通過(guò)TensorFlow綁定用戶(hù)形成軟硬一體化的體驗(yàn)。逐步擴(kuò)展邊界在獲得成功后再向訓(xùn)練等領(lǐng)域演進(jìn)并不斷升級(jí)硬件。這給我們工程實(shí)踐的啟示是面對(duì)一個(gè)復(fù)雜優(yōu)化問(wèn)題不要試圖一次性設(shè)計(jì)出完美方案。先找到那個(gè)能創(chuàng)造最大價(jià)值的核心點(diǎn)用最小的代價(jià)實(shí)現(xiàn)它、驗(yàn)證它。獲得正反饋后再?lài)@這個(gè)核心逐步構(gòu)建更強(qiáng)大的系統(tǒng)?；剡^(guò)頭看TPU的誕生不是一個(gè)偶然的科技奇跡而是一次基于深刻業(yè)務(wù)洞察、嚴(yán)謹(jǐn)工程權(quán)衡和長(zhǎng)期主義投入的系統(tǒng)性創(chuàng)新。它從谷歌數(shù)據(jù)中心里一個(gè)具體的算力與成本問(wèn)題出發(fā)最終通過(guò)軟硬協(xié)同的深度設(shè)計(jì)不僅解決了自身的問(wèn)題更推動(dòng)了整個(gè)AI基礎(chǔ)設(shè)施的發(fā)展。對(duì)于我們而言其價(jià)值不在于是否能用上TPU芯片而在于它展示了一種解決問(wèn)題的方法論當(dāng)你在通用道路上遇到難以逾越的墻時(shí)不妨停下來(lái)看看你的負(fù)載是否特別到值得為它修一條專(zhuān)屬的高速公路。這條路的開(kāi)頭可能很難但一旦走通風(fēng)景將截然不同。