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

📄 arXiv: 2609.26796v1 📥 PDF

作者: Quan Nguyen-Tri, Mukul Ranjan, Zhiqiang Shen

分类: cs.CL

发布日期: 2026-09-22

备注: Code available at: https://github.com/VILA-Lab/Flash-dLLM

💡 一句话要点

提出Flash-dLLM以解决扩散大语言模型推理效率问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散大语言模型 推理加速 KV缓存 并行解码 内存效率 草拟-验证策略 深度学习

📋 核心要点

  1. 现有的扩散大语言模型在推理过程中效率低下，主要由于缺乏有效的KV缓存和并行解码机制。
  2. 本文提出Flash-dLLM，通过I/O感知的KV缓存内核和草拟-验证解码策略，优化了推理过程。
  3. 实验结果显示，Flash-dLLM在多个基准测试中显著提升了推理速度和内存效率，超越了现有方法。

📝 摘要（中文）

扩散大语言模型（dLLMs）作为自回归语言模型的有力替代，面临推理效率低下的问题，主要源于缺乏有效的键值（KV）缓存和可扩展的并行解码机制。现有加速方法通常孤立地研究KV缓存和并行解码，忽视了在联合应用时出现的I/O瓶颈。本文提出Flash-dLLM，一个无训练推理加速框架，通过I/O感知的融合KV缓存内核和高效的KV缓存驱动的草拟-验证解码策略，显著提升了dLLM的推理速度和内存效率。实验结果表明，Flash-dLLM在数学推理和代码生成基准上均优于现有最先进的dLLM加速方法，尤其在GSM8K和HumanEval上分别实现了5.1倍和11.0倍的速度提升。

🔬 方法详解

问题定义：本文旨在解决扩散大语言模型推理过程中的效率低下问题，现有方法在KV缓存和并行解码的联合应用中存在I/O瓶颈，导致推理速度缓慢和内存使用不当。

核心思路：Flash-dLLM通过引入I/O感知的KV缓存内核，减少冗余内存移动，同时采用草拟-验证解码策略，使得模型在不需要辅助模型的情况下，既能草拟又能验证生成的文本，从而提升推理效率。

技术框架：Flash-dLLM的整体架构包括两个主要模块：I/O感知的KV缓存内核和草拟-验证解码策略。首先，优化的KV缓存内核减少了内存I/O瓶颈；其次，草拟-验证策略利用dLLM自身进行生成和验证，提升了整体解码速度。

关键创新：最重要的创新在于将KV缓存和并行解码结合，并通过I/O感知的设计来优化内存使用，显著提升了推理速度和内存效率。这一设计与现有方法的本质区别在于其综合考虑了I/O瓶颈。

关键设计：在实现中，采用了高效的内存管理策略，设置了适当的缓存大小，并设计了适合的损失函数以优化草拟和验证过程，确保生成质量与速度的平衡。具体参数设置和网络结构细节在实验部分进行了详细描述。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，Flash-dLLM在GSM8K和HumanEval基准测试中分别实现了5.1倍和11.0倍的速度提升，显著优于现有最强基线Elastic-Cache。这一成果不仅提升了推理速度，还改善了内存效率，展示了其在实际应用中的巨大潜力。

🎯 应用场景

Flash-dLLM的研究成果在多个领域具有广泛的应用潜力，包括自然语言处理、代码生成和数学推理等。其高效的推理机制可以被集成到各种智能应用中，提升用户体验和系统响应速度。未来，随着模型规模的不断扩大，该方法有望在更复杂的任务中发挥重要作用。

📄 摘要（原文）

Diffusion Large Language Models (dLLMs) have recently emerged as a promising alternative to autoregressive LLMs by enabling non-autoregressive text generation. However, their practical deployment remains limited by inefficient inference, largely due to the absence of effective Key-Value (KV) caching and scalable parallel decoding mechanisms. Existing acceleration methods typically study KV caching and parallel decoding in isolation, overlooking the I/O bottlenecks that arise when cache reuse and parallel token verification are jointly applied. In this work, we introduce $\textbf{Flash-dLLM}$, a training-free inference acceleration framework for fast and memory-efficient dLLMs. Flash-dLLM first identifies GPU memory I/O as a dominant bottleneck in KV-cache-enabled dLLM inference and addresses it with an I/O-aware fused KV-cache kernel that reduces redundant memory movement. Building on this optimized cache mechanism, Flash-dLLM further proposes an efficient KV-cache-driven draft-and-verify decoding strategy, where the dLLM itself serves as both drafter and verifier without requiring an auxiliary model. This unified design enables faster decoding while preserving generation quality and improving scalability to longer sequences and larger batch size. Extensive experiments on mathematical reasoning and code-generation benchmarks demonstrate that Flash-dLLM consistently outperforms existing state-of-the-art dLLM acceleration methods in both inference speed and memory efficiency. In particular, it achieves $5.1\times$ and $11.0\times$ speedups over prior strongest baseline Elastic-Cache on GSM8K and HumanEval, respectively.