Mask Tokens as Prophet: Fine-Grained Cache Eviction for Efficient dLLM Inference
This addresses the problem of inefficient long-context handling in dLLMs for resource-limited settings, offering a domain-specific incremental improvement over existing ARM-focused strategies.
The paper tackles the high memory cost of cache mechanisms in diffusion large language models (dLLMs) by introducing MaskKV, a training-free cache eviction framework that compresses the KV cache to 256 pairs (less than 5% of tokens) while retaining 94% of full-cache performance on LongBench and achieving up to 31x acceleration at 32k prompt length.
Diffusion large language models (dLLMs) present a promising alternative to dominant autoregressive models (ARMs) by the ability of parallel decoding at the expense of substantial computation and memory costs. Specifically, the cache mechanism for bidirectional attention in dLLMs demands large memory footprint, restricting their ability to handle long contexts under resource-limited settings. Existing cache eviction strategies are designed for ARMs and ignore the unique characteristics of dLLMs, thus leading to unsatisfactory performance. To address these challenges, we introduce MaskKV, a training-free cache eviction framework tailored to dLLMs, focusing on the effect of mask tokens in dLLMs. MaskKV is built on two key innovations: (1) a mask-query guided scoring mechanism that leverages attention weights to identify and evict less critical prompt tokens for each head; (2) an adaptive cache budgeting strategy that improves efficiency by reducing allocation in intermediate layers and concentrating resources on prompt-preferring heads. On LLaDA with MaskKV, compressing the KV cache to only 256 pairs (less than 5% of tokens) retains 94% of the full-cache performance on LongBench and achieves up to 31x acceleration at 32k prompt length. The code is publicly available at: https://github.com/jianuo-huang/MaskKV