PaperlayerView exact revision on arXiv

Source-linked AI summary

Mask-Guided KV Cache Eviction in Block Diffusion Language Models

Gleb Molodtsov, Ekaterina Alimaskina, Evgeny Uskov, Artur Zagitov, Aleksandr Beznosikov

arXiv:2610.06996v1cs.LGcs.AIcs.CL

TL;DR

Block diffusion’s growing KV cache raises both memory and repeated-read costs, while existing mask-query methods reduce reads but retain the full history. MaskAhead uses current-block queries for temporary reads and future-mask probes for retention, with value-aware scoring; Q-MaskAhead applies selection and attention directly to low-bit KV. On long-prompt QA, MaskAhead reduces KV memory by 9.5× on average with a 1.2-point mean F1 loss relative to dense inference, and Q-MaskAhead reaches 20.1× reduction with a 2.3-point mean F1 loss.

  • Problem

    Block diffusion repeatedly reads a growing KV cache during denoising, and existing mask-query methods reduce reads while retaining the full history.

    Source: §1 · p. 1

  • Method

    MaskAhead uses current-block masks for temporary reads and future-mask probes with value-aware scoring for persistent retention; Q-MaskAhead performs selection and attention from four-bit KV.

    Source: §2 · p. 3

  • Results

    MaskAhead achieves the highest F1 among evicting methods across six main QA settings, including one tie, with an average 9.5× KV-memory reduction and a 1.2-point mean F1 loss versus dense inference; Q-MaskAhead reaches 20.1× reduction with a 2.3-point mean F1 loss.

    Source: §6 · p. 9

  • Takeaways & Limitations

    Four-bit storage adds a 2.48× KV-memory reduction over bf16 MaskAhead on MATH-500 at a 0.6-point accuracy cost.

    Source: §5.3 · p. 8

    The formal equations describe decision horizons, not objectives that the implementation optimizes exactly.

    Source: §3.1 · p. 3

Abstract

from arXiv · show

Block diffusion language models keep a large key-value (KV) cache throughout generation and attend to it at every denoising step, limiting both memory capacity and generation speed. Reducing these costs requires deciding which past tokens to use for denoising the current block (selection) and which to keep in memory for future blocks (eviction). We propose MaskAhead, a training-free method that solves both tasks with a single mask-query-based ranking mechanism. Current-block masks guide selection, while probes of upcoming masked blocks guide eviction. Both rank KV entries by their estimated contribution to the attention output. Our quantized variant, Q-MaskAhead, computes selection and attention directly from low-bit KV, largely preserving the selected entries. Experiments on Fast-dLLM-v2, DreamReasoner, and LLaDA2.0-mini cover long-generation reasoning, long-prompt question answering, and needle-in-a-haystack retrieval. On long-prompt QA, MaskAhead reduces KV memory by $9.5\times$ on average with a 1.2-point mean F1 loss relative to dense inference. Q-MaskAhead increases the reduction to $20.1\times$ with a 2.3-point mean F1 loss. In a batch-32 systems profile, MaskAhead achieves $1.23\times$ end-to-end and $1.68\times$ decode-stage speedups over dense inference.

1 Introduction

Block diffusion restores a causal prefix for denoising, but its growing KV cache raises both resident-memory and repeated-read costs. MaskAhead uses current-block queries for temporary reads and future-mask probes to guide irreversible eviction.

1 Introduction

  • Block diffusion generates blocks sequentially while denoising their tokens in parallel, restoring a causal prefix and reusable KV cache.As the prompt and response grow, every new block retains more KV entries and each denoising step reads a longer prefix.Source: §1 · p. 1
  • Existing mask-query methods reduce cache reads but retain the full history, leaving memory reduction to irreversible eviction decisions.Block diffusion permits masked-position probes conditioned on the known prefix, although those probes cannot include tokens finalized in intermediate blocks.Source: §1 · p. 1
  • MaskAhead uses current-block queries to select temporary reads and upcoming masked-block queries to forecast which entries future blocks may need.Source: §1 · p. 1

2 Related Work

Prior work addresses sparse reads, cache eviction, and reduced-precision storage, often separately. MaskAhead combines temporary read selection with persistent eviction, while Q-MaskAhead adds low-bit storage.

2 Related Work

  • Sparse-read methods reduce repeated cache reads but keep the full prefix resident, whereas MaskAhead physically removes entries to enforce a storage budget.HERALD offloads the full cache to host memory; MaskAhead instead removes entries, while GroupKV uses hierarchical selection and predictive prefetching across cache tiers.Source: §2 · p. 2
  • Block diffusion requires recurrent, forward-looking eviction because finalized output continually appends to the exact cache and future masks are exposed only in the active block.MaskAhead tests future masks as a look-ahead signal and turns them into a training-free retention strategy.Source: §2 · p. 2
  • Q-MaskAhead combines four-bit storage with per-token selection and recurrent eviction, unlike prior diffusion quantization that quantizes the full cache representation.The contribution reports 97.44% read-set and 97.65% retention-set overlap between four-bit and bf16 storage, supporting a compact cache without a bf16 shadow copy.Source: §2 · pp. 2–3
  • MaskAhead separates temporary reads from irreversible retention, using current-block masks for reads and future-mask probes with value-aware eviction for a persistent KV budget.A no-write probe supplies future queries, and recurrent compaction enforces the storage budget.Source: §2 · pp. 2–3
  • MaskAhead achieves the highest F1 among evicting methods across all six main QA settings, including one tie, with 9.5× lower average KV memory than dense inference.Experiments cover Fast-dLLM-v2-7B, DreamReasoner-8B, and LLaDA 2.0 across long-generation reasoning, long-prompt QA, and needle-in-a-haystack retrieval.Source: §2 · p. 3

3 Problem Setup: Read and Retention Decisions

Read selection and retention are distinct decisions over a live prefix: one supplies the current block, while the other determines what remains available later. Future-block queries are only probes, and the tested horizon does not enumerate every query the retained cache will serve.

3.1 Two decisions over one prefix

  • Read selection chooses a temporary set under budget k_t, while retention replaces the persistent cache under budget C_t and makes excluded entries unavailable until recomputation.Their ideal objectives minimize different prefix-attention output errors; the stated equations describe decision horizons, not objectives the implementation optimizes exactly.Source: §3.1 · p. 3
  • At budget n=29, a shared read set covers 97.75% of within-block prefix-attention mass versus 98.50% for a fresh set; coverage falls to 74.69% one block later and 62.16% after four blocks.The fresh-set comparison includes entries appended after the shared set was selected, so the trace motivates refreshing reads but does not identify the best retention rule.Source: §3.1 · p. 3, Figure 1 · p. 4

3.2 Queries available before commitment

  • The no-write probe uses masks in the next M−1 blocks to estimate future cache demand, without appending KV or producing draft tokens.The tested setting uses E=4 and M=4 blocks; the retained cache also serves the first step of block t+4, which the probe does not cover, making the horizon an empirical design choice.Source: §3.2 · p. 4

4 MaskAhead: Two Decisions, Two Query Horizons

MaskAhead applies one value-aware scoring rule to two query horizons: current-block queries rank temporary reads, while current and future probes rank persistent retention. It gathers the read buffer before compacting the cache, and Q-MaskAhead applies the decisions directly to four-bit KV.

4 MaskAhead: Two Decisions, Two Query Horizons

  • MaskAhead builds separate rankings for remaining current-block reads and persistent KV retention, changing query evidence while keeping one scoring rule.Current-block masks select temporary reads, while current and future mask queries rank entries retained until the next eviction.Source: §4 · p. 4, Figure 2 · p. 5
  • The pooled score weights each value’s deviation from the attention centroid by its relevance to the query population, serving as a multi-entry removal proxy rather than an accuracy guarantee.The single-entry identity underlying the score does not solve the subset-selection problem.Source: §4 · p. 4
  • At eviction checkpoints, MaskAhead protects four initial positions and a 64-position recent window, fills remaining slots by score, then compacts persistent KV after gathering the current read buffer.Ranking and budgets are applied per layer and KV head; the first two layers remain dense for read selection.Source: §4 · p. 5
  • Q-MaskAhead ranks and attends directly from four-bit stored keys and values without a bf16 shadow cache, regrouping partial key groups after eviction.Keys use asymmetric per-channel quantization over groups of 32 positions; values use per-token groups of 32 channels with fp16 metadata.Source: §4 · p. 5

5 Experiments

Experiments across three models cover long-generation reasoning, long-prompt QA, and needle-in-a-haystack retrieval. They evaluate quality and KV memory under bounded caches, probe and score controls, quantization, batch profiling, and long-prefill heuristics.

5.1 Quality under a bounded cache

  • On Fast-dLLM reasoning, MaskAhead retains much of full-cache sparse-read quality while discarding prefix entries, but dense inference remains more accurate and evicting-method rankings vary by task.Expected Attention's auxiliary Gaussian query state can put KV memory above the dense cache despite eviction.Source: §5.1 · pp. 6–7, Table 2 · p. 6
  • Across six long-prompt QA settings, MaskAhead achieves the highest F1 among evicting methods, including one tie, with 9.5× average KV-memory reduction and a 1.2-point mean F1 loss versus dense inference.Q-MaskAhead raises the average reduction to 20.1× with a 2.3-point mean F1 loss.Source: Table 3 · p. 7, §5.1 · p. 7

5.2 What evidence does each decision need?

  • Future-mask probes improve retention evidence without materially changing MATH-500 KV memory; two- and four-block probes have similar point estimates, with no monotonic gain from longer horizons.The paired mask-versus-DapQ comparison on MATH-500 records 28 wins and 11 losses (p=0.0095, two-sided sign test).Source: §5.2 · p. 7, Table 4 · p. 7
  • Value-aware scores have the highest point estimates among evicting score combinations, but paired quality intervals against attention-mass scoring include zero on both tasks.The factorization varies read and retention scores independently at fixed queries and budgets.Source: Table 5 · p. 7, §5.2 · p. 8

5.3 Precision changes the memory and quality budgets

  • Four-bit storage adds a 2.48× KV-memory reduction over bf16 MaskAhead on MATH-500 at a 0.6-point accuracy cost; precision-related quality costs vary by task and representation.Rotated-and-centered keys recover much of untransformed K4V4's larger HotpotQA F1 loss at unchanged KV memory, though rotation and centering effects are not separated.Source: §5.3 · p. 8, Table 6 · p. 8

5.4 Batch-profile memory and execution time

  • At batch 32, MaskAhead reduces KV memory by 85.1% and achieves 1.23× complete-batch and 1.68× decode-time ratios versus dense bf16.Q-MaskAhead reduces KV memory by 89.9%, with 1.20× complete-batch and 1.64× decode-time ratios; K4V4-only storage reduces resident KV but not recorded decode time.Source: §5.4 · p. 8, Figure 3 · p. 9

5.5 Retention heuristics for long prefill

  • Question probe preserves substantially more answer quality than chunked prefill while using less KV memory, although the comparison changes both ranking and retention timing.Combining question probe with four-bit storage sharply lowers F1 on all three QA tasks; four-bit storage also incurs a larger NIAH retrieval loss than bf16 MaskAhead and DapQ.Source: §5.5 · p. 9

6 Conclusion

MaskAhead leads the reported evicting methods across the six main QA settings, while Q-MaskAhead trades some F1 for greater KV-memory reduction. The systems profile reports lower memory and decode times for both methods, with a separate batch-64 record.

6 Conclusion

  • MaskAhead has the highest F1 among evicting methods in all six main QA settings, including one tie, with 9.5× average KV-memory reduction and 1.2-point mean F1 loss; Q-MaskAhead reaches 20.1× and 2.3 points.The comparisons are relative to dense inference.Source: §6 · p. 9

A.1 Measurement scope and complete batch profile

  • Figure 4 covers successful batch sizes 1–32, while the separate batch-64 record reports OOM for dense bf16 and K4V4-only storage.At batch 64, MaskAhead uses 8.5 GiB resident KV and takes 242.4 s total; Q-MaskAhead uses 5.8 GiB and takes 245.0 s.Source: Figure 4 · p. 12, §A.1 · p. 12, Table A.1 · p. 13

A.2 Generated length in the quality runs

The budget results quantify memory–quality tradeoffs, while generated lengths vary across methods and fixed-output systems timings do not determine latency under quality-run stopping rules. An accompanying traffic accounting contrasts dense and sparse block reads but omits several execution costs.

A.2 Generated length in the quality runs

  • MaskAhead generates 617.8 tokens per MATH-500 example, 1.0% more than dense bf16, while Q-MaskAhead generates 646.8, 5.7% more.On GSM8K, both variants increase generated length by 0.6%; Q-MaskAhead adds up to 1.6 tokens per example on NarrativeQA.Source: §A.2 · p. 13, Table A.2 · p. 13
  • Because quality runs use configured stopping rules and response lengths differ, latency cannot be inferred from the fixed-length batch profile.The length difference is especially pronounced for four-bit storage.Source: §A.2 · pp. 13–14

A.3 Amortized attention traffic

  • Dense attention reads approximately T·N entries per block, versus N + (T−1)·k for MaskAhead, with probe and compaction work amortized across checkpoints.The accounting omits score computation, gathering, dense-read exceptions, kernel launches, and request scheduling, so cache reduction does not imply an equal reduction in total traffic or full-run allocation.Source: §A.3 · p. 14

B Read and Retention Budget Tradeoffs

  • Across six model/task settings, MaskAhead reduces KV memory by 9.5× on average with a 1.2-point mean F1 loss; Q-MaskAhead reaches 20.1× reduction with a 2.3-point mean loss versus dense inference.The mean F1 difference between bf16 MaskAhead and Q-MaskAhead is 1.1 points in favor of bf16.Source: B Read and Retention Budget Tradeoffs · p. 14

B.1 Reasoning capacity sweep

  • With the read floor fixed at 128, MaskAhead’s accuracy advantage over H2O is 11.6 points at retention floor 128, 3.8 at 256, and 0.2 at both 512 and 1024.The comparison uses the same 500 MATH-500 examples; KV includes growth between four-block eviction checkpoints.Source: §B.1 · p. 14, Table B.1 · p. 14

B.2 Common-harness comparison at a wider budget

At the wider budget, results vary by task: MaskAhead shares the top MATH-500 score with Expected Attention and leads on two QA tasks, while other methods lead or tie elsewhere. Increasing the read floor improves MATH-500 accuracy at nearly unchanged KV memory.

B.2 Common-harness comparison at a wider budget

  • At the wider budget on MATH-500, historical CAOTE-style retention reaches 57.6%, MaskAhead and Expected Attention 57.2%, and DapQ 56.2%.Expected Attention uses 2.346 times dense bf16 KV memory because its Gaussian state dominates the retained cache.Source: §B.2 · p. 14, Table B.2 · p. 15
  • MaskAhead scores 39.10 and 24.97 F1 on HotpotQA and NarrativeQA, exceeding DapQ’s 39.04 and 24.43; H2O leads MuSiQue at 18.02, and H2O, DapQ, and MaskAhead tie GSM8K at 85.44%.Q-MaskAhead scores 86.35% on GSM8K across 1,319 examples.Source: §B.2 · pp. 14–15, Table B.2 · p. 15

B.3 Read-budget controls at fixed retention

  • On MATH-500, raising the read floor from 128 to 512 increases accuracy from 54.6% to 58.4% while KV memory stays near 0.62 of dense bf16.The retention setting is fixed; on HotpotQA, both bf16 and four-bit curves improve as the read fraction increases.Source: §B.3 · p. 15, Figure 5 · p. 15

B.4 Long-context budget and read-fraction series

In long-context QA, MaskAhead improves HotpotQA and NarrativeQA over H2O across the tested budget fractions while using less KV memory, but read-fraction preferences vary by task. The accompanying controls distinguish attention-coverage diagnostics from complete-answer comparisons of query and score choices.

B.4 Long-context budget and read-fraction series

  • Across three long-context budget fractions, MaskAhead improves HotpotQA and NarrativeQA over H2O while using less KV memory; H2O leads MuSiQue by 0.30 F1 at C=20%.These runs use 32-read and 64-retained-entry floors with k=C/2, unlike the main QA and wider-budget settings.Source: §B.4 · p. 15, Table B.3 · p. 15
  • At fixed C=20%, raising HotpotQA reads from 2.5% to 10% increases F1 from 36.85 to 39.04 at the same 0.206 KV memory.NarrativeQA and MuSiQue have different read-fraction orderings, so one task does not determine a read budget for all QA workloads.Source: §B.4 · pp. 15–16, Table B.4 · p. 16

C Query Horizons and Ranking Criteria

C.1 Recorded within-block coverage

The appendix tests how read-set coverage, future-mask probes, scoring choices, and eviction timing affect block-diffusion cache decisions. Its diagnostics show where shared reads and retention choices differ from offline references, while leaving some scoring effects unresolved.

C.1 Recorded within-block coverage

  • The recorded comparison uses a block-shared read set against a same-budget offline top-k set selected from each later step’s realized queries.The MATH-500 recording starts with 30 examples, with fewer reaching later denoising steps; the offline set is an attention-coverage reference, not used to generate answers.Source: §C.1 · p. 16, Table C.1 · p. 16
  • The block-shared read set stays close to the same-budget offline reference within a block, but overlap with a fresh set falls to 47.2% at offset four.The candidate prefix expands across blocks; these diagnostics do not compare retention-policy answer quality.Source: §C.1 · p. 16, Table C.1 · p. 16

C.2 Probe horizon and query-source controls

  • 28 wins and 11 losses favor masks over DapQ on paired MATH-500 examples; two- and four-block probes score 54.4% and 54.6%, versus 53.2% for eight blocks.The sweep supports short lookahead rather than a monotonic benefit from more future positions; HotpotQA variants differ by at most 0.55 F1 points under the 32-token output protocol.Source: §C.2 · p. 16

C.3 Independent read and retention scores

  • Value-aware ranking for both read and retention decisions has the highest evicting-policy point estimate on MATH-500 and HotpotQA, but both comparisons’ confidence intervals include zero.Relative to attention mass for both decisions, differences are +2.6 MATH-500 points (95% CI [−0.6, 5.8]) and +0.568 HotpotQA F1 (CI [−0.008, 1.348]).Source: §C.3 · pp. 16–17

C.4 Protected positions and eviction interval

  • Evicting every two blocks yields 0.543 dense-relative MATH-500 KV memory and 51.8% accuracy; every eight blocks yields 0.718 memory and 53.8% accuracy.Table C.2 reports timing and protection controls at the main budgets, with KV including method state and normalized to dense bf16.Source: §C.4 · p. 17, Table C.2 · p. 17

D Four-Bit Storage and Ranking Stability

  • The precision experiments separately examine stored key/value precision, retention floors under K4V4, a key transform, and ranking overlap against bf16.Early-retention quantization is evaluated with prefill policies in Section E.2.Source: D Four-Bit Storage and Ranking Stability · p. 17

D.1 Key/value precision and retention floors

These experiments show how four-bit precision, retention floors, key transforms, and ranking stability relate to cache size and task scores. The results distinguish memory-saving operating points from answer-quality outcomes and rank overlap.

D.1 Key/value precision and retention floors

  • At the main MATH-500 retention policy, K4V4 gives 54.0% accuracy at 0.252 dense-relative KV memory.Reducing key precision to two bits costs more accuracy than reducing value precision at similar cache size; the lowest-memory format is a different accuracy operating point.Source: §D.1 · p. 17, Table D.1 · p. 17
  • With k_min = 128, raising the K4V4 retention floor from 128 to 256 improves MATH-500 accuracy from 50.0% to 54.0%.Larger floors add KV memory with little change in measured accuracy.Source: §D.1 · p. 17, Table D.2 · p. 17

D.2 Rotated and centered keys

  • At matched HotpotQA budgets and 0.050 KV memory, rotated-and-centered keys reach 37.33 F1 versus 34.75 for untransformed K4V4.Both variants change together, so the comparison does not isolate their individual effects; aggregates use 200 examples each, without a paired interval.Source: §D.2 · p. 17, Table D.3 · p. 18

D.3 Ranking overlap under stored precision

  • K4V4 and bf16 selected sets overlap by 97.44% on average, while retained sets overlap by 97.65%; high overlap measures rank stability, not equal outputs or answer quality.The intervals are [97.12, 97.77]% and [97.42, 97.90]%, respectively; the passage notes HotpotQA K4V4 quality loss.Source: §D.3 · p. 18

E Prefill Policies and Long-Prompt Retrieval

The appendix compares prefill policies and four-bit storage for long-prompt retrieval, measuring both task quality and memory. Results vary across prefill strategies, and adding K4V4 to early retention lowers F1 on all three QA tasks.

E.1 Prefill heuristics

  • Full prefill reaches 38.41 HotpotQA F1 at 0.999 maximum KV memory relative to dense bf16, while question probe reaches 37.16 F1 at 0.144 memory.Question probe uses question-conditioned evidence for early retention.Source: §E.1 · p. 18
  • Chunked prefill reaches 26.90 HotpotQA F1 at 0.185 KV memory, while all three prefill policies leave the same 0.107 normalized resident cache after response.Chunked prefill evicts between 512-token chunks; maximum KV memory during prompt processing and generation differs from post-response resident memory.Source: §E.1 · p. 18, E Prefill Policies and Long-Prompt… · p. 18
  • Question probe reaches 24.10 F1 at 0.125 KV memory on NarrativeQA and 15.35 F1 at 0.129 on MuSiQue.The prefill comparisons measure maximum KV memory over prompt processing and generation unless stated otherwise.Source: §E.1 · p. 18, E Prefill Policies and Long-Prompt… · p. 18

E.2 Precision boundary for early retention

  • Applying K4V4 at the same early-retention timing lowers KV memory further but also lowers F1 on all three QA tasks.The bf16 prefill result and four-bit generation-time storage are separate operating points; their combination requires its own quality evaluation. Table E.1 uses token F1 and 200 examples per task.Source: §E.2 · p. 18, Table E.1 · p. 18

E.3 NIAH passage retrieval

  • On 200 NIAH retrieval examples, MaskAhead scores 82.5% versus DapQ’s 82.0% at the same 0.1042 KV memory; full-cache MAGE scores 83.5%.Table E.2 reports accuracy as percent and distinguishes maximum from post-response resident KV, both relative to dense bf16.Source: §E.3 · p. 18, Table E.2 · p. 18
  • Q-MaskAhead scores 75.5% on NIAH at 0.0503 KV memory; this retrieval comparison uses full bf16 prefill, whose memory peak is near dense.The table’s resident KV is measured after the response, while maximum KV captures the peak.Source: §E.3 · p. 18, Table E.2 · p. 18
Loading 2610.06996v1…