diff --git a/README.md b/README.md index 667a7224..df321508 100644 --- a/README.md +++ b/README.md @@ -131,6 +131,7 @@ Some presses rely on a different logic: - `KVgradPress` ([source](kvpress/presses/kvgrad_press.py), [paper](https://openreview.net/forum?id=cg1wTCJjjk)): similar to `KVzipPress`, but scores KV pairs by an input × gradient attribution of their effect on the last hidden states, at the cost of one backward pass per chunk. - `RestoreKVPress` ([source](kvpress/presses/restorekv_press.py), [paper](https://arxiv.org/abs/2608.01247)): extends the KV cache with 8 learned restore tokens encoded by a LoRA module before budget-matched KVzip pruning. - `KVComposePress` ([source](kvpress/presses/kvcompose_press.py), [paper](https://arxiv.org/abs/2509.05165)): attention-guided eviction, aligning per-head selections into composite tokens to preserve cache structure. +- `QueryAwareKVzipPlusPress` ([source](kvpress/presses/query_aware_kvzip_plus_press.py), [paper](https://www.alphaxiv.org/abs/2610.afterglow-query-aware-kv-cache-eviction)): single-pass KVzip+ scores from the attention of the question, with Afterglow and a predictability gate (query-aware compression only) > [!NOTE] > `KVComposePress` performs an extra pass over the full context, temporarily creating a KV cache of ~2x the context length and creating memory overhead during prefill. @@ -141,6 +142,7 @@ Finally we provide wrapper presses that can be combined with other presses: - `PerLayerCompressionPress` ([source](kvpress/presses/per_layer_compression_press.py)): compress each layer with a different compression ratio (experimental) - `ComposedPress` ([source](kvpress/presses/composed_press.py)): compose multiple presses together by chaining their forward hooks - `KeyRerotationPress` ([source](kvpress/presses/key_rerotation_press.py)): rerotate pruned keys to have continuous RoPE embeddings +- `AfterglowPress` ([source](kvpress/presses/afterglow_press.py), [paper](https://www.alphaxiv.org/abs/2610.afterglow-query-aware-kv-cache-eviction)): propagate the scores of any `ScorerPress` forward with a decay, so that the whole answer span after a high-scoring token is kept (query-aware compression) - `ChunkKVPress` ([source](kvpress/presses/chunkkv_press.py), [paper](https://arxiv.org/abs/2502.00299)): compress by selecting important chunks, preserving semantic coherence - `EntropyGatedChunkKVPress` ([source](kvpress/presses/entropy_gated_chunkkv_press.py)): similar to `ChunkKVPress`, but reduces the length of chunks with high scores but low entropy, freeing budget for more chunks. - `ChunkPress` ([source](kvpress/presses/chunk_press.py), [paper](https://direct.mit.edu/tacl/article/doi/10.1162/tacl_a_00716/125280)): compress the KV cache on each sequence chunk separately. This can yield to more uniform compression across long sequences diff --git a/evaluation/evaluate_registry.py b/evaluation/evaluate_registry.py index 9552e9b0..7a47109d 100644 --- a/evaluation/evaluate_registry.py +++ b/evaluation/evaluate_registry.py @@ -14,6 +14,7 @@ from kvpress import ( AdaKVPress, + AfterglowPress, BlockPress, CAMPress, ChunkKVPress, @@ -42,6 +43,7 @@ ObservedAttentionPress, PyramidKVPress, QFilterPress, + QueryAwareKVzipPlusPress, RandomPress, RestoreKVPress, SnapKVPress, @@ -139,4 +141,8 @@ "merging_snapkv": MergingPress(SnapKVPress()), "merging_expected_attention": MergingPress(ExpectedAttentionPress(epsilon=1e-2)), "merging_kvzap_mlp": MergingPress(KVzapPress(model_type="mlp")), + "snapkv_afterglow": AfterglowPress(SnapKVPress()), + "adakv_snapkv_afterglow": AdaKVPress(AfterglowPress(SnapKVPress())), + "critical_adakv_snapkv_afterglow": CriticalAdaKVPress(AfterglowPress(SnapKVPress())), + "kvzip_plus_query_aware": QueryAwareKVzipPlusPress(), } diff --git a/kvpress/__init__.py b/kvpress/__init__.py index 40b7cb58..3e8fa683 100644 --- a/kvpress/__init__.py +++ b/kvpress/__init__.py @@ -5,6 +5,7 @@ from kvpress.attention_patch import patch_attention_functions from kvpress.pipeline import KVPressTextGenerationPipeline from kvpress.presses.adakv_press import AdaKVPress +from kvpress.presses.afterglow_press import AfterglowPress from kvpress.presses.base_press import SUPPORTED_MODELS, BasePress from kvpress.presses.block_press import BlockPress from kvpress.presses.cam_press import CAMPress @@ -42,6 +43,7 @@ from kvpress.presses.prefill_decoding_press import PrefillDecodingPress from kvpress.presses.pyramidkv_press import PyramidKVPress from kvpress.presses.qfilter_press import QFilterPress +from kvpress.presses.query_aware_kvzip_plus_press import QueryAwareKVzipPlusPress from kvpress.presses.random_press import RandomPress from kvpress.presses.restorekv_press import RestoreKVPress from kvpress.presses.scorer_press import ScorerPress @@ -103,4 +105,6 @@ "CapPress", "LUKVPress", "EntropyGatedChunkKVPress", + "AfterglowPress", + "QueryAwareKVzipPlusPress", ] diff --git a/kvpress/presses/afterglow_press.py b/kvpress/presses/afterglow_press.py new file mode 100644 index 00000000..3de1d97c --- /dev/null +++ b/kvpress/presses/afterglow_press.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Optional + +import torch + +from kvpress.presses.scorer_press import ScorerPress + + +def afterglow(scores: torch.Tensor, gamma: float, n_sink: int, gate: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + s'(t) = max(s(t), gate(t) * max_{n_sink <= j < t} gamma^(t - j) s(j)) along the last dimension, computed as one + cumulative max in log space. Scores must be >= 0. The first n_sink tokens are not propagation sources. + """ + if (scores < 0).any(): + raise ValueError("Afterglow requires scores >= 0, e.g. use scores.exp() for scorers with negative scores") + positions = torch.arange(scores.shape[-1], device=scores.device, dtype=torch.float32) * math.log(gamma) + sources = scores.float().clamp_min(1e-30).log() - positions + sources[..., :n_sink] = float("-inf") + inherited = (torch.cummax(sources, dim=-1).values + positions).exp() + if gate is not None: + inherited = inherited * gate + return torch.maximum(scores, inherited.to(scores.dtype)) + + +class AfterglowPress(ScorerPress): + """ + Afterglow: forward score propagation for query-aware compression. + + Query-aware scorers concentrate their scores on the first tokens of the span the question asks about, so the rest + of the span is evicted. Afterglow lets each token inherit the decayed scores of the tokens before it, per layer and + KV head: s'(t) = max(s(t), max_{j < t} gamma^(t - j) s(j)). Use it inside AdaKVPress or CriticalAdaKVPress, e.g. + CriticalAdaKVPress(AfterglowPress(SnapKVPress())), as CriticalKV sets the scores it protects to the largest float. + + Based on Afterglow (https://www.alphaxiv.org/abs/2610.afterglow-query-aware-kv-cache-eviction). + + Parameters + ---------- + press : ScorerPress + The scorer whose scores are propagated. Its scores must be >= 0. + gamma : float, default=0.95 + Decay of the propagated scores per token. + n_sink : int, default=4 + Number of initial tokens (attention sinks) that are not propagation sources. + """ + + def __init__(self, press: ScorerPress, gamma: float = 0.95, n_sink: int = 4): + self.press = press + self.gamma = gamma + self.n_sink = n_sink + assert isinstance(self.press, ScorerPress), "AfterglowPress requires a ScorerPress as input" + assert 0 < self.gamma < 1, "gamma must be in (0, 1)" + + def post_init_from_model(self, model): + self.press.post_init_from_model(model) + + @property # type: ignore[misc] + def compression_ratio(self): + return self.press.compression_ratio + + @compression_ratio.setter + def compression_ratio(self, value): + self.press.compression_ratio = value + + def score(self, module, hidden_states, keys, values, attentions, kwargs): + scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs) + return afterglow(scores, self.gamma, self.n_sink) diff --git a/kvpress/presses/query_aware_kvzip_plus_press.py b/kvpress/presses/query_aware_kvzip_plus_press.py new file mode 100644 index 00000000..64c408f1 --- /dev/null +++ b/kvpress/presses/query_aware_kvzip_plus_press.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Generator + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PreTrainedModel + +from kvpress.presses.afterglow_press import afterglow +from kvpress.presses.base_press import BasePress, is_prefilling +from kvpress.presses.kvzip_press import KVzipPress +from kvpress.utils import extract_keys_and_values, get_query_states + +LM_HEAD_CHUNK_SIZE = 1024 + + +@dataclass +class QueryAwareKVzipPlusPress(KVzipPress): + """ + Single-pass, query-aware KVzip+ with Afterglow and a predictability gate. + + Instead of KVzip's context reconstruction, the KVzip+ scores are computed during prefill from the attention of the + last `window_size` tokens, which contain the question (query-aware compression). Afterglow then lets each token + inherit the decayed scores of the tokens before it (see AfterglowPress), scaled by u(t)^alpha with + u(t) = 1 - p(x_t | x_bhgqk", queries.float(), keys.float()) / math.sqrt(head_dim) + self._mask_causal(attn_weights, window_size) + attn_weights = attn_weights.softmax(dim=-1) + + # KVzip+ normalization: divide by ||h|| of the query, multiply by ||W_O v|| of the key + attn_weights = attn_weights / hidden_states[:, -window_size:].float().norm(dim=-1)[:, None, None, :, None] + attn_weights = attn_weights * self._compute_value_output_norm(module, values).float()[:, :, :, None, :] + + if self.score_val is None: + n_layers = module.config.num_hidden_layers + self.score_val = torch.zeros(n_layers, bsz, num_key_value_heads, k_len, device=keys.device) + self.score_val[module.layer_idx] = attn_weights.amax(dim=(2, 3)) # max over group, queries + return output + + @torch.no_grad() + def unpredictability(self, model: PreTrainedModel, hidden_states: torch.Tensor, input_ids: torch.Tensor): + """u(t) = 1 - p(x_t | x_