Enabling KV Caching of Shared Prefix for Diffusion Language Models

📄 arXiv: 2606.07571 📥 PDF

作者: Younghun Go, Jaehoon Han, Changyong Shin, Chuck Yoo, Gyeongsik Yang

分类: cs.LG, cs.AI

发布日期: 2026-07-05

💡 一句话要点

提出双向前缀缓存技术以解决扩散语言模型的KV缓存问题

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

关键词: 扩散语言模型 KV缓存 双向前缀缓存 高吞吐量 自然语言处理 模型优化 性能提升

📋 核心要点

  1. 现有的KV缓存技术在扩散语言模型中面临挑战，因为动态更新令牌会影响整个上下文及其KV。
  2. 论文提出的双向前缀缓存技术bicache，能够动态识别安全的层深度以重用共享前缀KV，消除冗余计算。
  3. 实验结果表明，bicache相比现有技术提高了服务吞吐量36.3%-98.3%，且准确性保持在0-1.8%的差异内。

📝 摘要（中文）

共享前缀的关键值（KV）缓存对于高吞吐量的大型语言模型（LLM）服务至关重要，但在新兴的扩散语言模型（DLM）中面临重大挑战。在DLM中，双向注意力意味着动态更新任何令牌会改变整个上下文及其对应的KV。因此，现有的KV缓存技术假设一旦计算完成KV就保持不变，这会导致共享前缀KV的损坏。我们的实验表明，将这些技术应用于DLM会导致模型准确性几乎降至零。为了解锁高吞吐量的DLM服务，我们提出了双向前缀缓存技术bicache，这是DLM中共享前缀的首个KV缓存技术。bicache基于我们的综合分析得出的关键观察：共享前缀KV在浅层中保持稳定和可重用，而浅层的深度取决于每个请求中共享前缀令牌的比例。通过动态识别安全的层深度以重用共享前缀KV，bicache消除了冗余计算。评估结果表明，bicache相比现有技术显著提高了服务吞吐量，提升幅度为36.3%-98.3%，且没有准确性崩溃（仅0-1.8%的差异）。

🔬 方法详解

问题定义：论文要解决的问题是现有KV缓存技术在扩散语言模型中的不适用性，因其假设KV一旦计算完成便保持不变，导致共享前缀KV的损坏，影响模型准确性。

核心思路：论文的核心思路是提出双向前缀缓存技术（bicache），通过动态识别安全的层深度来重用共享前缀KV，从而避免冗余计算并提高吞吐量。

技术框架：整体架构包括对输入请求的分析，确定共享前缀的比例，动态调整KV缓存的层深度，以及实现高效的KV重用机制。主要模块包括请求解析、KV缓存管理和性能优化模块。

关键创新：最重要的技术创新点在于bicache能够在DLM中有效管理共享前缀KV的稳定性和可重用性，突破了现有技术的局限，避免了准确性崩溃。

关键设计：关键设计包括对共享前缀令牌比例的动态分析，安全层深度的计算，以及在浅层中保持KV的稳定性，确保在高吞吐量下的准确性。通过这些设计，bicache实现了显著的性能提升。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，bicache在服务吞吐量上相比现有技术提升了36.3%-98.3%，且在准确性方面仅有0-1.8%的差异，证明了其在高效KV缓存管理中的有效性和可靠性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和大规模文本生成等场景。通过提高扩散语言模型的服务效率，bicache可以在实时应用中显著提升用户体验，推动智能助手和自动化内容生成的进步。未来，该技术可能会影响更广泛的AI应用，提升其响应速度和准确性。

📄 摘要（原文）

Key-value (KV) caching for shared prefixes is essential for high-throughput large language model (LLM) serving, but it faces critical challenges in emerging diffusion language models (DLMs). In DLMs, bidirectional attention means that updating any token dynamically alters the entire context and its corresponding KVs. Thus, existing caching techniques developed for LLMs, which assume that KVs remain invariant once computed, corrupt the shared prefix KVs. Our experiments show that applying these techniques to DLMs causes model accuracy to collapse to near zero.To unlock high-throughput DLM serving, we propose bidirectional prefix caching, bicache, the first KV caching technique for shared prefixes in DLMs. bicache is designed based on key observations from our comprehensive analysis: shared prefix KVs remain stable and reusable in shallow layers, while the depth of shallow layers depends on the fraction of shared prefix tokens in each request. Thus, bicache dynamically identifies a safe layer depth for reusing shared prefix KVs and eliminates redundant computation. Evaluations demonstrate that bicache significantly improves serving throughput by 36.3%-98.3% compared to existing techniques without accuracy collapse (only 0-1.8% difference).