---
title: Token-Aware Cache Pruning
url: https://www.emergentmind.com/topics/token-aware-cache-pruning
type: topic
---

# Token-Aware Cache Pruning

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 [2407.14057][2412.12706][2511.08003][2409.18523]. 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 [2503.18278][2606.09916][2605.23640].

## 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$ [2407.14057]. 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)$ [2412.12706].

A related formulation appears in diffusion transformers, where pruning is posed as a constrained optimization over a binary mask $m_l^t\in\{0,1\}^n$ that decides which tokens are updated and which are reused from cache at block $l$ and timestep $t$. The objective minimizes the deviation between full and cached outputs subject to a budget on skipped tokens [2409.18523]. 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 [2503.18278][2511.08003].

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 $t_i$ at layer $\ell$ the score
$$
s_i^\ell=\frac{1}{H}\sum_{h=1}^H A^\ell_{h,i,N},
$$
where $A^\ell\in\mathbb{R}^{H\times N\times N}$ 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 [2407.14057].

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,
$$
\mathrm{Dissim}(v_1,v_2)=\left\|\frac{v_1}{\|v_1\|_2}-\frac{v_2}{\|v_2\|_2}\right\|_2=\sqrt{2-2\cos\theta},
$$
then fuses temporal and spatial signals as $I_t[i]=T_t[i]+w\cdot S_t[i]$ and derives a frame-wise retention ratio
$$
\alpha_t=\frac{\|T_t\|_2}{2\sqrt{f}}
$$
for $t>1$, with $\alpha_1=\|S_1\|_2/(2\sqrt{f})$ [2511.08003]. TopV instead formulates visual token pruning as an optimal-transport problem with transport plan $P\in\mathbb{R}^{N\times N}$ minimizing a visual-aware cost
$$
C_v(s_i,t_j)=\alpha C_f(s_i,t_j)+\beta C_s(s_i,t_j)+\gamma C_e(s_i),
$$
where the three components are feature similarity, relative spatial distance, and absolute central distance; token importance is then $I_i=\sum_j P^*_{ij}$ [2503.18278].

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 $A$, logits $Z$, value states $V$, and outputs $O$, arguing that output-aware scores better approximate the true effect of removing a cached token [2510.07651]. 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 $\mathrm{RelVar}(L)=\mathrm{Var}_{\mathrm{top}}(L)/\mathrm{init\_var}$, and fixes the first layer $L_{\mathrm{sel}}$ satisfying $\mathrm{RelVar}(L_{\mathrm{sel}})<\tau$ [2601.07667].

Online methods also use similarity-based redundancy criteria. Token Filtering maintains exponentially smoothed key and value anchors, computes per-layer similarity scores $s_k$ and $s_v$, fuses them by a variance-aware weight
$$
\alpha=\frac{1/\sigma_v^2}{1/\sigma_k^2+1/\sigma_v^2},
$$
and obtains the joint redundancy score $S_{\mathrm{KV}}=\alpha s_k+(1-\alpha)s_v$; attention is skipped when this score exceeds a learned threshold [2512.07090]. IntentKV introduces session-level QueryMemory with geometric decay,
$$
\mathbf{M}_t=e^{-\lambda}\mathbf{M}_{t-1}+\mathrm{Enc}(\mathbf{q}_t),
$$
and scores history positions by their attention alignment to this memory before applying a learned residual correction [2606.09916].

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 [2503.18278]. 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
$$
\mathrm{cos}_\ell=\mathrm{mean}_n\,\cos(V_\ell[n],V_0[n])
$$
and discards a layer’s stored visual KV pairs whenever $\mathrm{cos}_\ell<M$ [2511.08003].

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 [2407.14057]. 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 [2409.18523].

Several papers emphasize hardware and serving compatibility. SharpV states that neither stage requires exposed attention scores or backpropagation and that both stages run in $O(n\cdot d)$ rather than $O(n^2\cdot d)$ or worse, making the method fully compatible with FlashAttention hardware [2511.08003]. 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 $V$ and corresponding $K$ are passed forward [2503.18278].

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 $K[s^\dagger]=-10^4\cdot\mathbf{1}_D$ and $V[s^\dagger]=0$, while surviving rows retain their original slots, RoPE phases, and slot identities, preserving prefix-cache reuse across turns [2606.09916]. 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 [2605.23640].

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 [2407.14057]. ASL chooses a single selection layer adaptively and then propagates only the selected tokens through deeper layers [2601.07667]. OBCache replaces attention-only heuristics with output-aware second-order saliency scores [2510.07651]. 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 [2412.12706]. 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 [2505.22913], while SparK applies query-aware channel pruning with on-the-fly recovery of pruned channels during attention score computation [2508.15212].

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 [2503.18278]. 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 [2511.08003]. 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 [2604.17320].

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 [2409.18523]. 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 [2501.00375]. 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 [2502.00433].

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-$k$ rule under fixed budgets, explicitly targeting peak request tokens and KV read bandwidth [2606.09916]. 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 [2605.23640]. 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 [2606.17016].

## 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 [2407.14057]. 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 [2412.12706]. Mustafar reported that at $K_s=V_s=0.7$ 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 [2505.22913]. 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 [2510.07651]. 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 [2512.07090]. 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 [2606.09916].

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 [2503.18278]. 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 [2511.08003]. 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 [2604.17320].

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 [2409.18523]. 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 [2501.00375]. CAT Pruning, with $\alpha=0.3$ 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× [2502.00433].

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 [2605.23640]. 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 [2606.17016].

## 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 [2407.14057]; ASL chooses the pruning layer adaptively rather than fixing it a priori [2601.07667]; SharpV changes frame-wise retention ratios according to temporal change [2511.08003]; and IntentKV chooses a request-specific top-$k$ cut point rather than applying a global score threshold [2606.09916]. 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 [2503.18278][2511.08003].

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 [2407.14057]. 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 [2512.07090]. 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 [2510.07651]. SparK reduces KV storage but adds decode-time recovery overhead and has diminished returns at short contexts [2508.15212]. In token-precision compression, 8-bit and 4-bit quantized pruning incur negligible loss after token pruning, whereas 2-bit quantization collapses performance [2412.12706]. 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 [2604.17320]. TokenPilot identifies a system-level tension between text sparsity and prompt-cache continuity: unconstrained sequence mutations reduce text length but invalidate reusable prefixes [2606.17016].

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 [2510.07651][2606.09916][2604.17320]. 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 [2412.12706][2508.15212]. 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.

Source: https://www.emergentmind.com/topics/token-aware-cache-pruning