File size: 6,553 Bytes
3a464db
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
# Fast-dLLM v1: Training-free Acceleration of Diffusion LLM

[![Project](https://img.shields.io/static/v1?label=Project&message=Github&color=blue&logo=github-pages)](https://nvlabs.github.io/Fast-dLLM)
[![arXiv](https://img.shields.io/badge/Paper-arXiv-red.svg)](https://arxiv.org/abs/2505.22618)
<a href="https://fast-dllm.hanlab.ai"><img src="https://img.shields.io/static/v1?label=Demo&message=Fast-dLLM&color=yellow"></a> &ensp;

Fast-dLLM v1 is a **training-free** inference acceleration framework for diffusion-based Large Language Models (dLLMs). It supports efficient inference for models like **Dream** and **LLaDA** by enabling KV Cache and Parallel Decoding.

## Key Features

1. **Key-Value Cache for Block-Wise Decoding**
   We propose an efficient block-wise decoding KV Cache mechanism for Masked Diffusion Models (MDMs). By reusing attention Key-Value activations across multiple steps within each block, our approach avoids redundant computation and significantly accelerates inference. Furthermore, our DualCache extension also caches masked suffix tokens, enabling even greater speedup with negligible accuracy loss.

<div align="center">
  <img src="asset/kvcache.jpg" alt="KV Cache for block-wise decoding" width="800"/>
  <p>KV Cache for block-wise decoding</p>
</div>

2. **Confidence-Aware Parallel Decoding**
   Instead of decoding tokens sequentially, we introduce a confidence-aware parallel decoding scheme. At each step, only tokens with confidence over a threshold are unmasked in parallel, while uncertain ones remain masked for future steps. This selective approach effectively balances decoding efficiency and output quality.

<div align="center">
  <img src="asset/output.gif" alt="Decoding comparison" width="800"/>
  <p><b>Left:</b> Standard decoding (LLaDA). <b>Right:</b> Confidence-aware parallel decoding.</p>
</div>

<div align="center">
  <img src="asset/pseudo_code.jpg" alt="Pseudo code for our method" width="800"/>
  <p>Pseudo code for our method</p>
</div>

3. **Overall Performance**
   Overall, introducing the KV Cache mechanism yields significant speed improvements for all tasks and sequence lengths, typically achieving a 2x to 3.6x speedup compared to the vanilla backbone. When the parallel decoding strategy is applied individually, we see additional acceleration, often pushing speedups to 4x-6x for the evaluated settings, particularly as the generation length increases.

<div align="center">
  <img src="asset/overall_performance.jpg" alt="Overall performance" width="800"/>
  <p>Overall performance comparison</p>
</div>

## Demo

https://github.com/user-attachments/assets/32bbff97-6e60-4e14-95c0-2cbec136476f

<div align="center">
  <img src="asset/speedup.jpg" alt="End-to-end speedup over vanilla LLaDA baseline" width="800"/>
  <p>End-to-end speedup over vanilla LLaDA baseline</p>
</div>

## File Structure

```
v1/
β”œβ”€β”€ README.md               # This file
β”œβ”€β”€ requirements.txt        # Dependencies for inference & evaluation
β”œβ”€β”€ dream/                  # Dream model related code
β”‚   β”œβ”€β”€ model/              # Dream model definition
β”‚   β”œβ”€β”€ eval.py             # Evaluation harness integration
β”‚   β”œβ”€β”€ eval.md             # Evaluation guide
β”‚   β”œβ”€β”€ eval_gsm8k.sh       # GSM8K evaluation script
β”‚   β”œβ”€β”€ eval_humaneval.sh   # HumanEval evaluation script
β”‚   └── demo_multiturn_chat.py  # Multi-turn chat demo
└── llada/                  # LLaDA model related code
    β”œβ”€β”€ model/              # LLaDA model definition
    β”œβ”€β”€ generate.py         # Core generation with cache & parallel decoding
    β”œβ”€β”€ eval_llada.py       # Evaluation harness integration
    β”œβ”€β”€ eval.md             # Evaluation guide
    β”œβ”€β”€ eval_gsm8k.sh       # GSM8K evaluation script
    β”œβ”€β”€ eval_humaneval.sh   # HumanEval evaluation script
    β”œβ”€β”€ chat.py             # Command-line chat interface
    └── app.py              # Gradio web demo
```

## Installation

```bash
cd v1
pip install -r requirements.txt
```

## Usage

### 1. Using LLaDA Model

#### Interactive Chat
```bash
python llada/chat.py --gen_length 128 --steps 128 --block_size 32
```

Parameter descriptions:
- `--gen_length`: Maximum length of generated text
- `--steps`: Number of sampling steps
- `--block_size`: Cache block size
- `--use_cache`: Whether to use cache
- `--if_cache_position`: Whether to use dual cache
- `--threshold`: Confidence threshold

#### Web Demo
```bash
pip install gradio
cd llada
python app.py
```

#### Model Evaluation
| Benchmark         | Gen Length | LLaDA   | +Cache         | +Parallel      | +Cache+Parallel (Fast-dLLM) |
|-------------------|------------|---------|----------------|----------------|-----------------------------|
| **GSM8K (5-shot)**| 256        | 79.3<br>6.73<br>(1Γ—) | 79.5<br>21.23<br>(3.2Γ—) | 79.2<br>16.53<br>(2.5Γ—) | 78.5<br>**54.4<br>(8.1Γ—)** |
|                   | 512        | 77.5<br>3.23<br>(1Γ—) | 77.0<br>10.43<br>(3.3Γ—) | 77.6<br>18.63<br>(5.8Γ—) | 77.2<br>**35.3<br>(11.0Γ—)** |
| **HumanEval (0-shot)** | 256   | 41.5<br>30.5 (1Γ—) | 42.7<br>40.73<br>(1.3Γ—) | 43.9<br>101.53<br>(3.3Γ—) | 43.3<br>**114.1<br>(3.7Γ—)** |
|                   | 512        | 43.9<br>18.4 (1Γ—) | 45.7<br>29.33<br>(1.6Γ—) | 43.3<br>57.13<br>(3.1Γ—) | 44.5<br>**73.7<br>(4.0Γ—)** |

Each cell presents the accuracy (top row, in percentage) and the decoding throughput (middle row, in tokens per second) with relative speedup (bottom row) to the LLaDA baseline.

For detailed evaluation instructions, please refer to:
- [LLaDA Evaluation Guide](llada/eval.md)
- [Dream Evaluation Guide](dream/eval.md)

### 2. Using Dream Model

For detailed evaluation instructions on GSM8K and HumanEval benchmarks, please refer to [Dream Evaluation Guide](dream/eval.md).

## Citation

```bibtex
@misc{wu2025fastdllmtrainingfreeaccelerationdiffusion,
      title={Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding}, 
      author={Chengyue Wu and Hao Zhang and Shuchen Xue and Zhijian Liu and Shizhe Diao and Ligeng Zhu and Ping Luo and Song Han and Enze Xie},
      year={2025},
      eprint={2505.22618},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2505.22618}, 
}
```

## Acknowledgements

We would like to thank the authors of [LLaDA](https://github.com/llada-project/llada) and [Dream](https://github.com/dream-project/dream) for their excellent work and open-source contributions.