英偉達與港大等合作推出快速 KV 緩存，助力擴散模型提速

近日，英偉達、香港大學與麻省理工學院的研究團隊聯合發佈了一種名爲 Fast-dLLM 的創新技術，旨在提升擴散語言模型的推理效率。與傳統的自迴歸模型不同，擴散語言模型採用逐步去除文本噪聲的方式生成文本，因此能夠在一次迭代中生成多個單詞，整體效率更高。然而，在實際應用中，許多開源擴散語言模型的推理速度仍然不如自迴歸模型，主要受限於缺乏鍵值（KV）緩存支持以及並行解碼時生成質量下降。

KV 緩存是自迴歸模型中常用的加速推理技術，通過存儲和重用之前計算的注意力狀態，顯著減少重複計算，從而提高生成速度。但由於擴散語言模型採用雙向注意力機制，直接應用 KV 緩存並不容易。Fast-dLLM 架構的創新之處在於將文本生成過程劃分爲多個塊，每個塊包含一定數量的 token。通過這種塊狀生成方式，模型可以在生成一個塊之前預先計算並存儲其他塊的 KV 緩存，進而避免重複計算。

儘管 KV 緩存機制有效提升了推理速度，但在並行解碼時，生成質量往往會下降。這是因爲擴散型模型在解碼時假設條件獨立性，然而標記之間可能存在複雜的依賴關係。爲了解決這一問題，Fast-dLLM 提出了一種基於置信度的平行解碼策略。在每個解碼步驟中，模型會計算每個標記的置信度，並選擇置信度超過閾值的標記進行解碼。這種策略確保了高置信度的情況下能夠安全進行並行解碼，從而維持生成文本的連貫性和準確性。

爲驗證 Fast-dLLM 的性能，研究人員在 NVIDIA A10080GB GPU 上對 LLaDA 和 Dream 兩種擴散語言模型進行了全面評估，涵蓋了數學推理和代碼生成等任務。在 KV 緩存機制的測試中，塊大小爲32時，模型的吞吐量達到了54.4tokens/s，準確率爲78.5%。並行解碼測試中，採用動態閾值策略優於固定 token 數基線。整體來看，LLaDA 模型在 GSM8K 任務中僅用 KV Cache 加速3.2倍，並行解碼加速2.5倍，兩者結合的速度提升達到8.1倍，生成長度達到1024時的端到端加速更是高達27.6倍。所有測試結果顯示，Fast-dLLM 在加速的同時，保持了生成質量的穩定。