Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FlashDLM: Accelerating Diffusion Language Model Inference via Efficient KV Caching and Guided Diffusion

Zhanqiu Hu, Jian Meng, Yash Akhauri, Mohamed S. Abdelfattah, Jae-sun Seo, Zhiru Zhang, Udit Gupta

TL;DR

This work tackles the latency bottlenecks of diffusion language models (DLMs) by introducing two training-free techniques: FreeCache, which reduces KV recomputation through a progressively frozen-window caching strategy, and Guided Diffusion, which pairs the DLM with a lightweight autoregressive guider to safely unmask tokens in parallel. Together, these methods yield substantial end-to-end speedups (average 12.14x overall, up to 34x in long-context tasks) with negligible accuracy degradation, enabling DLMs to match or exceed autoregressive latency on multiple benchmarks. The approach is validated on Dream-7B-Instruct and LLaDA-8B-Instruct across reasoning and QA tasks, showing strong cross-domain performance and robust long-context support without any finetuning. The work highlights a practical pathway for deploying diffusion-based reasoning systems at scale and in high-throughput settings.

Abstract

Diffusion language models offer parallel token generation and inherent bidirectionality, promising more efficient and powerful sequence modeling compared to autoregressive approaches. However, state-of-the-art diffusion models (e.g., Dream 7B, LLaDA 8B) suffer from slow inference. While they match the quality of similarly sized autoregressive (AR) models (e.g., Qwen2.5 7B, Llama3 8B), their iterative denoising requires multiple full-sequence forward passes, resulting in high computational costs and latency, particularly for long input prompts and long-context scenarios. Furthermore, parallel token generation introduces token incoherence problems, and current sampling heuristics suffer from significant quality drops with decreasing denoising steps. We address these limitations with two training-free techniques. First, we propose FreeCache, a Key-Value (KV) approximation caching technique that reuses stable KV projections across denoising steps, effectively reducing the computational cost of DLM inference. Second, we introduce Guided Diffusion, a training-free method that uses a lightweight pretrained autoregressive model to supervise token unmasking, dramatically reducing the total number of denoising iterations without sacrificing quality. We conduct extensive evaluations on open-source reasoning benchmarks, and our combined methods deliver an average of 12.14x end-to-end speedup across various tasks with negligible accuracy degradation. For the first time, diffusion language models achieve a comparable and even faster latency as the widely adopted autoregressive models. Our work successfully paved the way for scaling up the diffusion language model to a broader scope of applications across different domains.

FlashDLM: Accelerating Diffusion Language Model Inference via Efficient KV Caching and Guided Diffusion

TL;DR

This work tackles the latency bottlenecks of diffusion language models (DLMs) by introducing two training-free techniques: FreeCache, which reduces KV recomputation through a progressively frozen-window caching strategy, and Guided Diffusion, which pairs the DLM with a lightweight autoregressive guider to safely unmask tokens in parallel. Together, these methods yield substantial end-to-end speedups (average 12.14x overall, up to 34x in long-context tasks) with negligible accuracy degradation, enabling DLMs to match or exceed autoregressive latency on multiple benchmarks. The approach is validated on Dream-7B-Instruct and LLaDA-8B-Instruct across reasoning and QA tasks, showing strong cross-domain performance and robust long-context support without any finetuning. The work highlights a practical pathway for deploying diffusion-based reasoning systems at scale and in high-throughput settings.

Abstract

Diffusion language models offer parallel token generation and inherent bidirectionality, promising more efficient and powerful sequence modeling compared to autoregressive approaches. However, state-of-the-art diffusion models (e.g., Dream 7B, LLaDA 8B) suffer from slow inference. While they match the quality of similarly sized autoregressive (AR) models (e.g., Qwen2.5 7B, Llama3 8B), their iterative denoising requires multiple full-sequence forward passes, resulting in high computational costs and latency, particularly for long input prompts and long-context scenarios. Furthermore, parallel token generation introduces token incoherence problems, and current sampling heuristics suffer from significant quality drops with decreasing denoising steps. We address these limitations with two training-free techniques. First, we propose FreeCache, a Key-Value (KV) approximation caching technique that reuses stable KV projections across denoising steps, effectively reducing the computational cost of DLM inference. Second, we introduce Guided Diffusion, a training-free method that uses a lightweight pretrained autoregressive model to supervise token unmasking, dramatically reducing the total number of denoising iterations without sacrificing quality. We conduct extensive evaluations on open-source reasoning benchmarks, and our combined methods deliver an average of 12.14x end-to-end speedup across various tasks with negligible accuracy degradation. For the first time, diffusion language models achieve a comparable and even faster latency as the widely adopted autoregressive models. Our work successfully paved the way for scaling up the diffusion language model to a broader scope of applications across different domains.

Paper Structure

This paper contains 26 sections, 6 equations, 4 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Diffusion Language Models
  4. KV Caching for Autoregressive LLM
  5. Speculative Decoding
  6. Methodology
  7. Preliminary of Diffusion Language Model
  8. Bottleneck of Diffusion Language Model Inference
  9. Token Incoherence in Parallel Generation
  10. Proposed Method
  11. FreeCache: Reducing Window Caching for DLM Inference
  12. Guiding Diffusion with Lightweight Autoregressive Model
  13. Experiments
  14. Comparison with the baseline diffusion language model
  15. Comparison with the Vanilla Autoregressive LLM models
  16. ...and 11 more sections

Figures (4)

  • Figure 1: (a) Quality degradation under aggressive unmasking. (b) Pareto front of accuracy vs. latency (log‐scale) for baseline and optimized cached methods.
  • Figure 2: Rationale behind the proposed FreeCache: The variation of K and V projections of clean tokens is small throughout the subsequent diffusion steps.
  • Figure 3: AR-guided Diffusion Model. The diffusion model performs a one-step diffusion process, followed by a one-time forward pass of the AR guider. The matched tokens are unmasked for the next step of diffusion.
  • Figure : Guided Diffusion