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