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

It identifies redundant memory read/write of the KV cache as a dominant bottleneck in dLLM inference and proposes Flash-Cache: QKV projection, rotary positional embedding and cache writing are fused into one kernel, with keys and values produced in SRAM and written directly into the KV cache. Prior acceleration work typically studies KV caching and parallel decoding in isolation and implements cache updates as separate PyTorch kernel launches per layer, moving intermediate tensors to and from HBM repeatedly; here a fused kernel is combined with scheduled Flash Attention, which uses a block table to align query blocks with their key-value blocks and handles the length divergence caused by alternating caching and update stages within a batch. The paper provides a per-layer HBM traffic analysis (Fig. 3(a)) and kernel design description, evaluated on a single NVIDIA A100 80GB with LLaDA-1.5; the authors also report a speedup on an RTX 3090 GPU (that multiplier is absent from this parsed text).

Selective cache update: only a small set of most-attended tokens is maintained and updated, rather than reusing all cached states uniformly. The paper reports that in the middle layers (layers 5 to 20) only 32 top-attended tokens can contribute roughly 50% of the attention weight, and accordingly restricts each step's query to a sliding unmasking window plus a fixed tracking budget, bounding per-step computation by that budget. The observation is paired with controlled ablations over tracking budget and thresholds (accuracy-throughput trade-off, with standard deviations across five random seeds reported as at most a few percentage points for accuracy and a few tokens/s for throughput).

Flash-Verify: the dLLM itself acts as both drafter and verifier, with no auxiliary model and no additional training. Earlier draft-and-verify schemes rely on a separate autoregressive verifier or multiple independent forward passes; here the draft view and the mask view are placed at the same positions, share positional embeddings and are isolated by a causal attention mask inside the fused Triton kernel, and a position is accepted only when both views agree and the mask-view confidence exceeds a threshold, with acceptance proceeding in causality order and stopping at the first mismatch. It reports markedly more tokens accepted per step (about 5.6–5.7 versus 2.8), reaching 210.6 tokens/s and 83.02% accuracy on GSM8K-512; the appendix gives a theorem and corollary bounding deviation in total variation from the model's own chain-rule joint.

The approach targets inference serving for masked diffusion language models, especially structured-output tasks such as mathematical reasoning and code generation, and batched or longer-sequence settings (the tables report scaling to batch size 32 and lower GPU memory use than Fast-dLLM). For resource-constrained deployers, being training-free and needing no external drafter means it can be layered directly onto existing dLLM checkpoints; for systems researchers, the fused cache kernel, block-table scheduling and two-view verification are reusable engineering components, and the appendix's deviation bound supplies an analysis tool for verification-style parallel decoding. The paper states explicitly that validation covers masked diffusion models, with continuous-space diffusion language models and open-ended generation listed as future directions.

Several numeric values were stripped from the body text parsed here (default confidence threshold, verify threshold, block size, tracking budget, and parts of the memory and throughput multipliers), so the account rests on the abstract and the table numbers still visible; the exact hyperparameter settings and the fused-kernel speedup on RTX 3090 would need checking against the original. The experiments list LLaDA-1.5, while the appendix speaks of two representative masked diffusion LLMs, and readers may want to confirm how those correspond. In addition, speedups such as 5.1× and 11.0× are closely tied to baseline implementation and hardware configuration, so the magnitude bears watching when reproduced elsewhere; open-domain long-form generation, continuous-space diffusion language models, and schemes that adapt the thresholds to running confidence statistics all remain open questions.