Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Fast, Memory-Efficient Diffusion LLMs
Quan Nguyen-Tri, Mukul Ranjan, Zhiqiang Shen
cs.CL
2026-09-23
Flash-dLLM 用融合 KV 内核减少显存搬运，再用模型自己做草稿–校验的并行解码。相对无缓存贪心，MBPP-512 加速 148.2 倍；相对 Elastic-Cache，GSM8K/HumanEval 分别快 5.1 倍和 11.0 倍。
扩散大语言模型按去噪步并行填掩码，理论上不该像自回归那样一个词一个词等。开源 dLLM 仍然慢，因为双向注意力和每步都在变的隐状态，让自回归那套 KV 缓存不能直接搬。已有加速把缓存和并行解码分开做：Fast-dLLM 在块边界整表刷新，Elastic-Cache 按注意力漂移自适应刷新，草稿–校验往往再外挂一个自回归模型。缓存复用和并行核验叠在一起时，真正卡住的是 GPU 显存读写，不是 FLOP。
Flash-dLLM 免训练，两块拼在一起。
Flash-Cache：常规实现每层把 QKV 投影、RoPE、写缓存、注意力拆成多次内核，中间张量在 HBM 里来回倒。融合核对查询集在 SRAM 里做投影和 RoPE，直接写进 KV 缓存。配合块表调度的 Flash Attention，让同一 batch 里「只算窗口」和「整段重算」的样本不用互相填充。再加选择性更新：中间层 top-32 被关注的已解码 token 能占总注意力约 50%，所以每步查询固定为滑动掩码窗 βm 加跟踪预算 βt，其余位置吃缓存。
Flash-Verify：先按置信度拆成高于 ε 直接提交的集合，和低于阈值的搜索集。第二次前向把搜索位置各放两份，一份填草稿、一份仍是 [MASK]，用因果掩码隔开，共享同一份 KV。两视图预测一致且掩码视图置信度高于 γ，才按从左到右接到第一处冲突为止。模型自己当草稿器和校验器，额外代价大约是 2βm，不是全序列。
两视图共享位置编码，但不能互相看见，所以校验不是「再看一眼同一张图」，而是同一上下文上草稿条件与掩码条件是否给出同一个词。默认超参：ε=0.9，γ=0.8，βt=80，βm=64，生成长度 512。主实验在单卡 A100 80GB 上跑 LLaDA-1.5。
相对无缓存贪心，Flash-Cache+Flash-Verify 在八个设置里都是最高吞吐，148.0–210.6 token/s，加速 22.3×–148.2×。GSM8K-512：83.02% 准确率、210.6 token/s、81.0×。相对当时最强基线 Elastic-Cache，GSM8K 快 5.1×（210.6 / 41.7），HumanEval 快 11.0×（185.6 / 16.8）。同一张更长的对照表里，FreeDave 是 42.8 token/s，Flash-Verify 配置是 210.6。
融合核在 RTX 3090 上单独有 1.37×。batch 32 仍近似线性，Fast-dLLM 在 24 就 OOM。batch 16 时显存约 26GB，对 Fast-dLLM 的 50GB 少约 48%。数学任务上组合配置离最高准确率不超过 1.78 点；256 token 代码任务会掉 3.6 点左右，吞吐优先时要认这笔账。
dLLM 的并行如果落不到 wall-clock 上，部署端还是会退回自回归。这篇把 IO 和并行解码当成同一件事：缓存内核先把带宽拿回来，自校验再把每步能安全提交的 token 翻倍，还不外挂模型。对已经在用 LLaDA/Dream 一类掩码扩散模型的人，这是可直接接的推理层，不用重训。
主表是 LLaDA-1.5，附录写还看了另一种掩码扩散模型，连续空间扩散没验证。基准是数学和代码，开放长文的置信度更平，阈值策略可能失效。γ 和 βm 全程固定。代码任务上最快配置不是最准配置。和 Llama3-8B 的吞吐曲线只作参照，不是同结构公平赛。