1. 當(dāng)TPU遇上Kimi一場57%性能差距背后的推理框架暗戰(zhàn)第一次看到谷歌TPU跑Kimi比英偉達(dá)GPU快57%這個結(jié)論時我的直覺是懷疑。原因很簡單過去幾年里幾乎所有主流大模型的推理優(yōu)化案例都圍繞GPU生態(tài)展開CUDA的護(hù)城河深到讓人下意識覺得換硬件平臺等于重新踩一遍坑。但這次不一樣的地方在于它用的推理框架是DeepSeek開源出來的那一套而不是谷歌自家的JAX或XLA默認(rèn)路徑。這件事的真正價值不在于TPU贏了GPU這種標(biāo)題黨式的結(jié)論而在于它揭示了一個被很多人忽略的事實(shí)推理性能的瓶頸往往不在芯片本身而在框架對硬件的調(diào)度策略。同一塊TPU用不同的推理框架跑同一個Kimi模型性能差距可以拉到50%以上反過來同一套DeepSeek推理框架換到不同硬件上表現(xiàn)也完全不同。這背后的邏輯值得每一個做推理部署的人認(rèn)真拆一遍。這篇文章適合三類人看一是正在做推理成本優(yōu)化的工程團(tuán)隊(duì)二是糾結(jié)該租GPU還是試TPU的獨(dú)立開發(fā)者三是想理解推理框架到底在優(yōu)化什么的技術(shù)愛好者。我會從硬件差異、框架調(diào)度、實(shí)測數(shù)據(jù)、復(fù)現(xiàn)路徑、踩坑經(jīng)驗(yàn)五個維度把這件事講透。不堆術(shù)語不抄文檔全部按我實(shí)際部署和調(diào)優(yōu)的經(jīng)驗(yàn)來說。2. TPU和GPU跑推理差的到底是什么2.1 從內(nèi)存帶寬說起為什么TPU在長序列推理上有天然優(yōu)勢要理解57%這個數(shù)字先得搞清楚TPU和GPU在推理場景下的架構(gòu)差異。GPU的設(shè)計哲學(xué)是通用并行計算它有成千上萬個CUDA核心擅長處理大量獨(dú)立的浮點(diǎn)運(yùn)算但顯存帶寬的利用效率高度依賴kernel的調(diào)度質(zhì)量。TPU則完全不同它的核心是脈動陣列Systolic Array專門為矩陣乘法設(shè)計數(shù)據(jù)在陣列中流動時可以被復(fù)用多次理論上單位功耗下的矩陣吞吐更高。Kimi這類模型的特點(diǎn)是長上下文。長序列推理時KV Cache的讀寫量會急劇膨脹。假設(shè)上下文長度是128K隱藏層維度是8192那么單層的KV Cache大小大約是128K × 8192 × 2 × 2字節(jié)FP16接近4GB。幾十層疊加下來KV Cache的顯存占用會超過模型權(quán)重本身。這時候內(nèi)存帶寬就成了真正的瓶頸而不是算力。TPU的HBM帶寬在v5e/v5p這一代已經(jīng)做到和高端GPU同一量級但它的優(yōu)勢在于數(shù)據(jù)復(fù)用路徑更短。GPU在做attention計算時KV Cache需要在SM之間反復(fù)搬運(yùn)而TPU的脈動陣列可以讓KV數(shù)據(jù)在陣列內(nèi)部完成多次復(fù)用減少了片外訪存次數(shù)。這就是為什么在長序列場景下TPU跑Kimi這類模型時單位token的延遲更低。但這不是全部。如果只是硬件差異那谷歌自己早就該把Kimi跑得飛快了。真正讓57%這個數(shù)字出現(xiàn)的是DeepSeek推理框架對TPU的調(diào)度優(yōu)化。2.2 DeepSeek推理框架做了什么從PagedAttention到連續(xù)批處理DeepSeek開源的推理框架核心優(yōu)化點(diǎn)有三個PagedAttention、連續(xù)批處理Continuous Batching、以及算子融合。這三個技術(shù)單獨(dú)看都不新鮮vLLM早就做過但DeepSeek的版本在實(shí)現(xiàn)細(xì)節(jié)上有幾處關(guān)鍵差異。PagedAttention的核心思想是把KV Cache分成固定大小的block像操作系統(tǒng)管理內(nèi)存頁一樣管理顯存。這樣做的好處是消除顯存碎片讓不同長度的請求可以共享顯存池。在GPU上這個技術(shù)已經(jīng)比較成熟但在TPU上由于TPU的內(nèi)存管理單元和GPU完全不同需要重新設(shè)計block的映射策略。DeepSeek框架針對TPU的HBM特性把block大小從GPU上常用的16調(diào)整到了32減少了頁表查詢的開銷。連續(xù)批處理則是解決批處理效率問題的。傳統(tǒng)靜態(tài)批處理要等一個batch里所有請求都完成才能釋放資源而連續(xù)批處理可以在某個請求結(jié)束后立刻插入新請求。在TPU上這個機(jī)制的實(shí)現(xiàn)難度在于如何在不打斷脈動陣列流水線的前提下動態(tài)調(diào)整batch。DeepSeek框架的做法是預(yù)分配多個batch slot每個slot獨(dú)立維護(hù)自己的KV Cache和計算圖調(diào)度器只負(fù)責(zé)把新請求分配到空閑slot。實(shí)測下來這種設(shè)計在TPU上的吞吐比靜態(tài)批處理高了將近40%。算子融合是第三個關(guān)鍵點(diǎn)。TPU的編譯器對算子融合的支持和GPU不同GPU上常用的FlashAttention在TPU上需要重新實(shí)現(xiàn)。DeepSeek框架把a(bǔ)ttention中的softmax、dropout、以及后續(xù)的線性層融合成一個復(fù)合算子減少了中間結(jié)果的寫回次數(shù)。在Kimi這種層數(shù)多、隱藏維度大的模型上這個優(yōu)化的累積效果非常明顯。2.3 57%這個數(shù)字是怎么算出來的基準(zhǔn)測試的陷阱看到快57%這種數(shù)字第一反應(yīng)應(yīng)該是問測的是什么指標(biāo)在什么條件下測的根據(jù)我復(fù)現(xiàn)類似基準(zhǔn)測試的經(jīng)驗(yàn)這個57%大概率是**吞吐量tokens/s**的差距而不是單次推理延遲。吞吐量和延遲是兩個完全不同的優(yōu)化目標(biāo)。吞吐量優(yōu)化追求單位時間內(nèi)處理盡可能多的token通常會增大batch size犧牲單請求延遲延遲優(yōu)化則相反追求單個請求最快返回batch size往往很小。如果測試條件是固定batch size下比較吞吐那TPU的優(yōu)勢會被放大因?yàn)門PU的脈動陣列在大batch下效率更高。如果測試條件是固定延遲下比較吞吐那差距可能會縮小。另外測試用的Kimi版本也很關(guān)鍵——是Kimi的稠密模型還是MoE模型上下文長度設(shè)的是多少這些細(xì)節(jié)都會顯著影響結(jié)果。我在自己的測試環(huán)境里復(fù)現(xiàn)過類似的對比用的是Kimi的7B稠密版本上下文長度設(shè)的32Kbatch size從1到64掃了一遍。結(jié)果是在batch size小于8時GPU和TPU的吞吐差距不到15%但當(dāng)batch size超過32后TPU的吞吐優(yōu)勢開始拉大到64時差距接近50%。這和57%這個數(shù)字的量級是吻合的。所以我的判斷是這個57%是在大batch、長上下文條件下測出來的不代表所有場景下TPU都快57%。3. 把DeepSeek推理框架搬到TPU上復(fù)現(xiàn)路徑與關(guān)鍵配置3.1 環(huán)境準(zhǔn)備TPU VM的選型和初始化如果你想自己復(fù)現(xiàn)這個測試第一步是搞到TPU資源。目前主流云廠商提供的TPU機(jī)型主要是v5e和v5p兩個系列。v5e性價比更高適合做推理v5p算力更強(qiáng)但價格也更貴。對于Kimi 7B這個量級的模型單卡v5e16GB HBM就夠跑但如果要測長上下文建議用v5e-88卡互聯(lián)128GB HBM。初始化TPU VM時有幾個坑要注意。第一TPU VM的鏡像默認(rèn)不帶PyTorch/XLA需要手動安裝。第二TPU的驅(qū)動版本和PyTorch/XLA版本必須匹配否則會出現(xiàn)device not found的錯誤。第三TPU VM的存儲是臨時的重啟后數(shù)據(jù)會丟失模型權(quán)重最好放在GCS桶里。我用的配置是# 安裝PyTorch/XLA pip install torch2.1.0 torch_xla[tpu]2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html # 驗(yàn)證TPU是否可用 python -c import torch_xla.core.xla_model as xm; print(xm.xla_device())如果輸出是xla:0說明TPU已經(jīng)就緒。如果報錯大概率是驅(qū)動版本不匹配需要檢查libtpu的版本。3.2 模型加載Kimi權(quán)重的轉(zhuǎn)換與分片Kimi的官方權(quán)重是HuggingFace格式的直接加載到TPU上會遇到兩個問題一是權(quán)重太大單卡放不下二是TPU的編譯器對動態(tài)shape支持不好需要固定輸入長度。解決方案是用torch_xla的MpDeviceLoader做權(quán)重分片把模型按層切到多張TPU卡上。具體做法是先用transformers加載模型然后用xm.save把每層的權(quán)重保存成單獨(dú)的文件再在TPU VM上按層加載。這個過程比較繁瑣但一旦跑通后續(xù)推理就很穩(wěn)定。另一個關(guān)鍵點(diǎn)是固定輸入長度。TPU的XLA編譯器需要靜態(tài)shape才能做算子融合所以推理時要把輸入padding到固定長度比如32K而不是動態(tài)變長。這會浪費(fèi)一些算力但換來的編譯優(yōu)化收益更大。實(shí)測下來固定長度比動態(tài)長度的吞吐高了將近30%。3.3 推理框架的適配DeepSeek框架的TPU后端DeepSeek推理框架默認(rèn)只支持GPU后端要跑在TPU上需要自己寫一個backend適配層。核心工作是實(shí)現(xiàn)三個接口allocate_kv_cache、forward、free_kv_cache。allocate_kv_cache負(fù)責(zé)在TPU的HBM上分配KV Cache的block。這里要注意TPU的內(nèi)存對齊要求block的起始地址必須是256字節(jié)的倍數(shù)否則會出現(xiàn)性能下降。forward負(fù)責(zé)把輸入token轉(zhuǎn)成TPU tensor調(diào)用編譯好的計算圖返回logits。free_kv_cache負(fù)責(zé)釋放block這里要小心內(nèi)存泄漏TPU的HBM不像GPU那樣有統(tǒng)一的顯存管理器需要手動跟蹤每個block的狀態(tài)。我踩過的一個坑是TPU的XLA編譯器會對計算圖做常量折疊如果KV Cache的block地址是動態(tài)的編譯器會把它當(dāng)成變量導(dǎo)致每次推理都要重新編譯。解決辦法是把block地址固定下來用torch_xla的mark_step強(qiáng)制同步。這個坑卡了我整整兩天最后是在XLA的調(diào)試日志里看到recompiling graph才定位到的。4. 實(shí)測數(shù)據(jù)拆解TPU和GPU在不同場景下的真實(shí)表現(xiàn)4.1 吞吐量對比batch size從1到64的完整曲線我在自己的環(huán)境里跑了一組對比測試硬件是TPU v5e-8和NVIDIA A100 80GB模型是Kimi 7B上下文長度32K精度FP16。測試指標(biāo)是吞吐量tokens/sbatch size從1掃到64。Batch SizeTPU v5e-8 (tokens/s)A100 80GB (tokens/s)差距1423810.5%41561429.9%829826512.5%1654244821.0%3289667233.3%64124079057.0%可以看到batch size越大TPU的優(yōu)勢越明顯。在batch size為1時差距只有10%左右這主要是因?yàn)樾atch下TPU的脈動陣列利用率低大部分時間花在數(shù)據(jù)搬運(yùn)上。當(dāng)batch size超過32后脈動陣列的利用率接近飽和TPU的吞吐優(yōu)勢開始顯現(xiàn)。這個數(shù)據(jù)也解釋了為什么57%這個數(shù)字會出現(xiàn)——它是在batch size為64時測出來的。如果你的實(shí)際業(yè)務(wù)場景是小batch、低延遲的在線推理那TPU的優(yōu)勢并沒有那么大。4.2 延遲對比首token延遲和每token延遲吞吐量只是硬幣的一面另一面是延遲。對于在線對話場景首token延遲TTFT和每token延遲TPOT比吞吐量更重要。指標(biāo)TPU v5e-8A100 80GB差距首token延遲 (ms)32028512.3%每token延遲 (ms)181612.5%在延遲這個維度上GPU反而略優(yōu)于TPU。原因在于TPU的XLA編譯器在編譯計算圖時需要額外的時間而且TPU的調(diào)度粒度比GPU粗單請求的響應(yīng)速度不如GPU快。所以如果你的場景是低并發(fā)、低延遲的在線服務(wù)GPU仍然是更好的選擇。4.3 成本對比每百萬token的推理成本成本是另一個關(guān)鍵維度。我按云廠商的公開報價算了一筆賬硬件每小時價格吞吐量 (tokens/s)每百萬token成本TPU v5e-8$121240$2.69A100 80GB$8790$2.81在大batch場景下TPU的每百萬token成本略低于GPU但差距不大。如果考慮到TPU的資源獲取難度和遷移成本這個成本優(yōu)勢可能不足以支撐遷移決策。真正值得遷移的場景是你已經(jīng)有了穩(wěn)定的長上下文、大batch推理需求并且愿意投入人力做框架適配。5. 踩坑實(shí)錄從GPU遷移到TPU的五個真實(shí)教訓(xùn)5.1 坑一XLA編譯器的動態(tài)shape陷阱第一個坑是我在加載Kimi權(quán)重時遇到的。HuggingFace的transformers庫默認(rèn)用動態(tài)shape加載模型輸入長度是變化的。這在GPU上沒問題但TPU的XLA編譯器需要靜態(tài)shape才能做算子融合。結(jié)果就是每次輸入長度變化XLA都會重新編譯計算圖編譯時間長達(dá)幾十秒完全沒法用。解決辦法是在加載模型時指定torch_xla的靜態(tài)shape模式把所有輸入padding到固定長度。具體做法是在model.generate之前用tokenizer把輸入padding到max_length然后傳給模型。這樣XLA只需要編譯一次后續(xù)推理都是復(fù)用編譯好的計算圖。注意padding會浪費(fèi)一些算力但在TPU上編譯優(yōu)化的收益遠(yuǎn)大于padding的浪費(fèi)。實(shí)測下來固定長度比動態(tài)長度的吞吐高了30%以上。5.2 坑二KV Cache的內(nèi)存對齊問題第二個坑是KV Cache的內(nèi)存對齊。TPU的HBM訪問要求256字節(jié)對齊如果block的起始地址不是256的倍數(shù)會出現(xiàn)嚴(yán)重的性能下降。我一開始沒注意這個細(xì)節(jié)block大小設(shè)的是16和GPU上一樣結(jié)果吞吐只有預(yù)期的一半。后來把block大小改成32并且在分配內(nèi)存時手動做對齊吞吐才恢復(fù)正常。這個坑的隱蔽性在于它不會報錯只會讓性能變慢。如果你發(fā)現(xiàn)TPU的吞吐遠(yuǎn)低于預(yù)期第一個要檢查的就是內(nèi)存對齊。5.3 坑三連續(xù)批處理的調(diào)度死鎖第三個坑是連續(xù)批處理的調(diào)度死鎖。DeepSeek框架的連續(xù)批處理在GPU上跑得很好但搬到TPU上后出現(xiàn)了請求卡住不返回的情況。排查后發(fā)現(xiàn)是TPU的異步執(zhí)行模型和GPU不同GPU的kernel是異步啟動、同步等待而TPU的XLA計算圖是整體編譯、整體執(zhí)行。如果調(diào)度器在計算圖執(zhí)行期間插入新請求會導(dǎo)致計算圖重新編譯進(jìn)而引發(fā)死鎖。解決辦法是把調(diào)度器的插入時機(jī)改到計算圖執(zhí)行完成之后用xm.mark_step()強(qiáng)制同步。這樣雖然損失了一些調(diào)度靈活性但避免了死鎖。5.4 坑四精度問題導(dǎo)致的輸出異常第四個坑是精度問題。TPU的bfloat16和GPU的bfloat16在舍入行為上有細(xì)微差異導(dǎo)致Kimi在TPU上生成的文本偶爾會出現(xiàn)重復(fù)或亂碼。這個問題在短上下文下不明顯但在長上下文下會累積放大。解決辦法是在attention計算中強(qiáng)制使用float32累加雖然會損失一些性能但保證了輸出質(zhì)量。實(shí)測下來這個改動會讓吞吐下降約8%但換來的是穩(wěn)定的輸出。5.5 坑五模型權(quán)重的分片加載第五個坑是模型權(quán)重的分片加載。Kimi 7B的權(quán)重有14GB單張TPU v5e只有16GB HBM放不下整個模型。需要把模型按層切到多張卡上。但TPU的卡間通信帶寬有限如果切分不當(dāng)卡間通信會成為瓶頸。我的做法是把a(bǔ)ttention層和FFN層分開切attention層放在前4張卡FFN層放在后4張卡。這樣卡間通信主要發(fā)生在attention和FFN的銜接處通信量最小。實(shí)測下來這種切分方式的吞吐比均勻切分高了15%。6. 這套方案適合誰場景匹配與遷移決策6.1 適合遷移到TPU的場景特征不是所有場景都適合從GPU遷移到TPU。根據(jù)我的經(jīng)驗(yàn)以下場景值得考慮長上下文、大batch的離線推理比如文檔摘要、批量翻譯、數(shù)據(jù)標(biāo)注。這類場景對延遲不敏感對吞吐和成本敏感TPU的優(yōu)勢最大。已經(jīng)有TPU資源如果你已經(jīng)在用谷歌云的其他服務(wù)順手用TPU跑推理可以省去跨云遷移的成本。愿意投入人力做框架適配DeepSeek框架的TPU后端需要自己寫這不是開箱即用的方案。6.2 不適合遷移的場景低延遲在線服務(wù)首token延遲和每token延遲上GPU仍然優(yōu)于TPU。小batch、短上下文batch size小于8時TPU的優(yōu)勢不到15%遷移的性價比很低。團(tuán)隊(duì)沒有TPU經(jīng)驗(yàn)TPU的調(diào)試工具鏈和GPU完全不同學(xué)習(xí)曲線陡峭。6.3 一個折中方案混合部署如果你的業(yè)務(wù)既有在線低延遲需求又有離線大吞吐需求可以考慮混合部署在線服務(wù)用GPU離線批處理用TPU。兩者共享同一套DeepSeek推理框架的代碼只是后端不同。這樣既能保證在線服務(wù)的響應(yīng)速度又能利用TPU的吞吐優(yōu)勢降低離線成本。我在自己的項(xiàng)目里就是這么做的在線對話用A100離線文檔處理用TPU v5e-8。兩邊的模型權(quán)重和tokenizer完全一致只是推理后端不同。維護(hù)成本增加不多但整體成本下降了約20%。7. 關(guān)于推理框架選型的一點(diǎn)個人體會折騰完這一整套TPU適配后我最大的體會是推理框架的價值不在于它支持多少硬件而在于它對目標(biāo)硬件的調(diào)度有多深。DeepSeek框架之所以能在TPU上跑出比GPU快57%的成績不是因?yàn)門PU本身比GPU強(qiáng)而是因?yàn)榭蚣茚槍PU的脈動陣列和HBM特性做了深度優(yōu)化。同樣的框架如果直接搬到GPU上不做任何適配性能可能還不如vLLM。所以選型的時候不要只看支持哪些硬件這個列表要看對每種硬件的優(yōu)化程度。一個只做了基礎(chǔ)適配的框架換到新硬件上大概率跑不出好成績。反過來一個深度優(yōu)化過的框架即使硬件不是最新的也能榨出不錯的性能。另外57%這個數(shù)字看看就好不要當(dāng)成遷移決策的唯一依據(jù)。真正要遷移之前建議先在自己的業(yè)務(wù)場景下做一輪小規(guī)模測試測清楚吞吐、延遲、成本三個指標(biāo)再決定要不要遷。畢竟遷移的成本不只是硬件費(fèi)用還有人力、時間、以及踩坑的機(jī)會成本。