MaskKV: Cache Eviction for Diffusion LLMs

• The paper introduces MaskKV, a training-free cache eviction framework that leverages mask-query guided scoring and adaptive budgeting to reduce memory and runtime overhead in diffusion LLMs.
• It proposes a novel mask-voting mechanism to identify key prompt tokens, outperforming traditional autoregressive eviction heuristics by exploiting bidirectional attention.
• Empirical results on LLaDA and Dream-7B indicate that MaskKV retains 94–98% of full-cache performance while delivering up to 31× acceleration and substantial memory savings.

MaskKV is a training-free cache eviction framework for diffusion LLMs (dLLMs) that targets the distinctive memory and runtime costs induced by bidirectional attention and iterative denoising. In dLLMs, the sequence initially consists of a prompt followed by many [MASK] tokens, and every denoising step reprocesses the full sequence rather than extending a causal history as in autoregressive models. MaskKV addresses this setting with two coupled mechanisms: a mask-query guided scoring rule that uses attention from mask tokens to rank prompt-token importance, and an adaptive cache budgeting strategy that allocates budget non-uniformly across layers and heads. On LLaDA with MaskKV, compressing the KV cache to only 256 pairs, described as less than 5% of tokens, retains 94% of the full-cache performance on LongBench and achieves up to 31x acceleration at 32k prompt length (Huang et al., 10 Oct 2025).

1. Problem formulation in diffusion LLM inference

The motivating claim behind MaskKV is that dLLMs require a different cache-compression logic from autoregressive models. Diffusion LLMs generate text by iterative denoising of a sequence initialized as

$x^{(T)} = [c_1, \dots, c_M, \underbrace{[MASK], \dots, [MASK]}_{L}],$

where $c=(c_1,\dots,c_M)$ is the prompt and $L$ is the target length. The prompt token set is $\mathcal{P}=\{1,\dots,M\}$, the response token set is $\mathcal{R}=\{M+1,\dots,M+L\}$, and layers are indexed by $\ell=1,\dots,D$ (Huang et al., 10 Oct 2025).

The paper formalizes the cacheable per-position bundle as

$\mathrm{Feat}_{\ell}^{(t)}(i) = \{ K_{\ell,i}^{(t)}, V_{\ell,i}^{(t)}, \mathrm{Attn}_{\ell,i}^{(t)}, \mathrm{FFN}_{\ell,i}^{(t)} \},$

with layer computation

$h_{\ell}^{(t)} = \mathrm{Attn}_{\ell}\!\bigl(Q_{\ell}^{(t)}, K_{\ell}^{(t)}, V_{\ell}^{(t)}\bigr) + \mathrm{FFN}_{\ell}. \tag{1}$

This framing makes explicit that the cache in dLLMs is not limited to a past-only KV history. It is repeatedly updated during denoising and interacts with prompt, decoded tokens, and masks under bidirectional attention. The paper identifies three resulting burdens: large KV memory footprint, repeated cache updates every denoising step, and substantial runtime overhead for long contexts.

MaskKV’s central premise is that autoregressive cache eviction methods do not transfer directly because they presuppose causal or past-only attention and importance estimates derived from token histories. In dLLMs, by contrast, future masked positions participate in attention, and the model’s generative process is organized around mask-token recovery. This difference is the foundation for treating mask tokens as informative observers of prompt importance rather than reusing autoregressive heuristics.

2. Diffusion-specific motivation and failure of autoregressive heuristics

The paper argues that existing autoregressive cache eviction methods fail in dLLMs for structural rather than merely empirical reasons. Methods such as SnapKV and H2O rely on causal attention patterns, whereas dLLMs use bidirectional attention in which prompt, response, and mask tokens interact symmetrically during denoising (Huang et al., 10 Oct 2025).

A key observation is that non-mask tokens in dLLMs often exhibit strong locality: they mostly attend to themselves or nearby neighbors, making their attention patterns a weak signal for global prompt importance. Mask queries behave differently. The paper reports that mask queries show sparse long-range retrieval, stable attention across denoising steps, and consistent focus on a small set of pivotal prompt tokens. This is the empirical basis for the method’s “Mask-Voting” component.

The paper also reports architectural asymmetries that undermine uniform cache allocation. First and last layers are much more important than middle layers, yielding a bimodal layer-importance profile. Heads likewise separate into information heads that read from the prompt and structure heads that focus on masks or planning. MaskKV treats these observations as design constraints: if cache is allocated uniformly across layers and heads, resources are wasted on components that contribute less to prompt-grounded generation.

This motivates the framework’s characterization as a task-aligned pruning method rather than a generic compression method. The task-aligned aspect is specific: cache retention is tied to mask-token evidence and to measured layer/head roles within diffusion inference.

3. Mask-query guided scoring and token selection

The first core component of MaskKV is mask-query guided scoring, also called “Mask-Voting.” Instead of ranking prompt tokens using prompt-token self-attention or historical autoregressive statistics, the method uses attention from mask queries at the initial step to estimate which prompt tokens are globally important (Huang et al., 10 Oct 2025).

The paper states that, for a mask token $m$, the final hidden state satisfies

$h_m^{(L)} = h_m^{(0)} + \sum_{l=1}^{L}\Delta_m^{(l)}. \tag{2}$

It further argues that multi-head attention is the only pathway through which prompt information can enter mask tokens. Within that framing, attention emitted by mask queries becomes the direct signal for prompt relevance.

The scoring rule begins with the attention matrix from all mask queries to all keys:

$c=(c_1,\dots,c_M)$0

Prompt-token importance is then aggregated as

$c=(c_1,\dots,c_M)$1

where $c=(c_1,\dots,c_M)$2 is the number of mask tokens and $c=(c_1,\dots,c_M)$3 is the number of prompt tokens. The resulting vector $c=(c_1,\dots,c_M)$4 defines a universal prompt-token importance ranking. Given a head- and layer-specific budget $c=(c_1,\dots,c_M)$5, the kept tokens are

$c=(c_1,\dots,c_M)$6

The interpretation is explicit in the paper: a higher $c=(c_1,\dots,c_M)$7 means that token $c=(c_1,\dots,c_M)$8 is repeatedly selected by mask queries and is therefore more valuable to preserve. On gsm8k with budget 128, the token-selection ablation reports 64.90 for SnapKV, 60.35 for Prompt-Voting, 66.64 for All-Voting, and 68.08 for Mask-Voting, which the paper uses to support the claim that mask-guided voting is the strongest selection rule among the tested variants.

4. Adaptive cache budgeting across layers and heads

The second core component is adaptive cache budgeting. MaskKV does not only score prompt tokens; it also decides where cache budget should be spent. The paper describes this as a two-stage hierarchical allocation: first across layers, then across heads (Huang et al., 10 Oct 2025).

For layer-wise allocation, the importance of layer $c=(c_1,\dots,c_M)$9 is measured by how much it transforms its input representations:

$L$0

Lower cosine similarity implies larger transformation and therefore greater importance. The observed bimodal profile motivates a partition into boundary layers $L$1 and middle layers $L$2. With total prompt budget $L$3, MaskKV first assigns a uniform base share

$L$4

then defines the remaining importance-driven budget

$L$5

That remainder is split between boundary and middle groups:

$L$6

$L$7

and each layer receives

$L$8

Head-wise redistribution then allocates each layer budget according to prompt dependence. For head $L$9 in layer $\mathcal{P}=\{1,\dots,M\}$0, the prompt preference score is

$\mathcal{P}=\{1,\dots,M\}$1

where $\mathcal{P}=\{1,\dots,M\}$2 is mask-to-prompt attention and $\mathcal{P}=\{1,\dots,M\}$3 is mask-to-mask attention. After normalization,

$\mathcal{P}=\{1,\dots,M\}$4

the final head budget is

$\mathcal{P}=\{1,\dots,M\}$5

Here $\mathcal{P}=\{1,\dots,M\}$6 provides a uniform floor, while the remaining budget is distributed according to prompt preference.

The paper’s ablations report the following progression: Mask-Voting alone, 35.39; + SqueezeAttention, 35.62; + PyramidKV, 33.52; + Boundary-Aware (online), 35.65; + Boundary-Aware (offline), 35.74; + adaKV head allocation, 35.70; + Prompt-Preference, 35.96; Full MaskKV, 36.27. This is used to support three claims: boundary-aware layer allocation helps, offline planning is better than online, and prompt-preference head allocation improves over AdaKV. The reported stable setting is $\mathcal{P}=\{1,\dots,M\}$7 and $\mathcal{P}=\{1,\dots,M\}$8; larger values make allocation too uniform, and smaller values make it overly aggressive and unstable.

5. Refresh policy, implementation choices, and system-level optimizations

MaskKV is not only a scoring-and-budgeting method; it is also specified as an inference system with explicit refresh rules and memory-saving optimizations. Prompt tokens are refreshed only every $\mathcal{P}=\{1,\dots,M\}$9 steps. Response tokens are refreshed either periodically every $\mathcal{R}=\{M+1,\dots,M+L\}$0 steps or when their value vectors change significantly (Huang et al., 10 Oct 2025):

$\mathcal{R}=\{M+1,\dots,M+L\}$1

$\mathcal{R}=\{M+1,\dots,M+L\}$2

$\mathcal{R}=\{M+1,\dots,M+L\}$3

The paper emphasizes an offline budgeting paradigm. Online layer budget allocation is said not to reduce peak memory effectively because importance scores from all layers would have to be collected while the full cache is still materialized. MaskKV therefore precomputes importance profiles on a calibration set, determines budgets in advance, and performs immediate layer-by-layer eviction.

Two implementation optimizations are highlighted. The first, Prompt-state Exclusion, removes redundant prompt features from the cache: only prompt KV states are kept, because prompt-side attention outputs and MLP outputs do not affect downstream generation according to the paper’s analysis. The second, Mask-Only Projection, restricts the final vocabulary projection to masked positions, since prompt logits are computed but never used.

For bf16 precision, the KV-memory formula is given as

$\mathcal{R}=\{M+1,\dots,M+L\}$4

where $\mathcal{R}=\{M+1,\dots,M+L\}$5 is sequence length, $\mathcal{R}=\{M+1,\dots,M+L\}$6 the number of heads, $\mathcal{R}=\{M+1,\dots,M+L\}$7 the head dimension, $\mathcal{R}=\{M+1,\dots,M+L\}$8 bytes, and the factor 2 accounts for K and V. The experimental configuration uses 8 × NVIDIA A100 80GB GPUs and evaluates LLaDA-8B-Instruct and Dream-7B-Instruct. Default settings are prompt refresh interval 50, response refresh interval 5, transfer ratio 0.25, and block length 8.

6. Empirical performance, scope, and relation to adjacent KV-cache paradigms

The main evaluation benchmark is LongBench, covering six domains: single-document QA, multi-document QA, summarization, few-shot learning, synthetic tasks, and code completion. The datasets listed are NarrativeQA, Qasper, MultiFieldQA-en, HotpotQA, 2WikiMultihopQA, MuSiQue, GovReport, QMSum, MultiNews, TREC, TriviaQA, SAMSum, PassageCount, PassageRetrieval-en, LCC, and RepoBench-P (Huang et al., 10 Oct 2025).

For LLaDA-8B-Instruct, the full-cache references are 38.61 for dLLM w/ Cache and 39.14 for dLLM w/o Cache. MaskKV achieves 31.25 at $\mathcal{R}=\{M+1,\dots,M+L\}$9, 35.60 at $\ell=1,\dots,D$0, 36.42 at $\ell=1,\dots,D$1, and 36.65 at $\ell=1,\dots,D$2. The paper summarizes this as retaining 94.33% of dLLM-Cache performance with 256 KV pairs. For Dream-v0-Instruct-7B, the corresponding full-cache references are 32.70 and 33.04, while MaskKV reports 25.51 at $\ell=1,\dots,D$3, 30.49 at $\ell=1,\dots,D$4, 30.88 at $\ell=1,\dots,D$5, and 31.34 at $\ell=1,\dots,D$6, summarized as retaining 98.66% of dLLM-Cache performance at 256 KV pairs.

The low-budget regime is particularly emphasized. At $\ell=1,\dots,D$7, MaskKV is reported to substantially outperform all baselines, and on LLaDA-8B it beats the best competing baseline by 7.02 points under budget 32. The paper also notes that in some cases it even exceeds the full-context dLLM-Cache baseline, likely because removing noisy context helps focus generation. In efficiency terms, with the proposed optimizations, at 32K context MaskKV achieves up to 31× faster decoding and 65% lower peak memory than LLaDA, and supports up to 8× longer prompts on an RTX 4090 GPU. A further observation is that later mask positions perform best for voting; specifically, the last, or back, mask position gives the best average.

The paper’s stated limitations are that experiments are limited to 7B/8B-scale open-source dLLMs, that the observed attention patterns and effectiveness are not yet validated on larger-scale or smaller lightweight models, and that evaluation is limited to text generation benchmarks. Extending the method to multimodal reasoning is proposed as a future direction.

MaskKV sits within a broader design space of KV-cache reduction methods. MiKV, an autoregressive method, argues that pure eviction can cause irreversible context loss, including safety breaches, hallucination, and incoherent generation, and instead keeps evicted KV pairs in low precision while retaining important KV pairs in higher precision (Yang et al., 2024). MKA, by contrast, approaches the KV bottleneck through hierarchical routing over local, session, and long-term memory tiers, with FastMKA fusing these routes into a single attention path for improved efficiency rather than evicting tokens (Liu et al., 21 Mar 2026). This suggests that MaskKV, MiKV, and MKA address related memory pressures through different mechanisms: diffusion-specific eviction guided by mask queries, mixed-precision retention in autoregressive inference, and hierarchical memory routing for long-context attention.