Fast-dLLM: Aceleração sem Treinamento de Difusão LLM ao Habilitar Cache KV e Decodificação 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

Resumo

Modelos de linguagem baseados em difusão (Diffusion LLMs) têm mostrado potencial para geração de texto não autorregressiva com capacidades de decodificação paralela. No entanto, a velocidade prática de inferência dos Diffusion LLMs de código aberto frequentemente fica atrás dos modelos autorregressivos devido à falta de Cache de Chave-Valor (KV Cache) e à degradação de qualidade ao decodificar múltiplos tokens simultaneamente. Para preencher essa lacuna, introduzimos um novo mecanismo de Cache KV aproximado por blocos, projetado especificamente para modelos de difusão bidirecionais, permitindo a reutilização do cache com uma queda de desempenho insignificante. Além disso, identificamos a causa raiz da degradação da qualidade de geração na decodificação paralela como a interrupção das dependências entre tokens sob a suposição de independência condicional. Para resolver isso, propomos uma estratégia de decodificação paralela baseada em confiança que decodifica seletivamente tokens que excedem um limiar de confiança, mitigando violações de dependência e mantendo a qualidade da geração. Resultados experimentais nos modelos LLaDA e Dream em múltiplos benchmarks de LLMs demonstram uma melhoria de até 27,6 vezes na taxa de transferência com perda mínima de precisão, reduzindo a diferença de desempenho em relação aos modelos autorregressivos e abrindo caminho para a implantação prática de 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.