谷歌TPU跑Kimi比英偉達GPU快57%！用的還是DeepSeek推理框架

n，比老黃的GB200快了57%。做出這個成績的，既不是造Kimi的月之暗面，也不是造TPU的谷歌，是一家叫Inferact的推理創業公司。Inferact創始團隊就是vLLM的原班人馬，今年拿了a16z領投的1.5億美元種子輪，估值8億美元。他們給TPU寫了一種叫megakernel的推理內核，把模型推理時原本要調度的幾百個小程序焊成一個大程序。配合TPU獨特的片上內存架構，megakernel把內存帶寬利用率逼到了硬件理論峰值附近。配上DeepSeek提出的DSpark加速框架，K3在TPU上跑出了前面提到的每秒709 token的成績。這套代碼，現在已經開源。

同樣16塊芯片，TPU快57% Inferact把TPU和英偉達GPU放在完全相同的條件下跑了一輪benchmark，TPU贏了。測試是用16塊TPU v7 Ironwood對陣16塊英偉達GB200，兩邊都跑Kimi K3，都用vLLM推理引擎，唯一的變量是芯片和底層的kernel實現。最終，TPU跑出的成績是每秒709個token，GB200跑出每秒452個，TPU的速度快了57%。這個速度裡有DSpark推測解碼的功勞。

DSpark是由DeepSeek提出的推理加速框架，其原理是讓一個小模型先快速猜出一串候選token，然後大模型再對這串候選token做批量驗證，猜對的直接跳過，猜錯的回退重來。Inferact在TPU上跑通了這套流程，acceptance length達到了6，也就是每輪猜出的token裡，平均有6個能直接通過大模型的驗證，單步decode的耗時大約在8.5毫秒。關掉推測解碼看裸速，差距同樣拉得開。在batch size為1的情況下，TPU每秒能跑249個token，GB200只有127個，差距接近一倍。

把batch size放大到8，TPU達到865個token每秒，GB200只有636個。另外在Qwen上，兩種芯片的差距更大。Inferact用4塊TPU v7跑Qwen 3.8 27B，每秒跑出1515個token，同樣配置的GB200只跑出了695個。而且這種提速並沒有造成精度損失，Inferact用greedy decoding做了驗證，Kimi K3在TPU上的GPQA-Diamond得分是94.4%，GSM8K是97.2%，和GPU上的結果完全一致。推理引擎一樣，模型一樣，芯片換一下，速度就差出了57%。這種程度的差距，按理說應該能從硬件規格表上找到解釋，但實際情況恰好相反。

兩款芯片的算力也在同一個量級，TPU v7的HBM帶寬是7380 GB/s，而GB200的HBM帶寬反而更高，達到8000 GB/s，帶寬更大的那塊芯片，跑出來的速度反而更低。這種速度的差距，實際上來源於芯片上面跑的那層軟件。把幾百個小kernel，焊成一個大程序 Inferact用一個叫megakernel的方案解決了TPU上的推理效率問題，核心思路，是把模型推理時原本要調度的幾百個獨立小程序，合併成一個大程序，從而消除程序之間切換時的帶寬浪費。大模型推理的速度瓶頸不在計算，在數據搬運。

每生成一個token，模型的全部權重都要從HBM主存搬進芯片內部的高速緩存裡算一遍，算完這一輪，下一個token再重新搬一遍。傳統的做法是把一次推理拆成幾百個獨立的kernel，每個kernel負責一小塊計算。這些kernel排著隊一個接一個地執行，但問題出依然會在交接的間隙。上一個kernel幹完了，下一個還沒有啟動起來，這段時間裡內存帶寬就這麼空轉著。CUDA Graphs和英偉達最新的PDL技術能縮短這個間隙，但kernel之間的邊界還在。Inferact的megakernel直接消除了這個邊界。

Kimi K3一共有92層MoE計算，Inferact把這92層的計算邏輯從頭到尾塞進了一個Pallas程序裡，一次調用就走完整個模型的前向傳播。邊界消失之後，跨層權重預取就變得順理成章了。第N層還在算MoE的時候，第N+1層的attention權重已經開始從HBM往片上緩存搬了。計算和數據搬運同時進行，內存帶寬幾乎沒有空轉的時刻。TPU的硬件特性讓這種編排做起來格外順暢。TPU v7的每個TensorCore帶有64 MiB的VMEM片上內存，這塊內存的生命週期完全交給軟件來管理，工程師可以精確控制什麼時候往裡面裝什麼數據、什麼時候把空間騰出來。

64 MiB的容量足夠同時裝下當前層的計算數據和下一層預取的權重。GPU那邊情況不同，英偉達Blackwell架構的片上內存分散在152個SM裡面，總量大約38 MiB，由硬件自動調度，要做類似的跨層編排限制更多。Inferact用Pallas語言手寫了整個megakernel的代碼，Pallas是谷歌給TPU做的底層編程語言，角色類似英偉達GPU上的CUDA。TPU默認的編譯工具鏈XLA擅長優化單層計算，但跨92層協調數據搬運的決策分支太多，XLA的自動調度找不到最優解。

Inferact的做法是繞過XLA，用Pallas語言手寫整個megakernel的代碼，由工程師自己來決定每一步搬什麼數據、存在哪塊緩存、什麼時候釋放。這樣做還帶來了一個額外好處，那就是編譯速度大幅提升，耗時從XLA的30分鐘以上降到了不到90秒，工程師改完一版kernel一分半鐘就能上機驗證。目前這套megakernel是針對Kimi K3的模型結構做的定製優化，如果換一個模型架構還需要重新適配。Inferact在博客裡提到，團隊接下來會把megakernel的支持擴展到更多模型架構上。

vLLM原班人馬的創業公司 Inferact去年11月成立，團隊幾乎就是vLLM的原始開發者，一群靠給英偉達GPU寫推理引擎成名的人。他們今年一月出來創業，轉頭把TPU的推理速度做到了GPU前面。vLLM是目前開源推理引擎領域的重量級框架，支持超過500種模型架構，社區貢獻者超過2000人，Meta、Google、Character.ai都拿它來做線上的推理服務。CEO Simon Mo是vLLM項目的原始維護者，伯克利EECS畢業，之前在Anyscale工作。聯合創始人Woosuk Kwon是整個vLLM項目的發起人，伯克利計算機科學博士，導師是Ion Stoica。

Kwon提出的PagedAttention算法解決了GPU顯存管理中的碎片化問題，這個算法至今仍然是vLLM的核心技術。首席科學家遊凱超曾獲清華大學特等獎學金，他在vLLM中主導了分佈式推理功能的開發。Inferact今年完成了1.5億美元的種子輪融資，公司估值8億美元。a16z和Lightspeed領投，真格、紅杉、Altimeter跟投。這次的TPU megakernel來自Inferact與Google Cloud的一項聯合工程合作。雙方的目標是讓TPU成為vLLM的一流支持對象，所有的優化成果都會回饋給上游的開源社區，tpu-megakernels代碼倉庫是這次合作的第一個公開成果。

參考鏈接： 700 TPS on Kimi K3: A Case for TPU Megakernels 版權所有，未經授權不得以任何形式轉載及使用，違者必究。