EntropyCache: Decoded Token Entropy Guided KV Caching for Diffusion Language Models

TL;DR

EntropyCache is a novel, training-free method for KV caching in diffusion language models that uses token entropy to efficiently decide when to recompute, significantly speeding up inference with minimal overhead.

Contribution

It introduces a new entropy-based, constant-cost decision mechanism for KV cache updates that is independent of context length and model size.

Findings

Achieves 15.2x-26.4x speedup on standard benchmarks.

Maintains competitive accuracy with minimal decision overhead.

Decision process accounts for only 0.5% of inference time.

Abstract

Diffusion-based large language models (dLLMs) rely on bidirectional attention, which prevents lossless KV caching and requires a full forward pass at every denoising step. Existing approximate KV caching methods reduce this cost by selectively updating cached states, but their decision overhead scales with context length or model depth. We propose EntropyCache, a training-free KV caching method that uses the maximum entropy of newly decoded token distributions as a constant-cost signal for deciding when to recompute. Our design is grounded in two empirical observations: (1) decoded token entropy correlates with KV cache drift, providing a cheap proxy for cache staleness, and (2) feature volatility of decoded tokens persists for multiple steps after unmasking, motivating recomputation of the $k$ most recently decoded tokens. The skip-or-recompute decision requires only $O(V)$ computation per step, independent of context length and model scale. Experiments on LLaDA-8B-Instruct and Dream-7B-Instruct show that EntropyCache achieves $15.2\times$-$26.4\times$ speedup on standard benchmarks and $22.4\times$-$24.1\times$ on chain-of-thought benchmarks, with competitive accuracy and decision overhead accounting for only $0.5\%$ of inference time. Code is available at https://github.com/mscheong01/EntropyCache.