CAI-DLLM: Convergence Aware Inference for Diffusion Language Models

Farhana Amin; Sabiha Afroz; Dimitrios S. Nikolopoulos

cs.AIAug 23, 2026

CAI-DLLM: Convergence Aware Inference for Diffusion Language Models

Authors: Farhana Amin, Sabiha Afroz, Dimitrios S. Nikolopoulos

Organizations: Virginia Tech

Abstract

Diffusion language models can generate many tokens in parallel, but they still require repeated denoising steps during inference. This makes generation costly, especially when the model continues to recompute tokens that are already stable. To address these limitations, we propose CAI-DLLM, a training-free inference method that uses first-step confidence to guide denoising and reduce inference time. Specifically, CAI-DLLM commits easy tokens earlier, allocates more denoising steps to harder tokens, and adjusts decoding schedules across output blocks. As it relies only on first-step confidence signals, it does not require retraining, extra predictors, or weight updates. We evaluate CAI-DLLM on LLaDA-8B-Instruct and Dream-7B-Instruct across math, code, reasoning, commonsense, and long-context tasks. CAI-DLLM achieves up to 18.2x wall clock inference speedup on LLaDA GSM8K while improving accuracy from 76.27% to 77.41%, and up to 13.1x speedup on Dream HumanEval while achieving higher pass@1 than no-cache inference, 48.17% compared with 46.95%. On harder reasoning tasks, speedups reach 44.8x, with a largest accuracy drop of 4.4 points, while energy consumption is reduced by up to 95.3%.

Figures & tables

Figure 1: Denoising work is uneven across layers, blocks, and tokens. (a) Hidden state drift varies across transformer layers. (b) Saved steps increase in later blocks as earlier context becomes available. (c) First-step confidence separates easy tokens with $c_{i}^{(0)}>0.50$ , medium tokens with $0.25\leq c_{i}^{(0)}\leq 0.50$ , and hard tokens with $c_{i}^{(0)}<0.25$ .

Figure 2: Overview of CAI-dLLM . ① The method runs a step-zero forward pass to obtain token confidence scores. ② These scores drive a confidence-aware controller that guides adaptive parallel decoding, block-adaptive scheduling, and token-level step budgets. ③ During iterative denoising, easy tokens are committed early while hard tokens receive more steps. ④ A grinding-phase detector stops low-yield late denoising. The method is training-free and does not change model weights.

Figure 3: Accuracy and speedup comparison across the seven main benchmarks for Dream-7B and LLaDA-8B. LongBench is reported in detail in Appendix D.2 .

Bench.	Method	TPS	Spd.	Score	Eng.
GSM8K	No Cache	19.8	$1.0\times$	79.68	3052
	DualCache	103.2	$9.3\times$	78.47	291
	ES-dLLM	97.9	$9.0\times$	77.48	219
	CAI-dLLM	369.5	$\mathbf{18.7\times}$	$77.26{\pm}1.18$	146
Human Eval	No Cache	44.3	$1.0\times$	46.95	374
	DualCache	258.1	$5.8\times$	45.12	60

Table 1: Results on Dream-7B-Instruct. TPS = tokens/sec; Eng. = energy in Wh. Bold = best. DualCache from 29 . CAI-dLLM scores show std. err. as

\pm

Bench.	Method	TPS	Spd.	Score	Eng.
GSM8K	No Cache	17.1	$1.0\times$	76.27	3559
	DualCache	196.4	$11.7\times$	77.18	264
	ES-dLLM	175.3	$10.6\times$	77.10	223
	CAI-dLLM	335.1	$\mathbf{18.2\times}$	$77.41{\pm}1.18$	169
Human Eval	No Cache	39.0	$1.0\times$	37.20	425
	DualCache	173.0	$4.4\times$	35.98	55

Table 2: Results on LLaDA-8B-Instruct. TPS = tokens/sec; Eng. = energy in Wh. Bold = best. Sample sizes in Appendix A . CAI-dLLM scores show std. err. as

\pm

Model	Config.	Pass@1	TPS
LLaDA	ES-dLLM	35.98	139.8
	+APD	35.37	424.3
	+APD+Block	35.37	434.2
	+APD+Block+Budget $\star$	35.37	435.3
Dream	DualCache	45.12	72.5
	ES-dLLM	41.46	56.7

Table 3: HumanEval ablation of CAI-dLLM components.

\star

marks the recommended configuration for each model.

Task	Model	Fixed	CAI	Reduction
GSM8K	LLaDA	256	114	55.5%
GSM8K	Dream	256	141	44.9%
HumanEval	LLaDA	512	315	38.5%
HumanEval	Dream	512	271	47.1%

Table 4: Average denoising steps per request and step reduction relative to fixed budget.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Reasoning	Code	Commonsense	LongBench
Gen. tokens	256	512	256	512
Block length	64	64	64	32
Batch size	8	8	8	8
Temperature	0	0	0	0

Appendix

Table 5: Generation settings, identical across all compared methods.

Benchmark	Split	Used	Full	Notes
GSM8K	test	1,319	1,319	Full test set
BBH	test	6,511	6,511	Full; 23 subtasks
HumanEval	test	164	164	Full test set
MBPP	test	257	257	Full sanitized set
WinoGrande	validation	500	1,267	Standard eval split
PIQA	validation	1,838	1,838	Full validation set

Appendix

Table 6: Sample sizes per benchmark. Subsets are used only to control evaluation cost, not due to any limitation of CAI-dLLM .

Task	Metric	Samples	Max input
NarrativeQA	F1	100	4,000
MultifieldQA-en	F1	100	4,000
GovReport	ROUGE-L	100	4,000
QMSum	ROUGE-L	100	4,000
Total	–	400	–

Appendix

Table 7: LongBench tasks used in evaluation.

Benchmark	Metric	Tool
GSM8K	Exact match	lm-eval
BBH	Exact match	lm-eval
HumanEval	pass@1	lm-eval
MBPP	pass@1	lm-eval
MathQA	Exact match	lm-eval
WinoGrande	Exact match	lm-eval

Appendix

Table 8: Evaluation metric and tool per benchmark.

Bench.	Dream acc.	LLaDA acc.	Std. err.
GSM8K	77.26%	77.41%	$\pm$ 1.18 pp
BBH	56.95%	52.33%	$\pm$ 0.53 pp
MBPP	56.60%	37.00%	$\pm$ 2.16 pp
HumanEval	48.17%	35.98%	$\pm$ 3.76 pp

Appendix

Table 9: CAI-dLLM accuracy and standard errors.

Method	Hours	% Total
No cache	77.38	77.9%
ES-dLLM	8.19	8.2%
DualCache	8.07	8.1%
CAI-dLLM	3.89	3.9%
CAI-dLLM no CG	1.85	1.9%
Total	99.38	100%

Appendix

Table 10: Inference GPU-hours by method. CG denotes confidence gating.

Model	Config.	Acc.	TPS
LLaDA	ES-dLLM	77.10	175.3
	+APD	76.80	490.2
	+APD+Block	76.95	496.0
	+APD+Block+Budget $\star$	77.41	335.1
	Full CAI, gate off	76.57	320.0
Dream	ES-dLLM	77.48	97.9

Appendix

Table 11: GSM8K ablation of CAI-dLLM components.

\star

marks the recommended configuration for each model.

Figure 4: Attention sparsity across denoising steps for sampled layers of LLaDA-8B on GSM8K. We sample early, middle, and late transformer layers. The plots show attention entropy, effective tokens attended, and top-10 attention mass. The strong layer-to-layer variation supports layer-aware analysis of diffusion language model inference.

Figure 5: Accuracy across APD schedule configurations on LLaDA-8B GSM8K. The selected cosine schedule with

\theta_{s}=0.90

\theta_{e}=0.70

, and

w=0.15

gives the best accuracy-speed trade-off.

Figure 6: Hardware behavior across batch sizes. Fixed-threshold parallel decoding and APD are compared with ES-dLLM. The H200 runs the workload in a compute-bound regime, making denoising-step reduction an effective optimization target.

Model	Method	Peak Mem.	vs DualCache
LLaDA	No cache	41,567 MB	+9,260 MB
LLaDA	DualCache	32,307 MB	baseline
LLaDA	ES-dLLM	32,503 MB	+196 MB
LLaDA	CAI-dLLM	32,787 MB	+480 MB
Dream	No cache	44,033 MB	+13,322 MB
Dream	DualCache	30,711 MB	baseline

Appendix

Table 12: Peak GPU memory on HumanEval. The extra memory for CAI-dLLM over DualCache comes from per-token confidence scores, budgets, and threshold tables. Dream’s larger overhead is due to fuller parallel output generation, not CAI-dLLM bookkeeping.

Task	Metric	No cache	DualCache	ES-dLLM	CAI-dLLM
NarrativeQA	F1	18.18	16.96	17.05	16.96
MultifieldQA-en	F1	47.24	45.35	47.95	43.72
GovReport	ROUGE-L	19.99	17.23	16.26	17.09
QMSum	ROUGE-L	13.69	14.59	14.38	14.66
Average score		24.63	23.53	23.69	23.11
Average TPS		2.7	23.4	24.2	51.4

Appendix

Table 13: Detailed LongBench results on Dream-7B. Average score is the unweighted mean of per-task scores.

Task	Model	Method	Avg Power	Energy (J)	Energy (Wh)	Saved
GSM8K	LLaDA	No cache	688.0 W	12,812,627	3559.1	–
GSM8K	LLaDA	DualCache	587.2 W	950,104	263.9	92.6%
GSM8K	LLaDA	ES-dLLM	441.7 W	801,455	222.6	93.7%
GSM8K	LLaDA	CAI-dLLM	558.3 W	607,297	168.7	95.3%
HumanEval	Dream	No cache	673.4 W	1,347,290	374.25	–
HumanEval	Dream	DualCache	550.5 W	215,714	59.92	84.0%

Appendix

Table 14: Detailed energy measurements. The Saved column reports reduction relative to the no-cache row for the same task and model.

Model	Method	$L=64$		$L=128$		$L=256$		$L=512$
		TPS	Speedup	TPS	Speedup	TPS	Speedup	TPS	Speedup
LLaDA	No cache	52.3	1.00 $\times$	53.8	1.00 $\times$	54.0	1.00 $\times$	40.6	1.00 $\times$
	DualCache	51.3	0.98 $\times$	50.2	0.93 $\times$	50.8	0.94 $\times$	51.4	1.27 $\times$
	ES-dLLM	45.4	0.87 $\times$	46.9	0.87 $\times$	46.8	0.87 $\times$	46.9	1.16 $\times$
	CAI-dLLM	108.0	2.07 $\times$	202.1	3.76 $\times$	123.6	2.29 $\times$	186.3	4.59 $\times$
Dream	No cache	61.5	1.00 $\times$	63.1	1.00 $\times$	61.4	1.00 $\times$	45.3	1.00 $\times$

Appendix

Table 15: Throughput and speedup over No cache across generation lengths on LLaDA-Instruct and Dream-Instruct. Best TPS for each generation length is shown in bold.

Explore similar work

ES-dLLM: Efficient Inference for Diffusion Large Language Models by Early-Skipping
Mar 10, 2026Zijian Zhu, Fei Ren, Zhanhong Tan +1Diffusion Language Model Inference Diffusion Language Models
IIIS, Tsinghua University · Polar Bear Tech.
DyLLM: Efficient Diffusion LLM Inference via Saliency-based Token Selection and Partial Attention
Mar 9, 2026Younjoo Lee, Seungkyun Dan, Junghoo Lee +2Diffusion Language Model Inference Diffusion Language Model Decoding
Seoul National University, Seoul, Republic of Korea
Dynamic-dLLM: Dynamic Cache-Budget and Adaptive Parallel Decoding for Training-Free Acceleration of Diffusion LLM
May 27, 2026Tianyi Wu, Xiaoxi Sun, Yanhua Jiao +5Diffusion Language Model Inference Diffusion Language Model Decoding
*Equal contribution. The work is done during Tianyi’s internship at CBG Celia DeviceAI Team, Huawei. · 1Harbin Institute of Technology, Shenzhen · 2Huawei +1

Mar 10, 2026J/K move · Enter open · S save

ES-dLLM: Efficient Inference for Diffusion Large Language Models by Early-Skipping

Zijian Zhu, Fei Ren, Zhanhong Tan, Kaisheng Ma

IIIS, Tsinghua University · Polar Bear Tech.

Diffusion Language Model Inference Diffusion Language Models Diffusion Model Inference Acceleration Token Pruning

Open paper AskPDF arXiv

Diffusion large language models (dLLMs) are emerging as a promising alternative to autoregressive models (ARMs) due to their ability to capture bidirectional context and the potential for parallel generation. Despite the advantages, dLLM inference remains computationally expensive as the full input context is processed at every iteration. In this work, we analyze the generation dynamics of dLLMs and find that intermediate representations, including key, value, and hidden states, change only subtly across successive iterations. Leveraging this insight, we propose ES-dLLM, a training-free inference acceleration framework for dLLM that reduces computation by skipping tokens in early layers based on the estimated importance. Token importance is computed with intermediate tensor variation and confidence scores of previous iterations. Experiments on LLaDA-8B and Dream-7B demonstrate that ES-dLLM achieves throughput of up to 226.57 and 308.51 tokens per second (TPS), respectively, on an NVIDIA H200 GPU, delivering 5.6

\times

to 16.8

\times

speedup over the vanilla implementation and up to 1.85

\times

over the state-of-the-art caching method, while preserving generation quality.

Abstract

Explore similar work

ES-dLLM: Efficient Inference for Diffusion Large Language Models by Early-Skipping

DyLLM: Efficient Diffusion LLM Inference via Saliency-based Token Selection and Partial Attention

Dynamic-dLLM: Dynamic Cache-Budget and Adaptive Parallel Decoding for Training-Free Acceleration of Diffusion LLM