Add AfterglowPress and QueryAwareKVzipPlusPress

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

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(),
 }

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",
 ]

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)

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_<t) from the prefill's own next-token distribution, so that predictable continuations inherit
+    less. As in KVzip, the budget is shared across layers and heads.
+
+    Based on Afterglow (https://www.alphaxiv.org/abs/2610.afterglow-query-aware-kv-cache-eviction).
+
+    Parameters
+    ----------
+    compression_ratio : float, default=0.0
+        Fraction of key-value pairs to remove during compression.
+    layerwise : bool, default=False
+        Whether to enable uniform compression ratios across layers.
+    n_sink : int, default=4
+        Number of initial tokens to preserve as attention sinks. They are not Afterglow sources.
+    window_size : int, default=64
+        Number of last tokens of the context used as queries (should cover the question).
+    gamma : float, default=0.95
+        Decay of the propagated scores per token.
+    alpha : float, default=0.5
+        Exponent of the predictability gate u(t)^alpha.
+    """
+
+    kvzip_plus_normalization: bool = True
+    window_size: int = 64
+    gamma: float = 0.95
+    alpha: float = 0.5
+
+    def __post_init__(self):
+        assert 0 <= self.compression_ratio < 1, "Compression ratio must be between 0 and 1"
+        assert 0 < self.gamma < 1, "gamma must be in (0, 1)"
+        assert self.kvzip_plus_normalization, "QueryAwareKVzipPlusPress always uses the KVzip+ normalization"
+        self._reset_internal_parameters()
+
+    def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dict, output: list):
+        """Only score during prefill; the cache is compressed once all layers are scored (see __call__)."""
+        hidden_states = kwargs["hidden_states"]
+        if self.compression_ratio == 0 or not is_prefilling(kwargs["cache_position"], hidden_states.shape[1]):
+            return output
+
+        keys, values = extract_keys_and_values(kwargs["past_key_values"], module.layer_idx)
+        bsz, num_key_value_heads, k_len, head_dim = keys.shape
+        num_key_value_groups = module.config.num_attention_heads // num_key_value_heads
+        window_size = min(self.window_size, k_len)
+
+        cos, sin = kwargs["position_embeddings"]
+        position_embeddings = (cos[:, -window_size:], sin[:, -window_size:])
+        queries = get_query_states(module, hidden_states[:, -window_size:], position_embeddings)
+        queries = queries.view(bsz, num_key_value_heads, num_key_value_groups, window_size, head_dim)
+
+        attn_weights = torch.einsum("bhgqd,bhkd->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_<t) from the final hidden states of the prefill, max-pooled over 3 tokens."""
+        u = torch.ones(input_ids.shape, device=hidden_states.device)
+        for start in range(0, input_ids.shape[1] - 1, LM_HEAD_CHUNK_SIZE):
+            end = min(start + LM_HEAD_CHUNK_SIZE, input_ids.shape[1] - 1)
+            logits = model.lm_head(hidden_states[:, start:end]).float()
+            target = logits.gather(-1, input_ids[:, start + 1 : end + 1, None])[..., 0]
+            u[:, start + 1 : end + 1] = 1 - (target - logits.logsumexp(dim=-1)).exp()
+        return F.max_pool1d(u[:, None], 3, stride=1, padding=1)[:, 0]
+
+    @contextmanager
+    def __call__(self, model: PreTrainedModel) -> Generator:
+        language_model = model.model.language_model if hasattr(model.model, "language_model") else model.model
+        captured = {}
+
+        def capture_input_ids(module, args, kwargs):
+            captured["input_ids"] = kwargs["input_ids"]
+
+        def capture_hidden_states(module, input, output):
+            captured["hidden_states"] = output
+
+        hooks = [
+            language_model.register_forward_pre_hook(capture_input_ids, with_kwargs=True),
+            language_model.norm.register_forward_hook(capture_hidden_states),
+        ]
+        try:
+            with BasePress.__call__(self, model):
+                yield
+            if self.score_val is not None:
+                gate = self.unpredictability(model, captured["hidden_states"], captured["input_ids"])
+                self.score_val = afterglow(self.score_val, self.gamma, self.n_sink, gate.pow(self.alpha)[:, None])
+                self.compress_post(model)
+        finally:
+            for hook in hooks:
+                hook.remove()
+            self._reset_internal_parameters()

tests/default_presses.py

@@ -25,6 +25,7 @@
     NonCausalAttnPress,
     PyramidKVPress,
     QFilterPress,
+    QueryAwareKVzipPlusPress,
     RandomPress,
     RestoreKVPress,
     SimLayerKVPress,
@@ -202,4 +203,11 @@ def post_init_from_model(self, model):
     },
     {"cls": CapPress, "kwargs": [{"compression_ratio": 0.5}, {"compression_ratio": 0.8}]},
     {"cls": TestLUKVPress, "kwargs": [{"compression_ratio": 0.5}, {"compression_ratio": 0.8}]},
+    {
+        "cls": QueryAwareKVzipPlusPress,
+        "kwargs": [
+            {"compression_ratio": 0.2, "window_size": 2},
+            {"compression_ratio": 0.8, "window_size": 2},
+        ],
+    },
 ]

tests/presses/test_afterglow_press.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+from transformers import DynamicCache
+
+from kvpress import AdaKVPress, AfterglowPress, CriticalAdaKVPress, KnormPress, SnapKVPress
+from kvpress.presses.afterglow_press import afterglow
+from tests.fixtures import unit_test_model  # noqa: F401
+
+
+def test_afterglow():
+    scores = torch.rand(2, 3, 50) ** 4
+    expected = scores.clone()
+    for t in range(50):
+        for j in range(4, t):
+            expected[..., t] = torch.maximum(expected[..., t], 0.9 ** (t - j) * scores[..., j])
+    assert torch.allclose(afterglow(scores, gamma=0.9, n_sink=4), expected, rtol=1e-4)
+    with pytest.raises(ValueError):
+        afterglow(-scores, gamma=0.9, n_sink=4)
+
+
+@pytest.mark.parametrize("wrapper", [None, AdaKVPress, CriticalAdaKVPress])
+def test_afterglow_press(unit_test_model, wrapper):  # noqa: F811
+    press = AfterglowPress(SnapKVPress(window_size=2))
+    press = press if wrapper is None else wrapper(press)
+    press.compression_ratio = 0.5
+    with press(unit_test_model):
+        input_ids = torch.randint(0, 1024, (1, 128), device=unit_test_model.device)
+        cache = DynamicCache()
+        unit_test_model(input_ids, past_key_values=cache)
+    if wrapper is None:
+        assert cache.get_seq_length() == 64
+    else:
+        n_masked = sum(len(layer.self_attn.masked_key_indices[0]) for layer in unit_test_model.model.layers)
+        assert n_masked == unit_test_model.config.num_hidden_layers * unit_test_model.config.num_key_value_heads * 64
+
+
+def test_afterglow_press_negative_scores(unit_test_model):  # noqa: F811
+    with pytest.raises(ValueError):
+        with AfterglowPress(KnormPress(compression_ratio=0.5))(unit_test_model):
+            unit_test_model(torch.randint(0, 1024, (1, 128), device=unit_test_model.device))