Publications

Enabling KV Caching of Shared Prefix for Diffusion Language Models

The 2026 Conference on Empirical Methods in Natural Language Processing (EMNLP 2026), Main Track, Acceptance rate 15.4%, Budapest, Hungary, 2026.

Younghun Go, Jaehoon Han, Changyong Shin, Chuck Yoo, Gyeongsik Yang

Download PDF

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

Paper summary

Reuse a prefix without freezing its context

Diffusion models can reuse parts of a shared prompt’s computation while allowing the changing context to update the rest.

A diffusion language model revises tokens through bidirectional attention. Changing a later token can therefore change the key/value tensors of an earlier shared prompt. Reusing that prompt’s cached tensors across every layer, as in autoregressive serving, can corrupt the context. BiCache instead asks which layers remain similar enough to reuse.

Partition caching by layer depth

Offline profiling maps the fraction of shared-prefix tokens to a reusable layer depth. At serving time, BiCache looks up that depth for each request and reuses prefix tensors across requests only in the selected shallow layers.

The prefix is shared, but its deeper interpretation still belongs to the current request.

Deeper layers compute their own tensors and refresh them periodically during denoising. This combines reuse across requests in shallow layers with reuse across steps within a request in deep layers.

Offline prefix profiling selects a reusable depth; serving combines cross-request shallow-layer reuse with deeper request-local caching.View full-size figure
Offline prefix profiling selects a reusable depth; serving combines cross-request shallow-layer reuse with deeper request-local caching.

Measure speed and answer quality together

The main tests use LLaDA on one NVIDIA B200, batch size 1, 256 output tokens and 128 denoising steps. Across ARC-Challenge-Chat, GPQA, MATH-500 and GSM8K, BiCache increases throughput by 36.3–82.8% over no caching; accuracy differs by at most 0.8 percentage points in this setting.

The authors also combine BiCache with Fast-dLLM, which caches a different part of the computation. Compared with Fast-dLLM alone, throughput rises by 51.6–98.3%; accuracy is unchanged on ARC-C, 0.2 and 0.8 points higher on GPQA and GSM8K, and 1.8 points lower on MATH. These results distinguish faster generation from changes in answer quality.

BiCache throughput gains on the evaluated LLaDA tasks, grouped by caching baseline.
ComparisonThroughput increase
BiCache versus no caching36.3–82.8%
BiCache + Fast-dLLM versus Fast-dLLM51.6–98.3%

Key takeaways

  1. Shared-prefix reuse depends on layer depth. Shallow-layer stability creates a caching opportunity whose depth depends on the prefix fraction.
  2. Reuse and freshness can coexist. Cross-request prefix caching is paired with request-local computation and periodic refresh deeper in the model.
  3. Prefix reuse can add to other diffusion-model optimizations. The Fast-dLLM experiment demonstrates a benefit from combining the caches.