From 7797071184101dc8f6a73d199ba8202f71c5c3e0 Mon Sep 17 00:00:00 2001 From: Shahar Ben-Ishay Date: Thu, 13 Aug 2026 18:58:15 +0000 Subject: [PATCH] Add EntropyGatedChunkKVPress Co-authored-by: Liran Azran Signed-off-by: Shahar Ben-Ishay --- README.md | 1 + evaluation/evaluate_registry.py | 2 + kvpress/__init__.py | 2 + .../presses/entropy_gated_chunkkv_press.py | 183 ++++++++++++++++++ .../test_entropy_gated_chunkkv_press.py | 135 +++++++++++++ tests/presses/test_presses.py | 23 ++- 6 files changed, 345 insertions(+), 1 deletion(-) create mode 100644 kvpress/presses/entropy_gated_chunkkv_press.py create mode 100644 tests/presses/test_entropy_gated_chunkkv_press.py diff --git a/README.md b/README.md index 58d12f34d..686f5d278 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,7 @@ Finally we provide wrapper presses that can be combined with other presses: - `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 - `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)): like `ChunkKVPress`, but reduces important chunks whose score mass is concentrated (using entropy from information theory) in a few tokens to their top-`rescue_size` tokens, reallocating the freed budget to 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 - `CriticalKVPress` and `CriticalAdaKVPress` ([source](kvpress/presses/criticalkv_press.py), [paper](https://arxiv.org/abs/2502.03805)): refine the scores using the L1 norm of Wo @ values, coupled with a two-stage selection. - `BlockPress` ([source](kvpress/presses/block_press.py), [paper](https://arxiv.org/abs/2504.15364)): segment input sequence into non-overlapping blocks and compress iteratively (⚠️ not a true chunked-prefill implementation) diff --git a/evaluation/evaluate_registry.py b/evaluation/evaluate_registry.py index d75cf0b09..1ac0b817f 100644 --- a/evaluation/evaluate_registry.py +++ b/evaluation/evaluate_registry.py @@ -25,6 +25,7 @@ DecodingPress, DMSPress, DuoAttentionPress, + EntropyGatedChunkKVPress, ExpectedAttentionPress, FastKVzipPress, FinchPress, @@ -87,6 +88,7 @@ "cur": CURPress(), "duo_attention": DuoAttentionPress(), "duo_attention_on_the_fly": DuoAttentionPress(on_the_fly_scoring=True), + "entropy_gated_chunkkv": EntropyGatedChunkKVPress(press=SnapKVPress(), chunk_length=10, rescue_size=4), "expected_attention": AdaKVPress(ExpectedAttentionPress(epsilon=1e-2)), "fastkvzip": FastKVzipPress(), "finch": FinchPress(), diff --git a/kvpress/__init__.py b/kvpress/__init__.py index 454d986dc..2606d4d66 100644 --- a/kvpress/__init__.py +++ b/kvpress/__init__.py @@ -19,6 +19,7 @@ from kvpress.presses.decoding_press import DecodingPress from kvpress.presses.dms_press import DMSPress from kvpress.presses.duo_attention_press import DuoAttentionPress +from kvpress.presses.entropy_gated_chunkkv_press import EntropyGatedChunkKVPress from kvpress.presses.expected_attention_press import ExpectedAttentionPress from kvpress.presses.expected_attention_with_stats import ExpectedAttentionStatsPress from kvpress.presses.fastkvzip_press import FastKVzipPress @@ -97,4 +98,5 @@ "MergingPress", "CapPress", "LUKVPress", + "EntropyGatedChunkKVPress", ] diff --git a/kvpress/presses/entropy_gated_chunkkv_press.py b/kvpress/presses/entropy_gated_chunkkv_press.py new file mode 100644 index 000000000..85fe08518 --- /dev/null +++ b/kvpress/presses/entropy_gated_chunkkv_press.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math +from dataclasses import dataclass +from typing import Optional + +import torch +from torch import nn + +from kvpress.presses.base_press import BasePress +from kvpress.presses.scorer_press import ScorerPress + + +@dataclass +class EntropyGatedChunkKVPress(BasePress): + """ + EntropyGatedChunkKV: chunk selection gated by within-chunk score entropy. + + Extends ChunkKVPress, which keeps or drops every chunk as a whole. A chunk whose + importance comes from a single high-scoring token therefore spends chunk_length + cache slots to preserve one useful token. This press measures the normalized + entropy of the token scores inside each chunk: coherent chunks (high entropy) are + kept whole, while important but spiky chunks (low entropy) are reduced to their + top rescue_size tokens, and the freed budget is spent on further chunks. The + number of retained tokens is exactly (1 - compression_ratio) * kv_len, matching + the budget of ChunkKVPress. + + Based on ChunkKV (https://arxiv.org/abs/2502.00299). + + Parameters + ---------- + press : ScorerPress + The underlying scoring method used to compute global importance scores. + chunk_length : int, default=10 + Length of each chunk for token selection. Shorter than the ChunkKVPress default + of 20: a finer granularity gives the gate more chunks to reallocate budget + between, which is where the gain comes from. + rescue_size : int, default=4 + Number of tokens kept from an important but spiky chunk. + entropy_threshold : float or None, default=None + Spikiness cutoff on the normalized within-chunk entropy, in [0, 1]. A chunk is + spiky when its entropy falls below this value. If None, the per-example median + entropy over all chunks is used. + + Notes + ----- + Chunk and token selection is shared across heads and computed from batch element 0, + the same convention as ChunkKVPress; it is intended for the batch-size-1 context + compression performed by the kvpress pipeline. Token scores are assumed to be + non-negative (as produced by e.g. SnapKVPress) and are clamped before the entropy + is computed. + """ + + press: ScorerPress + chunk_length: int = 10 + rescue_size: int = 4 + entropy_threshold: Optional[float] = None + + def __post_init__(self): + assert isinstance(self.press, ScorerPress), "EntropyGatedChunkKVPress requires a ScorerPress as input" + + def post_init_from_model(self, model): + self.press.post_init_from_model(model) + + @property + def compression_ratio(self): + return self.press.compression_ratio + + @compression_ratio.setter + def compression_ratio(self, value): + self.press.compression_ratio = value + + def compress( + self, + module: nn.Module, + hidden_states: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + attentions: torch.Tensor, + kwargs: dict, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.press.compression_ratio == 0: + return keys, values + assert attentions is None, "EntropyGatedChunkKVPress does not support attentions." + + eps = 1e-8 + kv_len = keys.shape[2] + c = self.chunk_length + + # Head-summed, non-negative per-token scores (batch element 0). + global_scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs) + tok = global_scores.sum(dim=1)[0].clamp(min=0).float() # (kv_len,) + + budget = max(1, int(kv_len * (1 - self.press.compression_ratio))) + if budget >= kv_len: + return keys, values + + # 1. Per-chunk semantic score S and normalized entropy H_tilde. + n_chunks = math.ceil(kv_len / c) + bounds = [(i * c, min(i * c + c, kv_len)) for i in range(n_chunks)] + n_complete = kv_len // c + remaining_tokens = kv_len % c + + # Per-chunk statistics are computed vectorized rather than in a Python loop, which + # would launch O(n_chunks) tiny kernels per forward pass. Complete chunks all hold + # exactly c tokens, so reshaping to (n_complete, c) makes each row one chunk and the + # row-wise reductions give its mean and normalized Shannon entropy. + X = tok[: n_complete * c].view(n_complete, c) + s_scores = X.mean(dim=1) + if c > 1: + p = X / (X.sum(dim=1, keepdim=True) + eps) + h = -(p * (p + eps).log()).sum(dim=1) + ht = (h / math.log(c)).clamp(0.0, 1.0) + else: + # Entropy is undefined for a single token, so such a chunk is treated as spiky, + # as for a length-1 trailing chunk below. Normalizing by log(1) = 0 would divide + # by zero here, which is why this case is handled separately. + ht = torch.zeros(n_complete, device=tok.device) + + # The trailing partial chunk does not fit the reshape and is handled separately. + # Entropy is undefined for a single token, so such a chunk is treated as spiky. + if remaining_tokens > 0: + ts = tok[n_complete * c :] + s_tail = ts.mean().unsqueeze(0) + if remaining_tokens >= 2: + pr = ts / (ts.sum() + eps) + hr = -(pr * (pr + eps).log()).sum() + ht_tail = (hr / math.log(remaining_tokens)).clamp(0.0, 1.0).unsqueeze(0) + else: + ht_tail = torch.zeros(1, device=tok.device) + s_scores = torch.cat([s_scores, s_tail]) + ht = torch.cat([ht, ht_tail]) + + med = s_scores.median() + if self.entropy_threshold is None: + tau = ht.median() + else: + tau = torch.tensor(float(self.entropy_threshold), device=tok.device) + + # 2. Greedy pass over chunks in decreasing semantic score. + # The loop is inherently sequential because the budget is consumed in order. Both + # gating masks are therefore computed vectorized and moved to CPU lists once: reading + # a GPU scalar per iteration would force a synchronize and serialize the loop. The + # topk calls stay on the GPU tensor so their tie-breaking is unchanged. + important_all = (s_scores >= med).tolist() + spiky_all = (ht < tau).tolist() + keep = torch.zeros(kv_len, dtype=torch.bool, device=tok.device) + for i in torch.argsort(s_scores, descending=True).tolist(): + if budget <= 0: + break + s, e = bounds[i] + n_i = e - s + ts = tok[s:e] + + if important_all[i] and spiky_all[i]: + # Important but spiky: keep only the highest-scoring tokens of the chunk. + r = min(self.rescue_size, budget, n_i) + keep[torch.topk(ts, r).indices + s] = True + budget -= r + elif n_i <= budget: + # Coherent chunk that fits in the remaining budget: keep it whole. + keep[s:e] = True + budget -= n_i + else: + # Last chunk to be considered: keep as much of it as the budget allows. + keep[torch.topk(ts, budget).indices + s] = True + budget = 0 + + # 3. Reducing spiky chunks may leave budget unspent. Top up with the highest-scoring + # remaining tokens so that exactly (1 - compression_ratio) * kv_len tokens are kept. + if budget > 0: + leftover = (~keep).nonzero(as_tuple=False).squeeze(-1) + if leftover.numel() > 0: + add = min(budget, leftover.numel()) + keep[leftover[torch.topk(tok[leftover], add).indices]] = True + + # 4. Gather the retained keys and values in positional order. + indices = keep.nonzero(as_tuple=False).squeeze(-1).sort()[0] + indices = indices.view(1, 1, -1, 1).expand(keys.shape[0], keys.shape[1], -1, module.head_dim) + keys = keys.gather(2, indices).contiguous() + values = values.gather(2, indices).contiguous() + return keys, values diff --git a/tests/presses/test_entropy_gated_chunkkv_press.py b/tests/presses/test_entropy_gated_chunkkv_press.py new file mode 100644 index 000000000..1d41e9788 --- /dev/null +++ b/tests/presses/test_entropy_gated_chunkkv_press.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +import pytest +import torch +from torch import nn + +from kvpress import EntropyGatedChunkKVPress +from kvpress.presses.scorer_press import ScorerPress + + +@dataclass +class FixedScorer(ScorerPress): + """Scorer returning pre-set token scores, so chunk statistics are fully controlled.""" + + scores: torch.Tensor = field(default_factory=lambda: torch.empty(0)) + + def score(self, module, hidden_states, keys, values, attentions, kwargs): + return self.scores + + +class DummyAttention(nn.Module): + def __init__(self, head_dim): + super().__init__() + self.head_dim = head_dim + + +def run_press(scores, press): + """Run compress on keys whose values encode their position, and return the kept positions.""" + kv_len = scores.shape[2] + n_heads, head_dim = scores.shape[1], 4 + positions = torch.arange(kv_len, dtype=torch.float32) + keys = positions.view(1, 1, kv_len, 1).expand(1, n_heads, kv_len, head_dim).contiguous() + values = keys.clone() + out_keys, out_values = press.compress(DummyAttention(head_dim), None, keys, values, None, {}) + assert torch.equal(out_keys, out_values) + return out_keys[0, 0, :, 0].long().tolist() + + +def expected_budget(kv_len, compression_ratio): + """The token budget the press targets, matching ChunkKVPress.""" + return max(1, int(kv_len * (1 - compression_ratio))) + + +def spiky_scores(n_chunks, chunk_length, n_heads=2): + """One dominant needle per chunk, the rest near-zero: every chunk is important and spiky.""" + kv_len = n_chunks * chunk_length + scores = torch.full((1, n_heads, kv_len), 0.01) + needles = [i * chunk_length + (i % chunk_length) for i in range(n_chunks)] + for rank, idx in enumerate(needles): + scores[0, :, idx] = 10.0 + rank # distinct so chunk ordering is deterministic + return scores, needles + + +@pytest.mark.parametrize("compression_ratio", [0.1, 0.25, 0.5, 0.75, 0.9]) +@pytest.mark.parametrize("chunk_length", [1, 4, 10, 20]) +@pytest.mark.parametrize("kv_len", [100, 251]) +def test_retains_exact_budget(compression_ratio, chunk_length, kv_len): + """The retained token count matches ChunkKVPress's budget exactly, including partial chunks.""" + torch.manual_seed(0) + scores = torch.rand(1, 2, kv_len) + press = EntropyGatedChunkKVPress( + press=FixedScorer(compression_ratio=compression_ratio, scores=scores), + chunk_length=chunk_length, + rescue_size=4, + ) + kept = run_press(scores, press) + assert len(kept) == expected_budget(kv_len, compression_ratio) + assert kept == sorted(set(kept)), "kept positions must be unique and in positional order" + + +def test_spiky_chunk_is_reduced_to_rescue_size(): + """An important but spiky chunk keeps only its needle, not all chunk_length tokens.""" + chunk_length, n_chunks, rescue_size = 10, 20, 1 + scores, needles = spiky_scores(n_chunks, chunk_length) + kv_len = n_chunks * chunk_length + # A budget of ~2 chunks: ChunkKV would spend it keeping 2 chunks whole, EG-ChunkKV rescues needles. + press = EntropyGatedChunkKVPress( + press=FixedScorer(compression_ratio=0.9, scores=scores), + chunk_length=chunk_length, + rescue_size=rescue_size, + entropy_threshold=0.5, + ) + kept = run_press(scores, press) + assert len(kept) == expected_budget(kv_len, 0.9) + + # Every rescued needle survives, and the highest-scoring chunks are not kept whole. + kept_set = set(kept) + top_needles = sorted(needles, key=lambda i: -float(scores[0, 0, i]))[:20] + assert kept_set.issuperset(top_needles[:10]), "the strongest needles must be retained" + per_chunk = [len([k for k in kept if k // chunk_length == c]) for c in range(n_chunks)] + assert max(per_chunk) < chunk_length, "no spiky chunk should be kept whole" + + +def test_entropy_threshold_degenerate_limits(): + """threshold=0 disables rescuing (chunks kept whole); threshold=1 rescues every important chunk.""" + chunk_length, n_chunks = 10, 20 + scores, _ = spiky_scores(n_chunks, chunk_length) + kwargs = dict(chunk_length=chunk_length, rescue_size=1) + + never_spiky = EntropyGatedChunkKVPress( + press=FixedScorer(compression_ratio=0.9, scores=scores), entropy_threshold=0.0, **kwargs + ) + always_spiky = EntropyGatedChunkKVPress( + press=FixedScorer(compression_ratio=0.9, scores=scores), entropy_threshold=1.0, **kwargs + ) + kept_whole = run_press(scores, never_spiky) + kept_rescued = run_press(scores, always_spiky) + + # Both spend exactly the same budget + budget = expected_budget(n_chunks * chunk_length, 0.9) + assert len(kept_whole) == len(kept_rescued) == budget + + # With rescuing disabled the budget goes to whole chunks; with it enabled the same budget + # is spread over strictly more chunks, which is the point of the press. + chunks_whole = len({k // chunk_length for k in kept_whole}) + chunks_rescued = len({k // chunk_length for k in kept_rescued}) + assert chunks_whole == 2, "without rescuing, a 20-token budget buys exactly 2 whole chunks" + assert chunks_rescued > chunks_whole + + +def test_compression_ratio_is_delegated_to_inner_press(): + """The wrapper exposes and forwards the inner ScorerPress's compression ratio.""" + inner = FixedScorer(compression_ratio=0.3, scores=torch.rand(1, 2, 64)) + press = EntropyGatedChunkKVPress(press=inner, chunk_length=8) + assert press.compression_ratio == 0.3 + press.compression_ratio = 0.7 + assert inner.compression_ratio == 0.7 + + +def test_requires_scorer_press(): + with pytest.raises(AssertionError): + EntropyGatedChunkKVPress(press="not-a-press") # type: ignore[arg-type] diff --git a/tests/presses/test_presses.py b/tests/presses/test_presses.py index d977bd1d2..9da801960 100644 --- a/tests/presses/test_presses.py +++ b/tests/presses/test_presses.py @@ -15,6 +15,7 @@ CriticalAdaKVPress, CriticalKVPress, DMSPress, + EntropyGatedChunkKVPress, FastKVzipPress, KeyRerotationPress, KnormPress, @@ -61,6 +62,18 @@ def test_chunkkv_press(unit_test_model): # noqa: F811 assert cache.get_seq_length() == 128 +def test_entropy_gated_chunkkv_press(unit_test_model): # noqa: F811 + press = SnapKVPress(compression_ratio=0.5) + for chunk_length in [2, 4, 8, 128]: + for rescue_size in [1, 4]: + composed_press = EntropyGatedChunkKVPress(press=press, chunk_length=chunk_length, rescue_size=rescue_size) + with composed_press(unit_test_model): + input_ids = torch.randint(0, 1024, (1, 256), device=unit_test_model.device) + cache = DynamicCache() + unit_test_model(input_ids, past_key_values=cache).past_key_values + assert cache.get_seq_length() == 128 + + @pytest.mark.parametrize("press_dict", default_presses) @pytest.mark.parametrize( "wrapper_press", @@ -74,6 +87,7 @@ def test_chunkkv_press(unit_test_model): # noqa: F811 CriticalAdaKVPress, DMSPress, MergingPress, + EntropyGatedChunkKVPress, ], ) def test_presses_run(unit_test_model, press_dict, wrapper_press): # noqa: F811 @@ -92,7 +106,14 @@ def test_presses_run(unit_test_model, press_dict, wrapper_press): # noqa: F811 return elif issubclass( wrapper_press, - (KeyRerotationPress, AdaKVPress, CriticalKVPress, CriticalAdaKVPress, MergingPress), + ( + KeyRerotationPress, + AdaKVPress, + CriticalKVPress, + CriticalAdaKVPress, + MergingPress, + EntropyGatedChunkKVPress, + ), ): press = wrapper_press(press=press) elif issubclass(wrapper_press, ChunkPress):