Prefilling-dLLM v2: Learned Chunk Routing for Long-Context Diffusion Language Models

Qi Han; Jing Xiong; Ziyuan Zhang; Ngai Wong

cs.CLOct 13, 2026

Prefilling-dLLM v2: Learned Chunk Routing for Long-Context Diffusion Language Models

Authors: Qi Han, Jing Xiong, Ziyuan Zhang, Ngai Wong

Organizations: University of Hong Kong · Zhejiang University

Abstract

Diffusion language models apply bidirectional attention over the whole context at every denoising step, so a longer context makes each step substantially more expensive. Sparsifying only at inference leaves the pretrained weights unchanged, so the model is never trained under the sparse pattern it runs with. We propose Prefilling-dLLM v2, which continues training a diffusion language model so that it routes with its native mask token. One mask is appended at the end of each chunk, and a learned role offset marks it as a routing slot. On every fourth layer, a query keeps a bidirectional local window and the top-K chunks scored from that mask; the other layers use a bidirectional sliding window. Starting from Dream-v0-Base-7B extended to 16K with YaRN, we compare Prefilling-dLLM v2 with bidirectional NSA, DSA, and SWA under the same layout, key budget, and dense initialization, using the same Fast-dLLM decoder. With the YaRN factor scaled to the evaluation length, Prefilling-dLLM v2 leads the other sparse methods on all three RULER tasks from 16K to 128K. Its mean accuracy at 128K is 7.4 points above NSA, the next sparse model, while DSA and SWA remain near zero. At the same length, decoding is 172x faster than with full dense attention.

Figures & tables

Figure 1: Mean accuracy of the four sparse methods on the three RULER tasks. The YaRN factor scales with the evaluation length.

Panel of Figure 2 — Figure 2: What query $q$ can see. Each chunk is three cells, and $q$ sits in the middle so the local window extends both ways. Teal cells are the local window and orange cells are selected. The bar above each NSA chunk is its summary. DSA reads scattered tokens, and Prefilling-dLLM v2 marks the last cell of every chunk as a mask.

Figure 3: Computation on a route layer. A chunk is $S-1$ content tokens followed by one routing mask. The role offset $r$ is added to that mask at the input. Its query pools the content keys with an intra-chunk softmax, and the weighted sum is the summary key $s_{c}$ . A content query adds a low-rank residual and scores these summaries. Inside a GQA group ( Ainslie et al., 2023 ) the scores are reduced by a maximum, and the top $K$ chunks are shared by every head in the group. Attention mass puts the selected chunk scores and the local-window log-sum-exp in one partition, then a second softmax distributes each chunk’s mass over its content tokens.

Dense
		16K				32K				64K				128K
Model	YaRN	S-N	MK	VT	Avg.	S-N	MK	VT	Avg.	S-N	MK	VT	Avg.	S-N	MK	VT	Avg.
Full dense	Fixed	100.0 ∗	100.0 ∗	91.7 ∗	97.2 ∗	100.0 ∗	96.5 ∗	75.7 ∗	90.7 ∗	94.0 ∗	68.5	40.7	67.7 ∗	85.0 ∗	90.0 ∗	34.1 ∗	69.7 ∗
Full dense	Scaled	100.0 ∗	100.0 ∗	91.7 ∗	97.2 ∗	100.0 ∗	99.5 ∗	80.0 ∗	93.2 ∗	100.0 ∗	97.5 ∗	69.3 ∗	88.9 ∗	100.0 ∗	100.0 ∗	62.3 ∗	87.4 ∗
Hybrid dense	Fixed	99.0	87.5	67.3	84.6	100.0 ∗	81.5	67.7	83.1	70.0	73.5 ∗	59.3 ∗	67.6	25.0	25.0	27.0	25.7
Hybrid dense	Scaled	99.0	87.5	67.3	84.6	99.0	80.0	69.0	82.7	98.0	81.5	62.7	80.7	85.0	83.5	53.7	74.1

Table 1: Three RULER tasks at 16K, 32K, 64K, and 128K. The upper block is dense attention. Full dense keeps dense attention in every layer through the same number of continued-training steps as the other variants. Hybrid dense uses three sliding-window layers and one dense layer. The lower block is sparse attention. Fixed uses YaRN factor 8. Scaled sets the factor to the evaluation length divided by 2048. MK is multi-key multi-query. Avg. is the mean of S-N, MK, and VT. In the dense block,

*

marks the best score in each column. In the sparse block, the best score is bold and the second best is underlined.

Model	Single-Doc. QA			Multi-Doc. QA			Summarization				Few-shot		Synthetic		Code		Avg.
Model	NarrativeQA	Qasper	MF-en	HotpotQA	2WikiMQA	Musique	GovReport	QMSum	MultiNews	SAMSum	TREC	TriviaQA	PCount	PRe	LCC	RB-P
Dense
Full dense	29.1 ∗	36.0	53.7 ∗	60.1 ∗	47.1 ∗	38.8	5.1	22.8	4.7 ∗	42.9 ∗	73.5 ∗	91.2 ∗	2.0 ∗	99.5 ∗	47.9 ∗	53.8		44.3 ∗
Hybrid dense	25.4	37.1 ∗	48.5	59.3	40.6	39.2 ∗	8.7 ∗	24.9 ∗	1.5	35.7	72.0	88.7	2.0 ∗	96.3	46.2	55.7 ∗	42.6
Sparse
SWA	18.0	20.8	26.4	28.3	28.1	12.8	19.8	20.7	13.7	41.9	64.0	86.2	1.7	14.5	45.8	50.8	30.9

Table 2: LongBench-v1 at the 16K training length. The upper block is dense attention. Full dense keeps dense attention in every layer through the same number of continued-training steps as the other variants. Hybrid dense uses three sliding-window layers and one dense layer. The lower block is sparse attention. In the dense block,

*

marks the best score in each column. In the sparse block, the best score is bold and the second best is underlined.

Figure 4: Cosine between the mean content key of each chunk and the summary key of each routing slot on the 16K single-needle probe, averaged over route layers and 100 examples. Each panel is binned to $32\times 32$ . On the clean probe the diagonal pairs a slot with its own chunk. Shuffle summaries permutes the summaries across chunks.

Figure 5: 16K single-needle accuracy for the routing mask, a new learned token, and EOS. Shuffle summaries permutes the summaries across chunks and leaves the text in place. Move needle places the intact needle in another chunk. Each bar is 100 examples.

Figure 6: Prefill time and decode time per generated token under the same Fast-dLLM decoder. Both vertical axes are logarithmic. The number above each full-dense and hybrid-dense bar is the ratio of that latency to the latency of Prefilling-dLLM v2. The YaRN factor is the evaluation length divided by 2048.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Variant	Layers	Local window	Remote budget	Total budget
Prefilling-dLLM v2	every 4th	radius 256	32 chunks of 64	$\approx 2560$
Prefilling-dLLM v2	other 21	radius 1280	–	$\approx 2560$
NSA	every 4th	radius 128	$32\times 64$ selected, one token per 64	$\approx 2560$
NSA	other 21	radius 1280	–	$\approx 2560$
DSA	every 4th	–	token top- $k$ 2560	2560
DSA	other 21	radius 1280	–	$\approx 2560$

Appendix

Table 3: Layer configurations. Each variant occupies two rows. The first row is every fourth layer, and the second row is the other 21 layers. Local window is a bidirectional radius. Remote budget is the set of keys gathered beyond that window. Total budget is the number of key positions on that row. The sparse rows share a budget of about 2560. A dense layer reads every key.

Hyperparameter	Value
GPUs	2
Steps	500
Text tokens per sequence	16128
YaRN factor	8
Optimizer	AdamW
Learning rate	$2\times 10^{-4}$

Appendix

Table 4: Dense context extension. The run continues Dream-v0-Base-7B under dense attention. The base weights stay frozen, and the listed low-rank adapters are the trainable parameters.

Variant	Hyperparameter	Value
NSA	Route layers	every 4th
	Compressed branch	block 64, stride 64
	Selected branch	32 blocks of 64
	Local branch	radius 128
	Gate	three sigmoids per head
	Other layers	SWA radius 1280

Appendix

Table 5: Continued training. Each variant is initialized from the dense step-500 checkpoint and trains for 1000 steps with a cosine schedule on four GPUs, at learning rate

10^{-4}

with 50 warmup steps. The base weights stay frozen. Low-rank adapters of rank 16, scale 32, and dropout 0.05 train the query, key, value, and output projections. The rows below are the parameters that differ across variants.

Figure 7: Training loss of the continued-training runs. Each line averages the previous 25 steps. The dashed vertical line marks the end of the DSA indexer warmup.

	16K				32K				64K				128K
Variant	S-N	MK	VT	Avg.	S-N	MK	VT	Avg.	S-N	MK	VT	Avg.	S-N	MK	VT	Avg.
Prefilling-dLLM v2	100.0	42.0	19.3	53.8	99.0	35.0	18.0	50.7	100.0	27.0	17.0	48.0	99.0	32.5	8.0	46.5
No residual	63.0	12.0	3.0	26.0	47.0	5.5	4.7	19.1	27.0	5.5	6.3	12.9	16.0	7.5	3.7	9.1
No role offset	95.0	13.5	13.0	40.5	94.0	9.5	15.7	39.7	93.0	13.0	18.0	41.3	85.0	13.5	13.3	37.3
Mean summary	99.0	50.5	20.0	56.5	100.0	51.0	18.7	56.6	96.0	39.0	21.3	52.1	62.0	29.5	16.3	35.9
No entropy	98.0	13.0	11.0	40.7	97.0	7.0	12.3	38.8	90.0	10.0	12.0	37.3	58.0	12.5	11.3	27.3

Appendix

Table 6: Component ablations on the three RULER tasks. The YaRN factor is held at 8 at every length. MK is multi-key multi-query. Avg. is the mean of S-N, MK, and VT. Prefilling-dLLM v2 is the full model, copied from its Fixed row in Table 1 . Each other row removes one component.

	16K				32K				64K				128K
Variant	S-N	MK	VT	Avg.	S-N	MK	VT	Avg.	S-N	MK	VT	Avg.	S-N	MK	VT	Avg.
Mask	100.0	42.0	19.3	53.8	99.0	35.0	18.0	50.7	100.0	27.0	17.0	48.0	99.0	32.5	8.0	46.5
EOS	98.0	10.5	4.0	37.5	97.0	4.5	7.0	36.2	81.0	7.0	7.3	31.8	57.0	5.0	6.3	22.8
New token	93.0	7.5	16.3	38.9	94.0	4.0	15.7	37.9	93.0	7.0	19.7	39.9	72.0	5.0	15.0	30.7

Appendix

Table 7: Routing-token ablations on the three RULER tasks. The YaRN factor is held at 8 at every length. MK is multi-key multi-query. Avg. is the mean of S-N, MK, and VT. Mask is the native mask used by Prefilling-dLLM v2, copied from its Fixed row in Table 1 . EOS and New token replace the mask in the routing slot, both in training and at evaluation.

Diffusion Language Model Inference

Diffusion language model inference studies how to generate, guide, and serve text efficiently through iterative denoising and parallel token updates.

See all 179 papers

Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models
Jun 9, 2026Jing Xiong, Qi Han, Shansan Gong +5Lost-in-the-Middle Effect Prefill-Decode Disaggregation
The University of Hong Kong · University of Michigan, Ann Arbor · LMSYS Org
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
RouteSparse: Input-Conditional Pattern Routing for Budgeted Long-Context Prefilling
Aug 29, 2026Chao Zhang, Yifan Ji, Ziyan Zhang +2Long-Context Language Model Inference LLM Inference Acceleration
College of Computer Science, Chongqing University · School of Information Science and Engineering, Chongqing Jiaotong University

Jun 9, 2026J/K move · Enter open · S save

Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models

Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh+4

The University of Hong Kong · University of Michigan, Ann Arbor · LMSYS Org

Lost-in-the-Middle Effect Prefill-Decode Disaggregation Long-Context Language Model Inference Diffusion Language Model Inference Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

Diffusion large language models (dLLMs) re-encode the entire prefix at every denoising step, causing recomputation that scales quadratically with context length and becomes prohibitive for long-context scenarios. We propose Prefilling-dLLM, a training-free prefill-decode disaggregation framework for dLLMs that partitions the prefix into N chunks, caches their KV representations once, and selects the top-K most relevant chunks with intra-chunk token sparsity for decoding, showing that sparse prefilling can outperform dense attention while reducing per-step complexity from quadratic in the full sequence length to quadratic only in the decode length. On LongBench and InfiniteBench, Prefilling-dLLM achieves state-of-the-art quality among dLLM acceleration methods, and an attention kernel that parallelizes decoding over the non-contiguously cached chunk KV yields 9.1--28.0x speedup at 8K--32K contexts. We further show that beginning-of-sequence tokens prepended to each chunk act as periodic attention anchors that eliminate the lost-in-the-middle phenomenon. Code is available at https://github.com/menik1126/Prefilling-dLLM.

Abstract

Diffusion Language Model Inference

Explore similar work

Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models

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

RouteSparse: Input-Conditional Pattern Routing for Budgeted Long-Context Prefilling