提出首个适用于扩散语言模型的共享前缀缓存技术,解决动态更新导致的精度崩溃问题。
Enabling KV Caching of Shared Prefix for Diffusion Language Models

- 基于浅层共享前缀键值对稳定可复用的观察,动态确定安全缓存深度。
- 相比现有方法,吞吐量提升36.3%~98.3%,精度损失仅0-1.8%。
- 适合需要高并发推理的扩散语言模型服务场景,如生成式AI应用。
共享前缀的键值(KV)缓存对实现大语言模型(LLM)的高吞吐服务至关重要,但在新兴的扩散语言模型(DLMs)中面临严峻挑战。在DLM中,双向注意力机制意味着任何标记的动态更新都会改变整个上下文及其对应的键值对。因此,传统针对LLM设计的缓存技术(假设键值对一旦计算即不变)会破坏共享前缀的键值对。实验表明,将这些技术应用于DLM会导致模型精度骤降至接近零。为实现高效的DLM服务,我们提出双向前缀缓存(bicache),这是首个面向DLM中共享前缀的键值缓存技术。bicache基于全面分析的关键发现:共享前缀的键值对在浅层保持稳定且可复用,而浅层深度取决于每个请求中共享前缀标记的比例。因此,bicache动态识别可安全复用共享前缀键值对的层深,消除冗余计算。评估结果表明,与现有技术相比,bicache在不引发精度崩溃的前提下,显著提升服务吞吐量36.3%~98.3%(精度差异仅0-1.8%)。
原文摘要 · Abstract (English)
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).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。