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.
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 sc . 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.
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.
Dense
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 ∗
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
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.
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×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
≈2560
other 21
radius 1280
–
≈2560
NSA
every 4th
radius 128
32×64 selected, one token per 64
≈2560
other 21
radius 1280
–
≈2560
DSA
every 4th
–
token top- k 2560
2560
other 21
radius 1280
–
≈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×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 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.
Jing Xiong, Qi Han, Shansan Gong +5
The University of Hong Kong · University of Michigan, Ann Arbor · LMSYS Org
Masked diffusion language models enable parallel token decoding, providing a promising alternative to the sequential nature of autoregressive generation. However, their iterative denoising process remains computationally expensive because it repeatedly processes the entire sequence at every step. We observe that across these diffusion steps, most token representations remain stable; only a small subset, which we term salient tokens, contributes meaningfully to the next update. Leveraging this temporal sparsity, we present DyLLM, a training-free inference framework that accelerates decoding by selectively computing only these salient tokens. DyLLM identifies saliency by measuring the cosine similarity of attention contexts between adjacent denoising steps. It recomputes feed-forward and attention operations only for salient tokens while reusing cached activations for the remainder. Across diverse reasoning and code-generation benchmarks, DyLLM achieves up to 9.6x higher throughput while largely preserving the baseline accuracy of representative open-source diffusion LLMs, LLaDA, and Dream.
Younjoo Lee, Seungkyun Dan, Junghoo Lee +2
Seoul National University, Seoul, Republic of Korea
Dynamic sparse attention can reduce the quadratic cost of long-context prefilling without changing model weights. MInference assigns each attention head one pattern offline and estimates that pattern's sparse indices for every prompt. This design is efficient, but it assumes that a head's preferred pattern and sparsity budget remain suitable across inputs. We introduce RouteSparse, which routes each head and prompt segment among a small library of GPU-efficient sparse patterns. A low-cost probe estimates pattern utility and uncertainty; a latency-aware router then selects a pattern and budget, while uncertain cases fall back to a denser mask. We formulate routing as constrained risk minimization, derive an attention-output error certificate from omitted probability mass, and evaluate the method on long-context retrieval, question answering, summarization, and language modeling. On Llama 3.1-8B-Instruct with 128K-token prompts, RouteSparse achieves 6.5× dense prefill speed with a 0.2-point RULER drop relative to dense attention, compared with 7.3× speed and a 1.6-point drop for fixed per-head routing. Ablations confirm that input-conditional routing, hardware profiling, and selective dense fallback each contribute to the quality--latency tradeoff.
Chao Zhang, Yifan Ji, Ziyan Zhang +2
College of Computer Science, Chongqing University · School of Information Science and Engineering, Chongqing Jiaotong University