擴散語言模型推理太慢？北大團隊提出ODB-dLLM框架，破解計算訪存雙重瓶頸
本研究由北京大學研究團隊完成。通訊作者為李萌，北京大學人工智能研究院和集成電路學院助理教授，博導，PKU SEC Lab 負責人，他的研究興趣集中于高效、安全人工智能加速算法和芯片，旨在通過算法到芯片的跨層次協同設計和優化，為人工智能構建高能效、高可靠、高安全的算力基礎。第一作者韋臨燁，北京大學集成電路學院博士一年級在讀，主要研究方向為多模態高效 AI 系統和加速器設計。
基于擴散的大語言模型 (dLLM) 憑借全局解碼和雙向注意力機制解鎖了原生的并行解碼和受控生成的潛力，最近吸引了廣泛的關注。例如 Fast-dLLM 的現有推理框架通過分塊半自回歸解碼進一步實現了 dLLM 對 KV cache 的支持，挑戰了傳統自回歸 LLMs 的統治地位。
然而，雙向注意力引入的 cache 刷新機制使這一推理過程成為預填充 / 解碼交錯的模式，將傳統自回歸模型解碼階段效率的訪存瓶頸限制轉化為計算 / 訪存瓶頸交替限制，占據了推理開銷中不可忽視的一部分。此外，認為整個序列 KV 狀態在分塊內不變的緩存近似方式還一定程度上劣化了模型性能。
針對這一缺陷，來自北大的研究團隊提出一種新的 dLLM 推理加速框架 ODB-dLLM（Orchestrating Dual-Boundaries: An Arithmetic Intensity Inspired Acceleration Framework for Diffusion Language Models）。它通過分析現有 dLLM 推理框架中交錯的計算和訪存瓶頸階段，引入了自適應長度預測策略和跳躍共享推測解碼，以優化 dLLM 在硬件平臺上的計算訪存特性，最大限度地提高推理效率。
- 論文標題：Orchestrating Dual-Boundaries: An Arithmetic Intensity Inspired Acceleration Framework for Diffusion Language Models
- 論文鏈接：https://arxiv.org/abs/2511.21759
- Github 鏈接：https://github.com/PKU-SEC-Lab/ODB-dLLM
研究背景和動機
傳統 dLLM 模型中任意兩個 token 間的雙向注意力破壞了 KV cache 應用的前提，讓整個推理過程陷入每一步都執行預填充的計算瓶頸問題，沒有取得超越自回歸模型的實際推理速度。Fast-dLLM 及類似工作將預設定長度的輸出序列劃分為多個分塊，只在每個分塊解碼的第一步執行預填充并存儲全局的的 KV 狀態，隨后計算分塊內的 tokens 并基于置信度閾值接收。當一個分塊內的所有位置都被解掩碼，轉到下一個分塊并刷新 KV cache。然而，這樣的推理框架仍然存在以下挑戰：
- 交錯的計算和訪存瓶頸模式都限制了效率。分塊半自回歸解碼的每個新塊都會觸發全局預填充以刷新 KV cache，占總推理延遲的 30-40%。隨著更細粒度的并行解碼設計來減少解碼步數，計算密集型的預填充階段將進一步占據主導地位。亟需一個算術強度感知的推理框架聯合緩解計算和訪存瓶頸。
- 以現有的 dLLM 無法感知實際任務需求，導致預設定輸出長度過長。dLLM 中的每個分塊都采用雙向注意力，以先前和后續令牌為條件。因此，需要固定的響應長度來控制剩余的輸出空間。為了確保足夠的容量，該長度通常設置為較大的默認值（例如 1024），這會在計算密集型的預填充階段引入大量開銷并導致冗余推理。
- 硬閾值延遲了對潛在正確 token 的接受，導致更長的解碼軌跡和更多的步數。在訪存密集型的解碼階段，模型并行計算分塊內的所有掩碼 token，但僅接受置信度超過閾值的部分，而其他 token 則被重新掩碼。大量正確 token 由于在當前步的置信度分數低于閾值而被推遲接收，從而導致額外的解碼步驟并增加不必要延遲。與此同時，直接降低閾值又會損害模型性能。
ODB-dLLM 框架
ODB-dLLM 推理框架的目標在于協調交錯的計算和訪存瓶頸限制。為了減輕預設定響應長度帶來的大量預填充開銷，本文提出了一種自適應長度預測策略來減少冗余計算。此外，本文引入了為 dLLM 特別定制的跳躍共享推測解碼方法，可在有限的算術強度下減少訪存受限階段的解碼步驟。
自適應長度預測
該策略在每一輪預填充后考察整個響應序列，檢測是否存在 [EOS] token。一旦識別出存在這一終止標記，就會提取其置信度分數并將其與預定義的截斷閾值進行比較。如果置信度超過閾值，則響應將在該位置被截斷。如果未檢測到 [EOS] token，或者對應的置信度均低于閾值，則響應長度保持不變。這種漸進式策略最大限度地減少了預填充階段的算術強度和冗余計算，隨著模型的不確定性和多樣性逐漸收斂還帶來了性能增益。
跳躍共享推測解碼
推測解碼通過增加計算強度鼓勵多樣化預測，為訪存受限場景下進一步提高解碼的并行度提供了解決方案。然而，dLLM 推理的雙向注意力將 token 間的因果依賴轉化為全局依賴，在傳統推測解碼中構造的遵循嚴格順序的 token tree 被全連通 token graph 所取代。因此，一方面候選 token graph 中任何獨立驗證的節點對都會解鎖其連接邊緣，從而實現更靈活的推測解碼探索；另一方面任意一個 token 的改變都需要計算一個完整的新塊，計算強度隨候選預測的增加而迅速增大并演化為訪存瓶頸，限制了推測的多樣性?；诖耍疚奶岢隽藘牲c關鍵設計：
1. 跳躍接收推測策略。推測解碼的草稿塊也同時作為中間驗證塊，實現草稿 token 跨塊的跳躍驗證和接收，從而在單次解碼迭代中解掩碼更多 token，在有限的推測塊預算下顯著提高解碼效率。
2. 解碼共享推測策略?？紤]各推測塊間的高度相似性，塊內已解碼 token 的 KV 狀態也可以只在當前塊內計算更新，并在多個推測塊之間共享。隨著解碼的進展，推測塊的計算成本逐漸降低，為更多推測快的拓展提供了機會。
實驗結果：更快、更準
研究團隊在兩種代表性的開源擴散語言模型 LLaDA-Instruct 和 LLaDA-1.5 上評估了 ODB-dLLM。在多個數據集上對比 Vanilla model 和 Fast-dLLM 推理框架結果如下：
- 團隊測量了預填充長度和延遲的減少，確認了計算密集階段有意義的算術強度節省。此外，每次預填充迭代期間響應長度的逐步收縮隱式地將模型限制在更準確的答案空間，從而比生成到 [EOS] token 后簡單地截斷產生更高的精度。
- 針對訪存密集階段的推測解碼，實驗發現跳躍接收策略增加了每步接受的 token 數量，并通過一次解碼中的多段接收減少了并行解掩碼 token 間的潛在沖突。解碼共享策略的引入進一步減少了解碼輪次，同時優先解掩碼之前未能接受但置信度較高的 token。跳躍共享推測解碼因而在速度和精度上都取得了顯著增益。
總結
受預填充和解碼階段異構算術強度的啟發，本研究提出了 ODB-dLLM，這是一個協調計算 / 訪存雙邊界以加速 dLLM 推理的框架。在預填充階段，團隊發現預定義的固定響應長度引入了繁重且冗余的計算開銷，ODB-dLLM 采用了自適應長度預測機制，可逐步減少預填充開銷和不必要的計算。在解碼階段，團隊分析了擴散語言模型的計算特性，并提出了一種特別定制的跳躍共享推測解碼方法，通過減少解碼迭代次數來提高效率。
論文與代碼均已公開，歡迎感興趣的同學閱讀、復現以及深入討論。
