Flash-dLLM：面向扩散语言模型的 I/O 感知 KV 缓存与并行解码加速框架
核心概要
Flash-dLLM 是一个免训练的扩散语言模型推理加速框架：用 I/O 感知的融合 KV 缓存内核（把 QKV 投影、RoPE 与缓存写入合并并在 SRAM 内完成、直接写回 KV 缓存）和调度式 Flash Attention 缓解 GPU 显存读写瓶颈，用选择性缓存更新只保留最受关注的少量 token，并让模型自身同时充当草稿者与验证者（Flash-Verify）；在 LLaDA-1.5 上的 GSM8K、MATH、HumanEval、MBPP 实验中，它报告相对最强基线 Elastic-Cache 在 GSM8K 与 HumanEval 上分别 5.1× 和 11.0× 的加速，并保持生成质量。
深度剖析
把 KV 缓存的冗余显存读写识别为扩散语言模型推理的主要瓶颈，并提出 Flash-Cache：将 QKV 投影、旋转位置编码与缓存写入融合为一个内核，键值在 SRAM 内产生后直接写入 KV 缓存。 已有加速方法多把 KV 缓存与并行解码分开研究，且缓存更新以 PyTorch 逐层多次 kernel 启动实现，中间张量反复往返 HBM；这里改为融合内核，并配合调度式 Flash Attention，用块表把查询块与键值块对齐，以应对缓存阶段与更新阶段交替造成的批内序列长度差异。 文中给出逐层 HBM 流量分析（Fig. 3(a)）与内核设计说明，并在单张 NVIDIA A100 80GB 上以 LLaDA-1.5 评测；作者称在 RTX 3090 上也取得了加速（该倍数在本次解析文本中缺失）。
选择性缓存更新：只维护并更新一小组最受关注的 token，而不是均匀复用全部缓存状态。 文中报告在中间层（第 5 至 20 层）仅 32 个最受关注 token 即可贡献约 50% 的注意力权重，据此把每步查询限定为滑动解掩码窗口与固定跟踪预算内的 token，每步计算量被预算约束住。 该观察配合跟踪预算与阈值的受控消融（精度—吞吐权衡，并报告五次随机种子的标准差，称精度差异至多数个百分点、吞吐差异至多若干 tokens/s）。
Flash-Verify：扩散语言模型自身同时充当草稿者与验证者，无需外部模型、无需额外训练。 此前的草稿—验证方案依赖单独的自回归验证器或多次独立前向；这里把草稿预测视图与掩码视图放在同一位置，共享受位置嵌入并由融合 Triton 内核内的因果注意力掩码隔开，位置只有在两个视图给出相同预测且掩码视图置信度超过阈值时才被接受，接受按因果顺序进行并在首个不一致处停止。 报告每步接受 token 数明显增加（约 5.6–5.7 对比 2.8），GSM8K-512 上达到 210.6 tokens/s 与 83.02% 精度；附录给出与模型自身链式联合分布之间总变差有界偏离的定理与推论。
启示与展望
这套方法面向掩码式扩散语言模型的推理服务，尤其是数学推理与代码生成这类结构化输出任务，以及批量推理和较长序列的场景（表中给出到批大小 32 的扩展结果，并报告比 Fast-dLLM 更低显存占用）。对资源有限的部署者，免训练且不需要外部草稿模型意味着可以直接叠加在已有 dLLM 检查点上；对系统研究者，融合缓存内核、块表调度与两视图验证是可复用的工程构件，附录的有界偏离保证也为验证式并行解码提供了分析工具。文中明确说明验证范围是掩码扩散模型，连续空间的扩散语言模型与开放式生成被列为后续方向。
本次解析的正文有若干数值被剥离（默认置信度阈值、验证阈值、块大小、跟踪预算、以及部分内存与吞吐倍数），只能依据摘要与仍可见的表格数字判断，具体超参设定与 RTX 3090 上的融合内核倍数需回原文核对；实验部分列出的是 LLaDA-1.5，而附录提到“两个代表性掩码扩散模型”，二者之间的对应关系值得读者自己确认。此外，5.1×/11.0× 这类加速比与基线实现和硬件配置密切相关，跨环境复现时的量级值得留意；开放域长文生成、连续空间扩散语言模型，以及随运行置信度自适应调整阈值的方案，都还是开放问题。