Fast-dLLM: Aceleración sin entrenamiento de Diffusion LLM mediante la habilitación de la caché KV y decodificación paralela

Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding

May 28, 2025

Autores: Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, Enze Xie

cs.AI

Resumen

Los modelos de lenguaje basados en difusión (Diffusion LLMs) han mostrado potencial para la generación de texto no autoregresivo con capacidades de decodificación paralela. Sin embargo, la velocidad de inferencia práctica de los Diffusion LLMs de código abierto a menudo se queda atrás en comparación con los modelos autoregresivos, debido a la falta de una caché de Clave-Valor (KV) y a la degradación de calidad al decodificar múltiples tokens simultáneamente. Para cerrar esta brecha, introducimos un novedoso mecanismo de caché KV aproximada por bloques, diseñado específicamente para modelos de difusión bidireccionales, que permite la reutilización de la caché con una caída de rendimiento insignificante. Además, identificamos que la causa principal de la degradación de la calidad en la decodificación paralela es la ruptura de las dependencias entre tokens bajo el supuesto de independencia condicional. Para abordar este problema, proponemos una estrategia de decodificación paralela consciente de la confianza, que decodifica selectivamente los tokens que superan un umbral de confianza, mitigando las violaciones de dependencia y manteniendo la calidad de la generación. Los resultados experimentales en los modelos LLaDA y Dream, evaluados en múltiples benchmarks de LLM, demuestran una mejora de hasta 27.6 veces en el rendimiento con una pérdida mínima de precisión, cerrando la brecha de rendimiento con los modelos autoregresivos y allanando el camino para la implementación práctica de los Diffusion LLMs.

English

Diffusion-based large language models (Diffusion LLMs) have shown promise for non-autoregressive text generation with parallel decoding capabilities. However, the practical inference speed of open-sourced Diffusion LLMs often lags behind autoregressive models due to the lack of Key-Value (KV) Cache and quality degradation when decoding multiple tokens simultaneously. To bridge this gap, we introduce a novel block-wise approximate KV Cache mechanism tailored for bidirectional diffusion models, enabling cache reuse with negligible performance drop. Additionally, we identify the root cause of generation quality degradation in parallel decoding as the disruption of token dependencies under the conditional independence assumption. To address this, we propose a confidence-aware parallel decoding strategy that selectively decodes tokens exceeding a confidence threshold, mitigating dependency violations and maintaining generation quality. Experimental results on LLaDA and Dream models across multiple LLM benchmarks demonstrate up to 27.6times throughput improvement with minimal accuracy loss, closing the performance gap with autoregressive models and paving the way for practical deployment of Diffusion LLMs.