本研究由北京大學(xué)研究團隊完成。通訊作者為李萌，北京大學(xué)人工智能研究院和集成電路學(xué)院助理教授，博導(dǎo)，PKU SEC Lab 負責人，他的研究興趣集中于高效、安全人工智能加速算法和芯片，旨在通過算法到芯片的跨層次協(xié)同設(shè)計和優(yōu)化，為人工智能構(gòu)建高能效、高可靠、高安全的算力基礎(chǔ)。第一作者韋臨燁，北京大學(xué)集成電路學(xué)院博士一年級在讀，主要研究方向為多模態(tài)高效 AI 系統(tǒng)和加速器設(shè)計。

基于擴散的大語言模型 (dLLM) 憑借全局解碼和雙向注意力機制解鎖了原生的并行解碼和受控生成的潛力，最近吸引了廣泛的關(guān)注。例如 Fast-dLLM 的現(xiàn)有推理框架通過分塊半自回歸解碼進一步實現(xiàn)了 dLLM 對 KV cache 的支持，挑戰(zhàn)了傳統(tǒng)自回歸 LLMs 的統(tǒng)治地位。

然而，雙向注意力引入的 cache 刷新機制使這一推理過程成為預(yù)填充 / 解碼交錯的模式，將傳統(tǒng)自回歸模型解碼階段效率的訪存瓶頸限制轉(zhuǎn)化為計算 / 訪存瓶頸交替限制，占據(jù)了推理開銷中不可忽視的一部分。此外，認為整個序列 KV 狀態(tài)在分塊內(nèi)不變的緩存近似方式還一定程度上劣化了模型性能。

針對這一缺陷，來自北大的研究團隊提出一種新的 dLLM 推理加速框架 ODB-dLLM（Orchestrating Dual-Boundaries: An Arithmetic Intensity Inspired Acceleration Framework for Diffusion Language Models）。它通過分析現(xiàn)有 dLLM 推理框架中交錯的計算和訪存瓶頸階段，引入了自適應(yīng)長度預(yù)測策略和跳躍共享推測解碼，以優(yōu)化 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

研究背景和動機

傳統(tǒng) dLLM 模型中任意兩個 token 間的雙向注意力破壞了 KV cache 應(yīng)用的前提，讓整個推理過程陷入每一步都執(zhí)行預(yù)填充的計算瓶頸問題，沒有取得超越自回歸模型的實際推理速度。Fast-dLLM 及類似工作將預(yù)設(shè)定長度的輸出序列劃分為多個分塊，只在每個分塊解碼的第一步執(zhí)行預(yù)填充并存儲全局的的 KV 狀態(tài)，隨后計算分塊內(nèi)的 tokens 并基于置信度閾值接收。當一個分塊內(nèi)的所有位置都被解掩碼，轉(zhuǎn)到下一個分塊并刷新 KV cache。然而，這樣的推理框架仍然存在以下挑戰(zhàn)：

• 交錯的計算和訪存瓶頸模式都限制了效率。分塊半自回歸解碼的每個新塊都會觸發(fā)全局預(yù)填充以刷新 KV cache，占總推理延遲的 30-40%。隨著更細粒度的并行解碼設(shè)計來減少解碼步數(shù)，計算密集型的預(yù)填充階段將進一步占據(jù)主導(dǎo)地位。亟需一個算術(shù)強度感知的推理框架聯(lián)合緩解計算和訪存瓶頸。

• 以現(xiàn)有的 dLLM 無法感知實際任務(wù)需求，導(dǎo)致預(yù)設(shè)定輸出長度過長。dLLM 中的每個分塊都采用雙向注意力，以先前和后續(xù)令牌為條件。因此，需要固定的響應(yīng)長度來控制剩余的輸出空間。為了確保足夠的容量，該長度通常設(shè)置為較大的默認值（例如 1024），這會在計算密集型的預(yù)填充階段引入大量開銷并導(dǎo)致冗余推理。

• 硬閾值延遲了對潛在正確 token 的接受，導(dǎo)致更長的解碼軌跡和更多的步數(shù)。在訪存密集型的解碼階段，模型并行計算分塊內(nèi)的所有掩碼 token，但僅接受置信度超過閾值的部分，而其他 token 則被重新掩碼。大量正確 token 由于在當前步的置信度分數(shù)低于閾值而被推遲接收，從而導(dǎo)致額外的解碼步驟并增加不必要延遲。與此同時，直接降低閾值又會損害模型性能。

ODB-dLLM 框架

ODB-dLLM 推理框架的目標在于協(xié)調(diào)交錯的計算和訪存瓶頸限制。為了減輕預(yù)設(shè)定響應(yīng)長度帶來的大量預(yù)填充開銷，本文提出了一種自適應(yīng)長度預(yù)測策略來減少冗余計算。此外，本文引入了為 dLLM 特別定制的跳躍共享推測解碼方法，可在有限的算術(shù)強度下減少訪存受限階段的解碼步驟。

自適應(yīng)長度預(yù)測

該策略在每一輪預(yù)填充后考察整個響應(yīng)序列，檢測是否存在 [EOS] token。一旦識別出存在這一終止標記，就會提取其置信度分數(shù)并將其與預(yù)定義的截斷閾值進行比較。如果置信度超過閾值，則響應(yīng)將在該位置被截斷。如果未檢測到 [EOS] token，或者對應(yīng)的置信度均低于閾值，則響應(yīng)長度保持不變。這種漸進式策略最大限度地減少了預(yù)填充階段的算術(shù)強度和冗余計算，隨著模型的不確定性和多樣性逐漸收斂還帶來了性能增益。

跳躍共享推測解碼

推測解碼通過增加計算強度鼓勵多樣化預(yù)測，為訪存受限場景下進一步提高解碼的并行度提供了解決方案。然而，dLLM 推理的雙向注意力將 token 間的因果依賴轉(zhuǎn)化為全局依賴，在傳統(tǒng)推測解碼中構(gòu)造的遵循嚴格順序的 token tree 被全連通 token graph 所取代。因此，一方面候選 token graph 中任何獨立驗證的節(jié)點對都會解鎖其連接邊緣，從而實現(xiàn)更靈活的推測解碼探索；另一方面任意一個 token 的改變都需要計算一個完整的新塊，計算強度隨候選預(yù)測的增加而迅速增大并演化為訪存瓶頸，限制了推測的多樣性?；诖?，本文提出了兩點關(guān)鍵設(shè)計：

1. 跳躍接收推測策略。推測解碼的草稿塊也同時作為中間驗證塊，實現(xiàn)草稿 token 跨塊的跳躍驗證和接收，從而在單次解碼迭代中解掩碼更多 token，在有限的推測塊預(yù)算下顯著提高解碼效率。

2. 解碼共享推測策略。考慮各推測塊間的高度相似性，塊內(nèi)已解碼 token 的 KV 狀態(tài)也可以只在當前塊內(nèi)計算更新，并在多個推測塊之間共享。隨著解碼的進展，推測塊的計算成本逐漸降低，為更多推測快的拓展提供了機會。

實驗結(jié)果：更快、更準

研究團隊在兩種代表性的開源擴散語言模型LLaDA-Instruct和 LLaDA-1.5 上評估了 ODB-dLLM。在多個數(shù)據(jù)集上對比 Vanilla model 和 Fast-dLLM 推理框架結(jié)果如下：

• 團隊測量了預(yù)填充長度和延遲的減少，確認了計算密集階段有意義的算術(shù)強度節(jié)省。此外，每次預(yù)填充迭代期間響應(yīng)長度的逐步收縮隱式地將模型限制在更準確的答案空間，從而比生成到 [EOS] token 后簡單地截斷產(chǎn)生更高的精度。

• 針對訪存密集階段的推測解碼，實驗發(fā)現(xiàn)跳躍接收策略增加了每步接受的 token 數(shù)量，并通過一次解碼中的多段接收減少了并行解掩碼 token 間的潛在沖突。解碼共享策略的引入進一步減少了解碼輪次，同時優(yōu)先解掩碼之前未能接受但置信度較高的 token。跳躍共享推測解碼因而在速度和精度上都取得了顯著增益。

總結(jié)

受預(yù)填充和解碼階段異構(gòu)算術(shù)強度的啟發(fā)，本研究提出了 ODB-dLLM，這是一個協(xié)調(diào)計算 / 訪存雙邊界以加速 dLLM 推理的框架。在預(yù)填充階段，團隊發(fā)現(xiàn)預(yù)定義的固定響應(yīng)長度引入了繁重且冗余的計算開銷，ODB-dLLM 采用了自適應(yīng)長度預(yù)測機制，可逐步減少預(yù)填充開銷和不必要的計算。在解碼階段，團隊分析了擴散語言模型的計算特性，并提出了一種特別定制的跳躍共享推測解碼方法，通過減少解碼迭代次數(shù)來提高效率。

論文與代碼均已公開，歡迎感興趣的同學(xué)閱讀、復(fù)現(xiàn)以及深入討論。