@googledevs：Google 與 UCSD 合作透過 DFlash 實現 TPU 上 LLM 推論 3.13 倍加速。 Google Cloud 與 UCSD 研究團隊攜…

Google 與 UCSD 合作透過 DFlash 實現 TPU 上 LLM 推論 3.13 倍加速。 Google Cloud 與 UCSD 研究團隊攜手，將「DFlash」（擴散式推測解碼）整合至 vLLM TPU 推論框架，打破傳統自迴歸推測解碼的序列瓶頸，平均達成 3.13 倍 token 每秒加速，峰值近 6 倍，特別在數學任務表現卓越。 打破自迴歸瓶頸 傳統 LLM 推論採用自迴歸方式，每生成一個 token 需完整前向傳遞，嚴重低度利用 TPU 等加速器的平行運算能力，尤其在低批次大小時。推測解碼透過小型「草稿」模型預測多個 token，再由「目標」模型平行驗證，若準確則以單步成本接受多 token，大幅降低延遲。但現有方法多依賴自迴歸草稿機制，需 K 個序列前向傳遞生成 K 個候選 token，此序列依賴形成執行瓶頸，限制加速潛力。DFlash 轉向區塊擴散（block diffusion），以 O(1) 單次前向傳遞生成整個區塊，消除序列猜測的時間消耗。 DFlash 在 TPU 上的擴散式草稿 DFlash 基於擴散 LLM（dLLM），利用目標模型提取的隱藏特徵，在單次前向傳遞生成高品質區塊草稿 token，複雜度從 O(K) 降至 O(1)，完美契合 TPU 的高頻寬矩陣乘法單元（MXUs）。UCSD 團隊在 Google Cloud 工程師指導下，將 DFlash 整合至 vLLM TPU 推論框架，優化記憶體頻寬與矩陣乘法飽和度，將草稿階段開銷最小化，最大化目標模型平行驗證吞吐量。 移植至 TPU/JAX 的三大挑戰 將 DFlash 從 GPU/PyTorch 移植至 TPU/JAX 生態需重新工程化，以匹配 TPU 架構優勢： 雙快取注意力解決方案：DFlash 的非因果區塊擴散與標準分頁注意力不相容，團隊設計雙快取架構——目標模型續用分頁 KV 快取與 Pallas 核心；草稿模型則用靜態 JAX 陣列，維持原設計並確保 TPU 原生效能。 智慧上下文管理：DFlash 草稿模型依賴目標模型中間隱藏狀態的上下文緩衝區，團隊採用 2 的冪次填充策略，優化 CPU-TPU 傳輸區塊，精準追蹤已消耗上下文，避免重複處理或資料遺失。 元資料同步：DFlash 具狀態性，依賴跨迭代的上下文緩衝、KV 快取位置與 RoPE 偏移；TPU vLLM 管道的元資料導致序列長度膨脹，團隊重新設計提案者嚴格同步真實接受 token 計數，恢復完美對齊，解鎖 TPU 上的精準區塊擴散。 基準測試：DFlash 對 EAGLE-3 的壓倒性勝利 在 TPU v5p 上，使用相同 Llama-3.1-8B 目標模型與官方檢查點（無微調），DFlash 端到端服務加速 2.29 倍，遠超 EAGLE-3 的 1.30 倍。EAGLE-3 自迴歸預測 2 token/步，需序列前向與 Python 協調開銷；DFlash 單步生成 10 個高品質 token，消除瓶頸。在 mbpp 程式碼任務，每 token 生成時間從 9.81ms 壓至 3.48ms，提升 2.83 倍。 TPU v5p 跨領域基準結果 UCSD 團隊建置獨立 JAX 基準，排除服務層開銷，評估 TPU v5p 上 DFlash： 平均加速 3.13 倍，數學推理峰值近 6 倍。 math500 任務：每 token 從 8.02ms 降至 1.40ms。 humaneval 程式碼評估：加速逾 3.5 倍。 結構化任務如數學與程式碼具高可預測性，接受長度長，飽和 TPU 平行驗證；對話任務隨機性高，加速較溫和。 推測效率洞見：K-Flat 突破 TPU v5p 等資料中心加速器驗證 1024 token 成本幾與 16 token 相同，因時間主導於模型權重載入而非注意力運算，「更寬驗證幾乎免費」。這推翻傳統觀念，瓶頸非驗證成本而是草稿品質，開發者可大膽擴大區塊大小，利用雙向上下文提升準確率。 擴展理論：品質勝於數量 儘管 K=16 已捕獲 90% 理論最大加速，擴至 K=128 僅增少於 1 token/步。提升每位置接受機率（a）效益為擴大 K 的 2–3 倍，未來焦點轉向智慧草稿訓練，而非更寬推測視窗。 任務驅動加速差異 接受機率依任務可預測性而定，區塊末尾 token 難猜測（位置衰減）。數學與程式碼衰減緩慢，維持高接受率；對話快速衰減。結構化推理更有效利用 TPU 潛力。 開源整合 vLLM 完整實作已提交 vLLM tpu-inference 儲存庫： PR #1868：DFlash 模型與提案者架構。 PR #1869：端到端推測解碼管道整合。 PR #1870：全面 CI 與端到端測試框架。 團隊正開發 torchax 提案者，支持 PyTorch 服務路徑。 未來前沿：TPU 推測系統擴展 DFlash 平行取樣奠基「Speculative Speculative Decoding (SSD)」，利用推測快取降低高吞吐延遲。計畫擴大草稿區塊捕獲豐富上下文，提升複雜推理接受率，使用 TPU RL Stack Tunix 與 MaxText；新 JAX 核心支持擴散目標模型，維持 vLLM-TPU 在非自迴歸生成前沿。 此成就源於 TPU Builder 計畫，提供硬體與 Google Cloud 點數予學術與開源社群。感謝 UCSD 團隊（Zhongyan Luo、Son Nguyen、Andy Huang）及 Google 工程師貢獻。詳見技術報告、Colab Notebook 與 vLLM GitHub。欲申請 TPU 研究，請聯絡 [email protected]。 原文：https://easyvibecoding.app/curated/1142