Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Fast, Memory-Efficient Diffusion LLMs

Flash-Verify: Cache-Based Draft-Verification Parallel Decoding

Conventional confidence-based decoding accepts only tokens with $c_i \ge \epsilon$ and discards the rest (source: §2.3). On uncertain tasks, per-step acceptance hits a ceiling.

Flash-Verify queries the full KV cache for $Q_t = T_t \cup M_t^{\beta_m}$ in the draft pass and sorts masked positions by confidence into a confident set $D_t$ and an exploratory set $S_t$ (source: Alg. 2). The verification pass builds new queries from three groups: an adjusted tracked set $T_v$, a draft view that fills $S_t$ with draft predictions, and a mask view that fills the same positions with [MASK]. The two views share position embeddings but cannot see each other through a causal mask (source: §2.3, Fig. 2).

The acceptance rule is simple:

$\hat{x}_i$ is the draft prediction, and $\tilde{x}_i$ and $\tilde{c}_i$ are the mask-view prediction and confidence. It accepts sequentially in decoding order only up to the first mismatch (source: §2.3). The extra cost scales with $2\beta_m$, not the full sequence, and no external model is needed.

There is also a theoretical guarantee (source: Appx. C). Let $p_i$ be the mask-view verification conditional distribution and $\pi_i$ the committed point mass; for each accepted token:

For a block $A = \{i_1,\dots,i_m\}$, $\text{TV}(\pi_A, p_A) = 1 - \prod_k \tilde{c}_{i_k} \le 1 - \gamma^m \le m(1-\gamma)$. Because the verification mask is causal, $p_A$ is the model’s true joint, not a product of marginals, and verification tokens introduce no parallel-unmasking dependency error (source: Appx. C, Remark 1). Lowering $\epsilon$ is not equivalent to verification, because $\epsilon$ splits the draft marginals in the all-masked state while $\gamma$ splits the distribution that is actually committed (source: Appx. C, Remark 2).

How It Works: A Concrete Example

Consider a mini example for a graduate student. The vocabulary contains only {A, B, [MASK]}, with generation length $N=4$, $\beta_m=2$, $\beta_t=1$, $\epsilon=0.9$, and $\gamma=0.8$ (variable definitions borrowed from §2.2-§2.3).

Let the input be $x_0 =$ [BOS] [MASK] [MASK] [MASK] [MASK], $D_{<1} = \{1\}$, and $M_1 = \{2,3,4,5\}$. Here $Q$ is the query positions, $K$ and $V$ are the cache, $D_t$ is newly decoded positions, $M_t^{\beta_m}$ is the front $\beta_m$ masks, $T_t$ is the tracked set, and $S_t$ is verification targets.

Step 1: Draft. Since $M_1^{\beta_m}=\{2,3\}$ and $T_1=\emptyset$, $Q_1=\{2,3\}$. The fused kernel reads only $x[Q_1]$ to build $q,k,v$ in SRAM and writes them directly to cache (source: Alg. 1). After attending over the full cache, suppose the model outputs $q_2=[0.92,0.05,0.03]$ and $q_3=[0.70,0.20,0.10]$. Confidences are $c_2=0.92$ and $c_3=0.70$. With $\epsilon=0.9$, $D_1=\{2: \hat{x}_2=A\}$ and $S_1=\{3: \hat{x}_3=A\}$.

Step 1: Verification. The verification query has the form [track, D_1=A, S_1 draft=A, S_1 mask=[MASK]]. Positions 3 in the draft view and mask view are masked so they cannot see each other. If the mask view outputs $p_3=[0.85,0.10,0.05]$, then $\tilde{x}_3=A$ and $\tilde{c}_3=0.85$. Since $\hat{x}_3=\tilde{x}_3$ and $0.85 \ge \gamma$, position 3 is also accepted. Under the prior approach, this token would have been discarded because of 0.70.

  flowchart LR
  Q["Q = T + Mβm<br/>e.g. {2,3}"] --> F["FusedForward<br/>Flash-Cache"]
  F --> D["D: c ≥ ε<br/>Accept immediately"]
  F --> S["S: c < ε<br/>Verification candidates"]
  S --> V["Verify: draft view vs mask view<br/>Causal mask"]
  V --> A["Agreement + c̃ ≥ γ<br/>Accept up to first mismatch"]

Step 1: Cache selection. Compute $a_i^1$ and pick the one previously decoded token most attended by the masked queries $\{2,3\}$. The newly resolved $\{2,3\}$ are automatically included before ranking, so $T_2$ stays at size 1. The next query is $Q_2 = M_2^{\beta_m} \cup T_2$.

This figure shows why this example works. As in the left panel, middle layers concentrate on a few tokens, and as in the right panel, higher confidence increases per-step acceptance (source: Fig. 1b-c).

The second figure corresponds to the actual verification pipeline from draft sorting through the $\epsilon$ and $\gamma$ threshold decisions (source: Fig. 2).

Tearing Down 1 Secret Weapon

If forced to choose between the fused kernel and self-verification, choose self-verification. The kernel trims constants, while verification changes the slope (source: §3.3).

The mechanism is clear. Confidence-based decoding discards a token from the marginal $q_i$ even when it is correct if it does not pass $\epsilon$ (source: Appx. C). Flash-Verify asks again under $p_{i_k}$ conditioned on earlier drafts. As contextual evidence accumulates, $\tilde{c}$ rises and the $r=0.48$ correlation turns into throughput (source: Fig. 1c, Fig. 6). At $\gamma=0.60$, it decodes 7.2 tokens/step with throughput similar to confidence-based decoding at 5.6 tokens/step but 3.5pp higher accuracy. As tokens/iteration grow, the accuracy gap narrows while the throughput advantage widens up to 1.33× (source: §3.3, Fig. 6).

Performance Validation: Key Results

Experimental Setup

The default is a single NVIDIA A100 80GB GPU, Triton 2.0 fused kernels, and LLaDA-1.5 evaluation (source: §3.1). Benchmarks are GSM8K 5-shot, MATH 4-shot, HumanEval 0-shot, and MBPP 3-shot, with generation lengths of 256 tokens and 512 tokens (source: §3.1, Appx. D). Default hyperparameters are $\epsilon=0.9$, $\gamma=0.8$, block size $\beta=16$, $\beta_t=80$, and $\beta_m=64$ (source: §3.1). Throughput is mean tokens/s up to the stop token under lm-eval-harness (source: Appx. D). Tokenization and the pretraining corpus follow the LLaDA-1.5 base model with no additional training or fine-tuning. The fine-tuning strategy is not applicable (source: §3.1, Abstract).

The figure above summarizes the batch-scaling experiment. Throughput rises nearly linearly up to batch size 32, Fast-dLLM OOMs at 24, and memory is lowest across the full range (source: Fig. 4).

Key Performance Metrics

• GSM8K-512: Flash-Cache + Flash-Verify ranks first in both accuracy and throughput at 83.02% accuracy and 210.6 tokens/s. That is 81.0× over greedy no-cache at 2.6 tokens/s (source: Tab. 1).
• MATH-512: 35.98%, 210.1 tokens/s, 42.0×. The highest accuracy is 37.76% from no-cache Flash-Verify (source: Tab. 1).
• HumanEval-512: 40.24%, 185.6 tokens/s, 58.0×. The highest accuracy is 42.07% from confidence-based Flash-Cache (source: Tab. 1).
• MBPP-512: 39.00%, 148.2 tokens/s, 148.2×. Greedy Flash-Cache leads in accuracy at 39.80% (source: Tab. 1).
• Scaling comparison: against dKV-Cache at 14.9 tokens/s, Fast-dLLM at 36.8 tokens/s, FreeDave at 42.8 tokens/s, and Elastic-Cache at 41.7 tokens/s, Flash-Cache + Flash-Verify reaches 210.6 tokens/s (source: Tab. 2).
• Memory: at batch size 16, about 26 GB versus about 50 GB for Fast-dLLM, about a 48% reduction. Flat preallocated cache removes padding overhead (source: §3.3, Fig. 4b).
• Masked window $\beta_m$: growing it from 16 to 96 lowers throughput, but Flash-Verify keeps a 1.4× to 1.5× lead over confidence-based decoding, with an accuracy gap of 0.2pp to 1.2pp (source: §3.3, Fig. 7).
• Generation length and prefill: both methods peak at 256 tokens and fall by about 10.5% to 12.8% at 1024 tokens. From 1-shot to 8-shot, the Flash-Verify advantage shrinks from 1.50× to 1.35×, but the absolute lead grows from 42.6 tokens/s to 78.6 tokens/s (source: Appx. F, Tab. 5).

Critical Comparison

The strongest comparison is the same-hardware rerun against Elastic-Cache and FreeDave (source: §3.1, Tab. 2). Elastic-Cache had the best baseline accuracy at 82.79% and FreeDave the best throughput at 42.8 tokens/s, but Flash-Cache + confidence-based decoding beats both at 82.87% and 149.4 tokens/s, and adding Flash-Verify pushes further to 83.02% and 210.6 tokens/s.

Conversely, in the 256-token code setting, the throughput leader is not the accuracy leader. On HumanEval-256, 39.63% versus a best of 43.29% leaves a 3.66pp gap, and on MBPP-256, 38.20% versus a best of 41.80% leaves a 3.60pp gap (source: §3.2). On math reasoning the gap is within 1.78pp, but on short code generation there is a region where aggressive parallel commits do not replace strict single-token refinement. Confidence-based decoding peaks slightly higher at about 83.4% accuracy but drops to about 131 tokens/s throughput (source: §3.3, Fig. 5b). In other words, the Pareto frontier lies farther out on the Flash-Verify side, but conservative decoding remains for a single extreme-accuracy point.

System-Level Metrics

The primary metrics are throughput in tokens/s and peak memory in GB, not latency (source: §3.3). From batch size 1 to 32, the greedy method goes from 18.5 tokens/s to 55.0 tokens/s for 2.97×, confidence-based decoding from 51.7 tokens/s to 139.5 tokens/s for 2.70×, and Flash-Verify from 56.0 tokens/s to 199.8 tokens/s for 3.57× (source: Appx. G, Tab. 6). Training FLOPs and $/1M tokens costs are not reported, nor are energy metrics. The reproduction checklist includes code release, Triton 2.0 specification, A100 80GB specification, and standard deviations over 5 seeds, but driver versions and exact full prompts rely on the appendix evaluation scripts (source: §3.1, Appx. D-E).

Our Take: Strengths, Limitations, and Why This Matters

Strengths. The authors’ claims are consistent. Cache reuse and parallel verification are not optimized separately but tied together on the same fused kernel and preallocated cache (source: §2.3). Having no external verifier matters for deployment. Unlike FlashDLM, which uses an external autoregressive verifier, and FreeDave, which uses two independent passes, there is no extra model or training (source: §4). It is also rare for a theory appendix to close a simple intuition with TV-distance guarantees (source: Appx. C).

Explicit limitations. The authors acknowledge evaluation focused on two masked diffusion LLMs and math/code structured outputs (source: Appx. A). Continuous-space diffusion, flat confidence distributions such as long-form writing or dialogue, and fixed $\gamma$ and $\beta_m$ remain unverified.

Potential limitations. First, a strong assumption. The guarantee bounds deviation from the model’s own sequential distribution, not accuracy against the data distribution (source: Appx. C, Assumption 1). Second, the parallel error of $D_t$ still relies on the confidence-based theorem $[46, Thm. 1]$. Third, cost transparency is weak. Training compute is not applicable, but inference cost and power are missing. Fourth, societal impact inherits harmful/biased generation from the base dLLMs unchanged and does not alter safety filters (source: Appx. B).

The reason it still matters is that it reframes the dLLM deployment bottleneck from FLOPs to I/O and scheduling. Holding throughput without OOM up to batch size 32 while cutting memory nearly in half can be used directly in serving-system design (source: Fig. 4, Tab. 6).

What’s Next?: Future Directions

The authors propose adaptive $\gamma$ and $\beta_m$, open-ended generation verification, and continuous-space diffusion as extensions (source: Appx. A). Three reasonable next steps are:

• Runtime confidence-statistics scheduler: instead of fixed thresholds, attach a controller that shrinks $\beta_m$ and raises $\gamma$ when prefixes are long or uncertainty is high. This can mitigate the relative slowdown from 1-shot to 8-shot in the appendix (source: Appx. F).
• Long-context and batched-serving integration: connect the block-table design to a PagedAttention-style serving stack and test pushing the saturation point beyond batch size 16. Batch size 16 already delivers about 93.2% of batch size 32, so the efficiency point is clear (source: Appx. G).
• Separate accuracy-sensitive code tasks: to close the roughly 3pp gap on 256-token code, test hybrid commits that conservatively revert only positions with high verification failure rates (source: §3.2).

Tables from the paper

Tables converted mechanically from the arXiv e-print LaTeX source. The numbers are the paper’s own and did not pass through a model.

Table 1. Accuracy and decoding efficiency of LLaDA-1.5 across different benchmarks and decoding configurations. Each cell reports accuracy (top) and throughput with speedup over greedy decoding without caching (bottom; blueblue: tokens/s, orangeorange: speedup). Bold indicates the highest accuracy in each row, while yellow!20yellow shading indicates the highest throughput.

Table 2. Comparison of accuracy and decoding throughput across different methods.

Table 3. The hyper-parameters of Flash-dLLM under various settings.

Table 4. Mean accuracy (%) and throughput over five random seeds. Accuracy is shown on the first line and throughput is shown in blue on the second line. Values are mean $\pm$ standard deviation.

Table 5. Ablation of generation and prefill lengths.

Table 6. Effect of batch size on throughput. Using LLaDA-1.5 for GSM8K-512 (5-shot). Speedup is measured relative to Flash-Cache with greedy decoding at the same batch size.

Table 7. Example LLaDA-1.5 (without cache) — A qualitative comparison of different decoding methods.

Table 8. Example Flash-dLLM: A qualitative comparison of different decoding methods, LLaDA-1.5

Figures in this post are taken from the original arXiv:2609.26796 (CC BY 4.0). Only size and format were changed.