[论文] [论文] Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Fa...

论文概要

研究领域: LLM
作者: Quan Nguyen-Tri, Mukul Ranjan, Zhiqiang Shen
发布时间: 2026-09-22
arXiv: 2609.26796

中文摘要

扩散大语言模型（dLLM）近来成为自回归 LLM 的有前景替代，支持非自回归文本生成。然而其实际部署仍受制于低效推理，主因是缺乏有效的键值（KV）缓存与可扩展的并行解码机制。现有加速方法多将 KV 缓存与并行解码孤立研究，忽视了二者联合应用时的 I/O 瓶颈。本文提出 Flash-dLLM——免训练的 dLLM 推理加速框架，兼顾速度与内存效率。它首先识别出 GPU 内存 I/O 是启用 KV 缓存的 dLLM 推理的主导瓶颈，用 I/O 感知的融合 KV 缓存内核减少冗余内存搬运；在此之上提出 KV 缓存驱动的"草稿-验证"解码策略，dLLM 自身同时充当草稿器与验证器，无需辅助模型。统一设计在保持生成质量的同时更快解码，并提升对更长序列与更大 batch 的可扩展性。在数学推理与代码生成基准上，Flash-dLLM 稳定优于现有最先进加速方法：在 GSM8K 与 HumanEval 上分别取得 5.1× 和 11.0× 于此前最强基线 Elastic-Cache 的加速。

原文摘要

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.

自动采集于 2026-09-24

#论文 #arXiv #LLM #小凯