Token-Aware Cache Pruning

• Token-aware cache pruning is a technique that optimizes transformer inference by dynamically retaining, reusing, compressing, or discarding tokens based on their computed importance.
• It leverages various estimation methods—such as attention, similarity, and output perturbation—to balance inference speed and memory usage across long prompts and multimodal models.
• Practical implementations span long-context LLMs, visual transformers, and diffusion models, demonstrating substantial improvements in speed and resource efficiency while managing trade-offs in accuracy.

Token-aware cache pruning is a family of inference-time techniques that reduce the memory, bandwidth, and latency costs of transformer inference by retaining, reusing, compressing, or selectively discarding only those tokens, token channels, or token-associated cache entries judged relevant for the current computation. The underlying bottleneck is consistent across domains: long prompts make prefilling and time-to-first-token expensive; autoregressive decoding makes key-value (KV) cache size grow with context length and batch size; multimodal models inherit large visual prefixes; and diffusion models accumulate redundant token updates across denoising steps (Fu et al., 2024, Zhang et al., 2024, Qin et al., 11 Nov 2025, Lou et al., 2024). Recent work applies token-aware pruning to long-context LLMs, VLMs and VideoLLMs, diffusion transformers, privacy-aware shared caching, and long-horizon agent systems, with token importance defined by attention, similarity, transport cost, second-order output perturbation, temporal dynamics, quantization sensitivity, lifecycle state, or privacy masks (Yang et al., 24 Mar 2025, Li et al., 6 Jun 2026, Wu et al., 22 May 2026).

1. Computational setting and formal objectives

Token-aware cache pruning arises from the two-stage structure of transformer inference. In LLMs, prefilling computes the KV cache for the prompt and predicts the first output token, while decoding reuses that cache to generate subsequent tokens. For long prompts, prefilling becomes a major bottleneck, and the standard prefilling cost scales as approximately $\sum_{\ell=1}^L O(N^2 d)=O(LN^2 d)$, whereas dynamic token pruning replaces this with $\sum_{\ell=1}^L O(N_\ell^2 d)$ for a layer-dependent alive-token count $N_\ell$ (Fu et al., 2024). In long-context compression work, the KV cache is also expressed directly as a memory-budget problem: with $T$ retained tokens, hidden size $d$, and precision $b$, the total KV-cache memory is $M(T,b)=2\cdot d\cdot T\cdot b$ bits, and compression is formulated as choosing $(T,b)$ under a hard budget $M_{\max}$ to minimize downstream error $E(T,b)$ (Zhang et al., 2024).

A related formulation appears in diffusion transformers, where pruning is posed as a constrained optimization over a binary mask $\sum_{\ell=1}^L O(N_\ell^2 d)$0 that decides which tokens are updated and which are reused from cache at block $\sum_{\ell=1}^L O(N_\ell^2 d)$1 and timestep $\sum_{\ell=1}^L O(N_\ell^2 d)$2. The objective minimizes the deviation between full and cached outputs subject to a budget on skipped tokens (Lou et al., 2024). In multimodal VLMs, token pruning is often executed once during prefilling so that all subsequent decoding steps inherit a smaller fixed KV cache, while in VideoLLMs the cache may be pruned hierarchically across both input visual tokens and internal visual key-value pairs (Yang et al., 24 Mar 2025, Qin et al., 11 Nov 2025).

These formulations make clear that token-aware cache pruning is not restricted to simple token deletion. It includes token eviction, one-shot token selection, adaptive layer choice, low-bit token retention, unstructured channel sparsification, and cache layouts that preserve prefix compatibility or privacy constraints. This suggests that the topic is best understood as a general inference-time resource-allocation problem rather than as a single pruning heuristic.

2. Importance estimation and token selection rules

The central design choice is how token importance is estimated. In attention-based long-context pruning, LazyLLM assigns each prompt token $\sum_{\ell=1}^L O(N_\ell^2 d)$3 at layer $\sum_{\ell=1}^L O(N_\ell^2 d)$4 the score

$\sum_{\ell=1}^L O(N_\ell^2 d)$5

where $\sum_{\ell=1}^L O(N_\ell^2 d)$6 denotes multi-head self-attention probabilities and the score measures the token’s contribution to predicting the next token. Pruning is then performed by layer-specific percentile thresholding, and different subsets may be selected at different generation steps (Fu et al., 2024).

Other methods replace attention-weight heuristics with more structured criteria. SharpV defines spatial and temporal importance for each visual token using the Euclidean distance between unit-normalized vectors,

$\sum_{\ell=1}^L O(N_\ell^2 d)$7

then fuses temporal and spatial signals as $\sum_{\ell=1}^L O(N_\ell^2 d)$8 and derives a frame-wise retention ratio

$\sum_{\ell=1}^L O(N_\ell^2 d)$9

for $N_\ell$0, with $N_\ell$1 (Qin et al., 11 Nov 2025). TopV instead formulates visual token pruning as an optimal-transport problem with transport plan $N_\ell$2 minimizing a visual-aware cost

$N_\ell$3

where the three components are feature similarity, relative spatial distance, and absolute central distance; token importance is then $N_\ell$4 (Yang et al., 24 Mar 2025).

In LLM KV eviction, OBCache explicitly targets the perturbation in attention outputs rather than accumulated attention weights. It derives closed-form saliency scores for isolated value pruning, isolated key pruning, and joint key-value pruning from attention weights $N_\ell$5, logits $N_\ell$6, value states $N_\ell$7, and outputs $N_\ell$8, arguing that output-aware scores better approximate the true effect of removing a cached token (Gu et al., 9 Oct 2025). ASL addresses a different axis of the problem by selecting the pruning layer itself: it tracks rank vectors of attention-pooled tokens across a lookback window, computes the relative variance $N_\ell$9, and fixes the first layer $T$0 satisfying $T$1 (Taniguchi et al., 12 Jan 2026).

Online methods also use similarity-based redundancy criteria. Token Filtering maintains exponentially smoothed key and value anchors, computes per-layer similarity scores $T$2 and $T$3, fuses them by a variance-aware weight

$T$4

and obtains the joint redundancy score $T$5; attention is skipped when this score exceeds a learned threshold (Lee et al., 8 Dec 2025). IntentKV introduces session-level QueryMemory with geometric decay,

$T$6

and scores history positions by their attention alignment to this memory before applying a learned residual correction (Li et al., 6 Jun 2026).

Across these systems, token awareness is therefore instantiated at multiple levels: immediate next-token attention, temporal change, transport-based representativeness, output perturbation, rank stability, anchor similarity, and cross-turn intent.

3. Cache layouts, pruning mechanics, and systems compatibility

Selection rules alone do not determine runtime behavior; cache layout and execution policy are equally consequential. A major distinction is between fixed-size caches established during prefilling and dynamically changing caches during generation. TopV prunes only once at the end of prefilling, so the visual token set used during decoding is fixed, the KV buffers are consistently reduced by the same factor, and there is no dynamic resizing or re-padding at each generation step (Yang et al., 24 Mar 2025). SharpV also separates input and internal pruning: Stage-1 drops low-information visual tokens before they reach the LLM, while Stage-2 computes the layer-wise cosine similarity

$T$7

and discards a layer’s stored visual KV pairs whenever $T$8 (Qin et al., 11 Nov 2025).

Other systems retain the ability to revive pruned content. LazyLLM keeps a sparse main KV cache containing only “alive” tokens at each layer and an AuxCache storing hidden states for tokens pruned at that layer. A token needed later can be revived either from the KV cache, if already present, or from the corresponding auxiliary hidden state without replaying the full forward pass (Fu et al., 2024). In diffusion transformers, TokenCache reuses cached residual updates from the next timestep for tokens marked as low-importance, interleaving such pruned steps with full inference steps under a Two-Phase Round-Robin schedule (Lou et al., 2024).

Several papers emphasize hardware and serving compatibility. SharpV states that neither stage requires exposed attention scores or backpropagation and that both stages run in $T$9 rather than $d$0 or worse, making the method fully compatible with FlashAttention hardware (Qin et al., 11 Nov 2025). TopV makes a similar point: because pruning is not based on explicit per-head or per-query attention scores, FlashAttention’s fused implementation is not broken; only the pruned $d$1 and corresponding $d$2 are passed forward (Yang et al., 24 Mar 2025).

For long-horizon agents and shared serving, layout preservation becomes a first-class concern. IntentKV performs eviction through slot-map redirection rather than compaction: dropped positions are remapped to a sentinel dead slot with $d$3 and $d$4, while surviving rows retain their original slots, RoPE phases, and slot identities, preserving prefix-cache reuse across turns (Li et al., 6 Jun 2026). CachePrune, designed for cross-user KV sharing, stores token-level reusable non-sensitive spans, indexes them through rolling hashes and SHA-256 metadata, reconstructs a unified KV cache with “holes,” and leaves the unchanged engine to fill missing entries and recompute flagged tokens (Wu et al., 22 May 2026).

These systems show that token-aware cache pruning is as much about cache semantics as about token scoring. Whether a method compacts, redirects, revives, quantizes, or leaves holes determines its interaction with FlashAttention, prefix caches, multi-tenant serving, and decode-time memory management.

4. Representative instantiations across model classes

In long-context LLM inference, the dominant pattern is selective retention of prompt or history tokens under a fixed KV budget. LazyLLM dynamically chooses which tokens to compute in both prefilling and decoding, with token subsets allowed to vary across steps (Fu et al., 2024). ASL chooses a single selection layer adaptively and then propagates only the selected tokens through deeper layers (Taniguchi et al., 12 Jan 2026). OBCache replaces attention-only heuristics with output-aware second-order saliency scores (Gu et al., 9 Oct 2025). Quantized pruning treats the number of cached tokens and their precision as a joint design space rather than separate axes, chaining any token-eviction method with low-bit KV quantization (Zhang et al., 2024). Mustafar and SparK move below the token axis into unstructured sparsity: Mustafar prunes key and value elements per token by magnitude and stores them in a bitmap-based sparse format with a custom sparse attention kernel (Joo et al., 28 May 2025), while SparK applies query-aware channel pruning with on-the-fly recovery of pruned channels during attention score computation (Liao et al., 21 Aug 2025).

In multimodal VLMs and VideoLLMs, token-aware cache pruning is shaped by large visual prefixes. TopV performs optimal-transport-based visual token pruning once during prefilling and carries a reduced visual token set through all subsequent decoding steps (Yang et al., 24 Mar 2025). SharpV uses a two-stage design, combining adaptive visual token pruning based on spatial-temporal information with later self-calibrated pruning of degraded visual key-value pairs inside the model (Qin et al., 11 Nov 2025). QUOTA integrates token pruning with low-bit deployment by using W4A4 calibration signals to derive a deterministic layer-wise visual-token keep-ratio schedule that is executed under the same W4A4 operators and a 4-bit KV cache used at inference time (Li et al., 19 Apr 2026).

In diffusion models, token-aware cache pruning targets redundancy across denoising steps rather than across autoregressive history. TokenCache trains a Cache Predictor to decide which tokens to update and which to reuse from the next timestep, prunes only selected blocks, and schedules pruning with Two-Phase Round-Robin to control error accumulation (Lou et al., 2024). DaTo measures temporal token dynamics from adjacent diffusion steps, dedicates self-attention to high-dynamic tokens, and recovers pruned tokens from similar base tokens while caching whole feature maps across timesteps (Zhang et al., 2024). CAT Pruning likewise combines token selection with caching, using relative noise magnitude, spatial clustering, and staleness balancing to decide which tokens are recomputed and which reuse cached hidden states and predicted noise (Cheng et al., 1 Feb 2025).

In multi-turn agents and serving systems, token-aware cache pruning extends beyond model-internal attention. IntentKV scores live history tokens using session-level QueryMemory and prunes with a top-$d$5 rule under fixed budgets, explicitly targeting peak request tokens and KV read bandwidth (Li et al., 6 Jun 2026). CachePrune derives reusable token spans by masking sensitive tokens and retaining only non-sensitive self-contextualized substrings, thereby turning token-level privacy constraints into cache-management rules (Wu et al., 22 May 2026). TokenPilot addresses a related serving problem: it stabilizes cacheable prompt prefixes through ingestion-aware compaction and delays structural eviction until context segments become “evictable,” balancing text sparsity against prompt-cache continuity (Xu et al., 15 Jun 2026).

5. Empirical trade-offs in speed, memory, and accuracy

Across long-context LLMs, the reported gains are substantial but heterogeneous. LazyLLM on Llama 2-7B reported time-to-first-token speedups of 2.34× on multi-document question answering, 1.36× on single-document QA, 1.46× on summarization, 2.19× on few-shot learning, 2.89× on synthetic tasks, and 1.94× on code completion; on multi-document QA the score changed from 22.43 to 22.31, and only 63.9% of prompt tokens were processed by the end of generation, yielding a 1.56× end-to-end speedup (Fu et al., 2024). Quantized pruning found that under identical memory budgets, storing 4× tokens at 4-bit consistently outperformed storing 1× tokens at 16-bit on LongBench and Needle-in-a-Haystack; the latter showed gains of up to 10–20 percentage points, and on the RULER-8k retrieval subset the score on Llama-3 rose from 67.5 to 82.2 (Zhang et al., 2024). Mustafar reported that at $d$6 the KV cache size was approximately 45% of dense inference, LongBench average remained approximately 41.0, and throughput reached up to 2.23× compared to dense inference (Joo et al., 28 May 2025). OBCache improved average exact-match in passkey retrieval from 52.3% to 65.8% for H₂O+OBCache-Joint and raised H₂O LongBench average accuracy from 46% to 49% at 10% KV retention (Gu et al., 9 Oct 2025). Token Filtering at 50% pruning increased LLaMA-2-13B throughput at batch size 128 from 2.7 to 5.2 tokens per second while maintaining substantially better commonsense reasoning and MMLU accuracy than FLAP and SLiMGPT under the same pruning ratio (Lee et al., 8 Dec 2025). IntentKV, at an 8k KV budget, reduced mean peak request tokens by 23.9% on Qwen3-8B and 30.7% on Qwen2.5-14B; on the 100 longest BCP queries on Qwen2.5-14B, worst-case raw KV reads fell from 411M to 31M (Li et al., 6 Jun 2026).

In multimodal models, the gains often combine prefill reduction with decode-time KV savings. TopV removed up to 51% of visual-token FLOPs for LLaVA-7B and 48% for InternVL2-2B, achieved 1.68× average speedup on LLaVA-7B and 2.1× on InternVL2-2B, reduced dynamic memory by 49% and 61% respectively, and still reported average gains of +0.39% across 8 tasks and +1.2% across 10 tasks at around 50% pruning (Yang et al., 24 Mar 2025). SharpV reported an end-to-end token budget of approximately 12%, with visual retention around 32% and memory retention around 39%; time-to-first-token decreased from approximately 0.61 s to 0.37 s on the 0.5B model and from approximately 1.05 s to 0.64 s on the 7B model, GPU memory fell from approximately 18 GB to 9 GB and from 32 GB to 22 GB, and the average accuracy drop remained within 0.5 points, with occasional 0–1% gains on VideoMME and NExTQA (Qin et al., 11 Nov 2025). QUOTA, under W4A4 weights and activations with a 4-bit KV cache and only 30% final visual-token keep ratio, achieved 95.65% average retention, compared with about 94.3% for representative stage-wise quantize→prune and prune→quantize pipelines under the same low-bit regime (Li et al., 19 Apr 2026).

In diffusion models, token-aware caching trades speed against image quality rather than task accuracy. TokenCache reduced GFLOPs from 237.3 to 144.5 on DiT-XL/2 and from 133.8 to 80.0 on MDT, with latency improvements from 1167 ms to 883 ms and from 4606 ms to 3050 ms respectively (Lou et al., 2024). DaTo reported a 9.01× speedup on Stable Diffusion v1.5 over ImageNet while reducing FID from 27.64 to 27.31, and a 7.00× acceleration on COCO-30k while reducing FID from 12.15 to 9.98 (Zhang et al., 2024). CAT Pruning, with $d$7 retained tokens, reduced MACs from 168.28 T to 90.28 T at 28 steps and from 300.50 T to 136.70 T at 50 steps, corresponding to approximately 46.3% and 54.5% reductions and end-to-end speedups of approximately 1.90× and 2.10× (Cheng et al., 1 Feb 2025).

In serving and agent systems, the objective shifts toward reuse, cost, and privacy. CachePrune reported a 44% improvement in cache hit rate over fixed-chunk baselines and up to 4.5× faster time-to-first-token than no sharing, while direct recovery of any token marked sensitive remained 0% under its strict threat model (Wu et al., 22 May 2026). TokenPilot reported cost reductions of 61% on PinchBench in both isolated and continuous modes, 56% on Claw-Eval isolated mode, and 87% on Claw-Eval continuous mode, while maintaining competitive overall performance (Xu et al., 15 Jun 2026).

6. Limitations, misconceptions, and research directions

A common misconception is that token-aware cache pruning is equivalent to static prompt truncation. The recent literature shows otherwise. LazyLLM can prune a token in one step and revive it later from AuxCache (Fu et al., 2024); ASL chooses the pruning layer adaptively rather than fixing it a priori (Taniguchi et al., 12 Jan 2026); SharpV changes frame-wise retention ratios according to temporal change (Qin et al., 11 Nov 2025); and IntentKV chooses a request-specific top-$d$8 cut point rather than applying a global score threshold (Li et al., 6 Jun 2026). Another misconception is that effective pruning requires exposed attention scores and therefore conflicts with fused attention kernels. TopV and SharpV explicitly position themselves against this assumption, stating compatibility with FlashAttention because pruning decisions are made without explicit exposed attention matrices (Yang et al., 24 Mar 2025, Qin et al., 11 Nov 2025).

The main technical limitations are also recurrent. Over-pruning risks information loss, and LazyLLM reports that early layers are more sensitive to pruning than later ones (Fu et al., 2024). Token Filtering requires a warm-up period because early decoding steps suffer from “cold-start” anchors, and its tail-focused heuristic must be tuned per model and scale (Lee et al., 8 Dec 2025). OBCache notes structural bias toward early tokens, the limitations of a diagonal Hessian approximation, and the mismatch between a fixed perturbation window and truly future queries (Gu et al., 9 Oct 2025). SparK reduces KV storage but adds decode-time recovery overhead and has diminished returns at short contexts (Liao et al., 21 Aug 2025). In token-precision compression, 8-bit and 4-bit quantized pruning incur negligible loss after token pruning, whereas 2-bit quantization collapses performance (Zhang et al., 2024). QUOTA argues that naive stage-wise combinations of quantization and pruning are brittle because calibration assumes a full graph while deployment executes a pruned one, creating a calibration–execution mismatch (Li et al., 19 Apr 2026). TokenPilot identifies a system-level tension between text sparsity and prompt-cache continuity: unconstrained sequence mutations reduce text length but invalidate reusable prefixes (Xu et al., 15 Jun 2026).

Several broader directions are visible. One is the shift from heuristic attention accumulation toward richer notions of relevance, including output-aware saliency, session-level intent, and quantization risk (Gu et al., 9 Oct 2025, Li et al., 6 Jun 2026, Li et al., 19 Apr 2026). A second is the move from token-only pruning to multi-axis compression, where token count, precision, channels, and layout are optimized jointly rather than separately (Zhang et al., 2024, Liao et al., 21 Aug 2025). A third is the increasing importance of layout-preserving and policy-aware cache management in serving systems, where privacy, prefix reuse, and hardware compatibility matter as much as saliency. This suggests that future token-aware cache pruning systems will likely combine adaptive token selection with quantization, recoverable sparsity, and cache layouts designed explicitly for real serving stacks rather than for isolated model kernels alone.