PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2505.22618v3cs.CL

TL;DR

Diffusion LLM inference is limited by missing KV caching and quality degradation during simultaneous token decoding. Fast-dLLM introduces block-wise approximate caching and confidence-aware parallel decoding, achieving up to 27.6× throughput improvement with minimal accuracy loss across evaluated models and benchmarks.

  • Problem

    Open-source Diffusion LLMs lack KV caching and can lose generation quality when decoding multiple tokens in parallel, limiting practical inference speed.

  • Method

    Fast-dLLM combines block-wise approximate KV caching for bidirectional diffusion models with confidence-thresholded parallel decoding.

  • Results

    Up to 27.6× throughput improvement is reported across LLaDA and Dream on GSM8K, MATH, HumanEval, and MBPP with minimal or no accuracy degradation.

  • Takeaways & Limitations

    Fast-dLLM closes part of the practical performance gap between Diffusion LLMs and autoregressive models while maintaining output quality.

  • Takeaways & Limitations

    At larger batch sizes, PrefixCache struggles to match LLaMA because diffusion decoding retains higher computational overhead from full attention operations.

Abstract

from arXiv · show

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 \textbf{27.6$\times$ throughput} improvement with minimal accuracy loss, closing the performance gap with autoregressive models and paving the way for practical deployment of Diffusion LLMs.

1. Introduction

Fast-dLLM addresses two practical barriers that keep open-source Diffusion LLMs slower than autoregressive models: missing KV caching and quality loss during parallel decoding. It combines approximate block-wise caching with confidence-aware token selection, achieving substantial throughput gains with minimal accuracy reduction.

  • Motivation: Open-source Diffusion LLMs lag autoregressive models because they lack KV caching and degrade in quality when decoding multiple tokens simultaneously.Diffusion LLMs perform best with one-token-at-a-time generation, while parallel decoding can reduce quality.
  • Key-Value Cache: Fast-dLLM reuses approximate KV activations across block-wise decoding steps and updates the cache after each block.The method is tailored to bidirectional attention and also supports DualCache for prefix and suffix tokens.
  • Parallel Decoding: Confidence-aware parallel decoding selects tokens whose confidence exceeds a global threshold instead of always selecting a fixed top-K set.The strategy targets dependency disruptions caused by conditional independence while preserving generation quality.
  • Results: 13.3× inference speed-up is achieved while preserving output quality through confidence-aware parallel decoding.This result is reported as part of the method’s contribution summary.
  • Results: 27.6× end-to-end speedup is achieved over the vanilla LLaDA baseline with long prefilling and a maximum generation length of 1024.The combined caching and parallel-decoding approach is evaluated on an NVIDIA A100 with a single batch size.
  • Results: Experiments on LLaDA and Dream across GSM8K, MATH, HumanEval, and MBPP show order-of-magnitude speedups with minimal or no accuracy degradation.The contribution summary describes the results as consistent across the evaluated models and benchmarks.

2. Preliminary

Masked Diffusion Models progressively replace tokens with [MASK] and reverse this process through iterative denoising. Parallel recovery accelerates generation, but conditional independence can disrupt dependencies among simultaneously generated tokens and reduce coherence.

  • Masked Diffusion Models: Masked Diffusion Models progressively replace tokens with a special [MASK] token during the forward noising process.They are also called absorbing state discrete diffusion models.
  • Masked Diffusion Models: Diffusion time t controls interpolation between the original sequence at t = 0 and a fully masked sequence at t = 1.The masking level determines the position along this forward process.
  • Generation Process: The reverse process is commonly approximated with τ-leaping so multiple masked tokens can be recovered in one step.This reduces the one-token-per-step computation of the analytical reverse process.
  • Generation Process: For masked positions, the model provides a vocabulary distribution for predicting a non-[MASK] token.This distribution is used within the reverse denoising process.
  • Parallel Decoding: Parallel decoding samples token positions using marginal distributions under a conditional independence assumption, despite dependencies in the true joint distribution.The mismatch can degrade sequence quality and coherence, especially when many tokens are unmasked simultaneously.

3. Methodology

Fast-dLLM combines block-wise approximate KV caching with confidence-aware parallel decoding to reduce redundant attention computation while preserving generation quality. Its decoding strategy is theoretically motivated by conditions under which independent marginal decoding matches greedy sequential decoding.

  • Key-Value Cache for Block-Wise Decoding: Fast-dLLM uses block-wise decoding so attention activations from previously decoded blocks can be cached and reused.The cache is initialized for the prompt, reused within each block, and updated after block completion.
  • Key-Value Cache for Block-Wise Decoding: High similarity between adjacent-step KV activations supports approximate cache reuse with minimal accuracy loss.The similarity is consistently close to 1 within blocks, indicating negligible changes in prefix keys and values during decoding.
  • Key-Value Cache for Block-Wise Decoding: DualCache additionally caches suffix tokens, producing further acceleration because suffix keys and values also change negligibly during block decoding.The suffix consists entirely of masked tokens under the block-wise decoding scheme.
  • Confidence-Aware Parallel Decoding: Confidence-aware decoding selectively unmasks tokens whose confidence exceeds a threshold, while always unmasking the highest-confidence token when needed to ensure progress.This replaces aggressive simultaneous unmasking based only on independent marginal probabilities.
  • Confidence-Aware Parallel Decoding: Theorem 1 states that greedy parallel decoding matches greedy sequential decoding when (n + 1)epsilon ≤ 1, or epsilon ≤ 1/(n + 1).The theorem compares the product-of-marginals distribution with the true joint distribution in a high-confidence regime.
  • Confidence-Aware Parallel Decoding: Factor-based decoding dynamically selects the largest number of parallel tokens satisfying (n + 1)(1 − c(n)) < f.It sorts confidence estimates and decodes the top-n tokens, linking the parallelism level to the theoretical confidence bound.

4. Experiments

Experiments evaluate Fast-dLLM across text and multimodal models, tasks, sequence lengths, and decoding strategies. The results show substantial throughput gains with small accuracy changes, while cache configuration and workload characteristics affect the trade-off.

  • Evaluation Setup: Fast-dLLM is evaluated on LLaDA, LLaDA-1.5, Dream, and multimodal LLaDA-V across math, code-generation, and vision-language benchmarks.The text benchmarks are GSM8K, MATH, HumanEval, and MBPP; multimodal evaluation uses MathVista and MathVerse.
  • Main Results: 2× to 3.6× speedups typically result from KV caching, while parallel decoding alone often reaches 4×–6×, especially for longer generations.Throughput is measured end-to-end in output tokens per second until the <eos> token.
  • Main Results: 11× on LLaDA GSM8K and 9.2× on LLaDA MBPP at length 512 are achieved when KV Cache and parallel decoding are combined.Dream-Base reaches 7.8× on MBPP and 5.6× on GSM8K at length 512; the techniques are reported as complementary.
  • Accuracy and Generalization: Accuracy remains within 1–2 points of the backbone across benchmarks and settings, with slight improvements in several cases.Longer sequences benefit proportionally more from cache reuse and batch computation, and LLaDA-1.5 has stronger accuracy with comparable or higher throughput.
  • Accuracy and Generalization: 9.9× speedup is obtained on multimodal LLaDA-V with minimal accuracy degradation after retaining full blocks and using refresh-based updates.Reducing block size from 96 to 8 otherwise drops MathVista accuracy by over 8%; MathVerse accuracy slightly improves under Fast-dLLM.
  • Cache Analysis: 27.6× DualCache speedup is reported for 8-shot, generation length 1024, exceeding the prefix cache’s 18.6× in the same scenario.DualCache generally performs better at longer generation lengths while maintaining competitive accuracy with minor trade-offs.

5. Related Work

Discrete diffusion research develops non-autoregressive generation through theoretical foundations, integration with pretrained language models, and inference optimizations. KV Cache methods address the difficulty of reusing attention states in bidirectional diffusion models.

  • Integration with Pre-trained Language Models: Diffusion-NAT combines discrete diffusion denoising with BART-based non-autoregressive decoding, reporting generation speed 20× faster than comparable AR transformers.This integration leverages pretrained language-model knowledge while iteratively refining masked tokens.
  • Key-Value Cache: KV Cache stores and reuses previously computed attention states, but full attention makes it difficult to apply directly to diffusion language models such as LLaDA.Block diffusion addresses this limitation by generating block-by-block so earlier blocks’ keys and values can be reused.
  • Non-Autoregressive Generation: Non-autoregressive generation accelerates inference by generating multiple tokens simultaneously and has expanded from translation to correction, summarization, dialogue, and speech recognition.Its central distinction from autoregressive generation is simultaneous rather than sequential token generation.

6. Conclusion

Fast-dLLM addresses Diffusion LLM inference limitations by combining approximate KV caching with confidence-aware parallel decoding. Across LLaDA and Dream benchmarks, it reports up to 27.6× speedup with minimal accuracy loss.

  • 6. Conclusion: Fast-dLLM targets missing KV Cache support and quality degradation during parallel decoding in Diffusion LLMs.Its block-wise generation enables an approximate cache suited to bidirectional attention.
  • 6. Conclusion: Fast-dLLM combines block-wise approximate KV Cache with Confidence-Aware Parallel Decoding to support safe and efficient multi-token generation.The decoding strategy addresses token-dependency disruption caused by the conditional independence assumption.
  • 6. Conclusion: 27.6× speedup is achieved across multiple benchmarks and LLaDA and Dream model baselines with minimal loss in accuracy.The reported result supports Fast-dLLM as a practical approach for Diffusion LLM deployment.

A. Proof

The proof establishes an argmax guarantee under high-confidence marginal predictions and bounds distributional distances. It also shows the confidence condition is tight and notes assumptions and worst-case scope.

  • A. Proof: The proof begins by showing that x* is the unique maximizer of the product-of-marginals distribution q.The argument maximizes each marginal term independently under the condition (n+1)ε≤1.
  • A. Proof: Under (n+1)ε<1, x* is also the unique maximizer of the joint distribution p, establishing the main argmax equivalence result.The proof compares p(x*|E) against every alternative sequence using high-confidence marginal conditions and event bounds.
  • A. Proof: The bound ε≤1/(n+1) is tight: when ε exceeds it, high-confidence marginals can coexist with different argmax outcomes for p and q.A binary-vocabulary construction demonstrates this failure case explicitly.
  • A. Proof: The proof derives bounds for Lp distance and total variation distance between p and q.The supplied result gives 2D1(p,q) < (3n−1)/2 ε.
  • A. Proof: The theorem assumes a coherent, well-defined joint probability mass function whose marginals are consistently derived, although practical MDMs may only approximate this property.The analysis also emphasizes that its guarantees are worst-case and practical distributions may behave better than the bounds suggest.

B. Case Study

The case study examines response quality under caching strategies, block sizes, and confidence thresholds. Across the reported arithmetic examples, outputs remain correct and clearly explained, while higher thresholds increase computation.

  • B.1. Effect of Caching Strategies on Response Quality: All Original, PrefixCache, and DualCache responses correctly compute the arithmetic answer 3 bolts with similar step-by-step reasoning.The comparison reports only minor phrasing differences, with semantic fidelity and interpretability unaffected.
  • B.2. Effect of Block Size: DualCache produces correct, clearly explained answers for block sizes 8, 16, and 32, with no meaningful output differences on the arithmetic prompt.The case indicates robustness to block size for this example.
  • B.3. Effect of Dynamic Threshold: Higher thresholds increase NFE from 9 to 20, while answer quality remains stable across threshold values 0.7, 0.8, and 0.9.For straightforward arithmetic questions, threshold adjustment mainly affects efficiency rather than correctness.

B.4. Multimodal Generation with LLAda-V

Fast-dLLM substantially accelerates LLaDA-V visual description while preserving detailed, faithful, and stylistically fluent captions. In the case study, it delivers comparable multimodal output in a fraction of the decoding time.

  • The visual case study uses identical visual input and the prompt “Please describe the image in detail.”
  • 63.0 seconds: the baseline produces a detailed, poetic rural-landscape description highlighting the barn, pink sky, and tranquil atmosphere.
  • 6.8 seconds: Fast-dLLM generates a richly detailed caption, achieving a nearly 10× speedup over the baseline.The caption adds grounding details such as gray roof shingles while aligning with appearance and mood cues.
  • Fast-dLLM enables high-quality vision-language generation at significantly improved efficiency for faster and more interactive multimodal applications.

C.1. Further Experiments with LLaDA-V

LLaDA-V performance is sensitive to block length, so the experiments retain a block length of 96 and vary cache refresh intervals. Longer intervals improve throughput but slightly reduce accuracy.

  • Highest MathVista accuracy occurs with block length 96 under a fixed decoding length of 48 steps.
  • Block lengths of 8 or 4 reduce MathVista accuracy by over 8%.
  • The method therefore keeps block length fixed at 96 and updates the cache only every r decoding steps using the most recent full block.
  • 28.2 tokens/s: throughput at refresh interval 32, compared with 15.9 tokens/s at interval 2.Increasing the refresh interval produces consistent throughput gains, while accuracy drops slightly with larger intervals.

C.2. Performance Comparison between Threshold and Factor Strategy

Factor-based confidence-aware decoding generally trades a small accuracy reduction for substantially higher throughput than threshold decoding. Across LLaDA experiments, LLaDA-1.5 further improves the accuracy–efficiency trade-off in most settings, with a longer-sequence regression on MATH.

  • Threshold versus Factor Strategy: Factor decoding achieves 1.4–1.5× higher throughput than threshold decoding across GSM8K and MATH settings.
  • Threshold versus Factor Strategy: 78.5 tokens/sec: factor decoding on GSM8K with 256 tokens, versus 54.4 tokens/sec for threshold decoding.The corresponding speedups are 11.7× and 8.1×, respectively.
  • Threshold versus Factor Strategy: Factor decoding typically sacrifices 1–3% accuracy for 40–50% higher throughput across both benchmarks and generation lengths.
  • LLaDA versus LLaDA-1.5: LLaDA-1.5 improves GSM8K accuracy by +2.2% at 256 tokens and +3.2% at 512 tokens over original LLaDA.
  • LLaDA versus LLaDA-1.5: 59.4 tokens/sec: LLaDA-1.5 throughput on GSM8K with 256 tokens, compared with 54.1 tokens/sec for LLaDA.
  • LLaDA versus LLaDA-1.5: On MATH, LLaDA-1.5 has comparable accuracy, improves throughput at 256 tokens, but falls from 47.1 to 41.1 tokens/sec at 512 tokens.

C.4. Analysis of Parallel Token Counts across Decoding Steps

Parallel token generation increases during the middle of decoding and declines near completion, where variability rises. These dynamics motivate aggressive parallelism during high-confidence phases and conservative decoding near boundaries; caching also improves LLaDA throughput, especially for shorter generations and larger batches.

  • Parallel Token Counts: Average parallel token generation rises during early-to-middle decoding, peaks around steps 30–60, then slightly declines near generation end.
  • Parallel Token Counts: Later decoding steps show greater variance across samples, partly because only a few remaining tokens differ across outputs.
  • Parallel Token Counts: High-confidence middle phases offer computational savings through increased parallelism, while conservative behavior near boundaries helps maintain quality.
  • Experimental Setup: The experiments use an NVIDIA A100, 256-token prefill, generation lengths of 16–64, and batch sizes from 1 to 32.
  • Experimental Setup: The caching comparison isolates acceleration from caching because parallel decoding is affected by dummy input tokens.
  • Caching and Throughput: PrefixCache improves LLaDA throughput across batch sizes and generation lengths, especially with shorter generations and larger batches.
  • Caching and Throughput: Over 211 tokens/s: PrefixCache throughput at generation length 16 and batch size 32.
  • Caching and Throughput: LLaDA throughput plateaus after batch size 8, whereas LLaMA benefits substantially from larger batches.
Loading 2505.22618v3…