Enabling KV Caching of Shared Prefix for Diffusion Language Models

2026-05-26

AI Summary

Existing KV caching methods suffer from degraded accuracy in diffusion language models due to the bidirectional attention mechanism, which invalidates shared prefix caching. This work proposes BiCache, the first approach enabling efficient reuse of KV caches for shared prefixes in diffusion models. Through empirical analysis, the authors observe that shallow-layer KV states remain stable and thus amenable to reuse. BiCache leverages this insight to dynamically determine a safe depth for cache reuse and adaptively adjusts its caching strategy based on the proportion of shared prefixes across requests. Compared to existing methods, BiCache achieves a 36.3%–98.3% improvement in inference throughput with negligible accuracy loss (0–1.8%), thereby overcoming the limitations imposed by conventional unidirectional model caching assumptions.

Abstract

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).

Problem

• KV caching
• shared prefix
• diffusion language models
• bidirectional attention
• model serving

Innovation

• diffusion language models
• KV caching
• bidirectional attention
• prefix sharing
• high-throughput serving