From 7797071184101dc8f6a73d199ba8202f71c5c3e0 Mon Sep 17 00:00:00 2001 From: Shahar Ben-Ishay Date: Thu, 13 Aug 2026 18:58:15 +0000 Subject: [PATCH 1/6] 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): From 9301ce0dc4b3789e617a6e48332c3ea37e91be1d Mon Sep 17 00:00:00 2001 From: Liran Azran Date: Tue, 18 Aug 2026 08:06:43 +0300 Subject: [PATCH 2/6] Refactor EntropyGatedChunkKVPress to subclass ChunkKVPress; test/registry cleanup - kvpress/presses/entropy_gated_chunkkv_press.py: subclass ChunkKVPress instead of BasePress, dropping the inherited press field, __post_init__, post_init_from_model, and compression_ratio property/setter; keep chunk_length default at 10. Remove the redundant explanatory comments and hoist the epsilon to a module-level EPSILON constant. - tests/presses/test_entropy_gated_chunkkv_press.py: deleted. The dedicated test duplicated coverage already provided by the test_presses_run wrapper matrix. - tests/presses/test_presses.py: removed the redundant test_entropy_gated_chunkkv_press function for the same reason; the press stays covered via the EntropyGatedChunkKVPress entry in the wrapper_press matrix. - evaluation/evaluate_registry.py: dropped the explicit chunk_length=10, rescue_size=4 from the registry entry; both are defaults on the press now, so EntropyGatedChunkKVPress(press=SnapKVPress()) is enough. - README.md: tightened the one-line description to match the other press entries. Signed-off-by: Liran Azran --- README.md | 2 +- evaluation/evaluate_registry.py | 2 +- .../presses/entropy_gated_chunkkv_press.py | 43 ++---- .../test_entropy_gated_chunkkv_press.py | 135 ------------------ tests/presses/test_presses.py | 12 -- 5 files changed, 10 insertions(+), 184 deletions(-) delete mode 100644 tests/presses/test_entropy_gated_chunkkv_press.py diff --git a/README.md b/README.md index 686f5d278..65f0cf84f 100644 --- a/README.md +++ b/README.md @@ -140,7 +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 +- `EntropyGatedChunkKVPress` ([source](kvpress/presses/entropy_gated_chunkkv_press.py)): similar to `ChunkKVPress` but reduces the chunk length for chunks with high scores but low entropy - `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 1ac0b817f..f39646b33 100644 --- a/evaluation/evaluate_registry.py +++ b/evaluation/evaluate_registry.py @@ -88,7 +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), + "entropy_gated_chunkkv": EntropyGatedChunkKVPress(press=SnapKVPress()), "expected_attention": AdaKVPress(ExpectedAttentionPress(epsilon=1e-2)), "fastkvzip": FastKVzipPress(), "finch": FinchPress(), diff --git a/kvpress/presses/entropy_gated_chunkkv_press.py b/kvpress/presses/entropy_gated_chunkkv_press.py index 85fe08518..d7bcde2a8 100644 --- a/kvpress/presses/entropy_gated_chunkkv_press.py +++ b/kvpress/presses/entropy_gated_chunkkv_press.py @@ -8,12 +8,13 @@ import torch from torch import nn -from kvpress.presses.base_press import BasePress -from kvpress.presses.scorer_press import ScorerPress +from kvpress.presses.chunkkv_press import ChunkKVPress + +EPSILON = 1e-8 @dataclass -class EntropyGatedChunkKVPress(BasePress): +class EntropyGatedChunkKVPress(ChunkKVPress): """ EntropyGatedChunkKV: chunk selection gated by within-chunk score entropy. @@ -52,25 +53,10 @@ class EntropyGatedChunkKVPress(BasePress): 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, @@ -84,7 +70,6 @@ def compress( return keys, values assert attentions is None, "EntropyGatedChunkKVPress does not support attentions." - eps = 1e-8 kv_len = keys.shape[2] c = self.chunk_length @@ -102,30 +87,22 @@ def compress( 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) + p = X / (X.sum(dim=1, keepdim=True) + EPSILON) + h = -(p * (p + EPSILON).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() + pr = ts / (ts.sum() + EPSILON) + hr = -(pr * (pr + EPSILON).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) @@ -139,10 +116,6 @@ def compress( 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) diff --git a/tests/presses/test_entropy_gated_chunkkv_press.py b/tests/presses/test_entropy_gated_chunkkv_press.py deleted file mode 100644 index 1d41e9788..000000000 --- a/tests/presses/test_entropy_gated_chunkkv_press.py +++ /dev/null @@ -1,135 +0,0 @@ -# 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 9da801960..5eebe0ef1 100644 --- a/tests/presses/test_presses.py +++ b/tests/presses/test_presses.py @@ -62,18 +62,6 @@ 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", From 149cb7b1c7ee45fb41c298a6a6f6e03b8a20b76f Mon Sep 17 00:00:00 2001 From: Liran Azran Date: Tue, 18 Aug 2026 08:23:02 +0300 Subject: [PATCH 3/6] docs: clarify EntropyGatedChunkKVPress description in README Signed-off-by: Liran Azran --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 65f0cf84f..da068ea78 100644 --- a/README.md +++ b/README.md @@ -140,7 +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)): similar to `ChunkKVPress` but reduces the chunk length for chunks with high scores but low entropy +- `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 - `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) From fff1d07bad0a3946094a5956875bd64e70f543f7 Mon Sep 17 00:00:00 2001 From: Liran Azran Date: Tue, 18 Aug 2026 10:37:58 +0300 Subject: [PATCH 4/6] Refine EntropyGatedChunkKVPress: signed-score correctness, naming, cleanup Negative-score correctness: - Drop clamp(min=0) on the per-token scores; ranking, median, argsort and all top-k selection now use the raw scores, so signed scorers (e.g. KeyDiffPress) are ordered correctly instead of silently collapsing to ties. - Compute the within-chunk entropy from a per-chunk min-shift (subtract the chunk minimum only when it is negative) so the scores form a valid distribution. Remove the entropy_threshold kwarg: - Always use the per-example median entropy as the spikiness cutoff, drop the entropy_threshold field, its docstring entry. Guard chunk_length and simplify: - Assert chunk_length > 1 in __post_init__ and remove the c == 1 entropy branch it makes unreachable (the length-1 partial-chunk case is still handled separately). - Remove the redundant budget >= kv_len early return. Naming and readability: - Rename locals to intent-revealing names (chunk_len, scores, chunk_token_scores, chunk_scores, chunk_entropy, score_threshold, high_score_chunks, low_entropy_chunks, low_entropy_chunk_length) and restructure the greedy loop around a single n_kept quantity. Signed-off-by: Liran Azran --- .../presses/entropy_gated_chunkkv_press.py | 121 ++++++++---------- 1 file changed, 56 insertions(+), 65 deletions(-) diff --git a/kvpress/presses/entropy_gated_chunkkv_press.py b/kvpress/presses/entropy_gated_chunkkv_press.py index d7bcde2a8..f0d33eb81 100644 --- a/kvpress/presses/entropy_gated_chunkkv_press.py +++ b/kvpress/presses/entropy_gated_chunkkv_press.py @@ -3,7 +3,6 @@ import math from dataclasses import dataclass -from typing import Optional import torch from torch import nn @@ -23,7 +22,7 @@ class EntropyGatedChunkKVPress(ChunkKVPress): 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 + top low_entropy_chunk_length 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. @@ -37,25 +36,25 @@ class EntropyGatedChunkKVPress(ChunkKVPress): 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 + low_entropy_chunk_length : 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. + compression performed by the kvpress pipeline. Ranking and top-k selection use the + raw scores, so signed scorers (e.g. KeyDiffPress) are ordered correctly; the entropy + gate rebases negative chunks to form a valid distribution but is most meaningful for + non-negative scores (e.g. SnapKVPress). """ chunk_length: int = 10 - rescue_size: int = 4 - entropy_threshold: Optional[float] = None + low_entropy_chunk_length: int = 4 + + def __post_init__(self): + super().__post_init__() + assert self.chunk_length > 1, "EntropyGatedChunkKVPress requires chunk_length > 1" def compress( self, @@ -71,74 +70,66 @@ def compress( assert attentions is None, "EntropyGatedChunkKVPress does not support attentions." kv_len = keys.shape[2] - c = self.chunk_length + chunk_len = 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,) + # Head-summed per-token scores (batch element 0), kept raw so ranking works for signed scorers. + scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs) + scores = scores.sum(dim=1)[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 - - 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) + EPSILON) - h = -(p * (p + EPSILON).log()).sum(dim=1) - ht = (h / math.log(c)).clamp(0.0, 1.0) - else: - ht = torch.zeros(n_complete, device=tok.device) + n_chunks = math.ceil(kv_len / chunk_len) + bounds = [(i * chunk_len, min(i * chunk_len + chunk_len, kv_len)) for i in range(n_chunks)] + n_complete = kv_len // chunk_len + remaining_tokens = kv_len % chunk_len + + chunk_token_scores = scores[: n_complete * chunk_len].view(n_complete, chunk_len) + chunk_scores = chunk_token_scores.mean(dim=1) + if (chunk_token_scores < 0).any(): + chunk_min = chunk_token_scores.amin(dim=1, keepdim=True) + chunk_token_scores = chunk_token_scores - chunk_min.clamp(max=0.0) + p = chunk_token_scores / (chunk_token_scores.sum(dim=1, keepdim=True) + EPSILON) + h = -(p * (p + EPSILON).log()).sum(dim=1) + chunk_entropy = (h / math.log(chunk_len)).clamp(0.0, 1.0) # The trailing partial chunk does not fit the reshape and is handled separately. if remaining_tokens > 0: - ts = tok[n_complete * c :] - s_tail = ts.mean().unsqueeze(0) - if remaining_tokens >= 2: - pr = ts / (ts.sum() + EPSILON) - hr = -(pr * (pr + EPSILON).log()).sum() - ht_tail = (hr / math.log(remaining_tokens)).clamp(0.0, 1.0).unsqueeze(0) + tail_scores = scores[n_complete * chunk_len :] + chunk_scores_tail = tail_scores.mean().unsqueeze(0) + if remaining_tokens == 1: + chunk_entropy_tail = torch.zeros(1, device=scores.device) 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) + if (tail_scores < 0).any(): + tail_scores = tail_scores - tail_scores.min().clamp(max=0.0) + pr = tail_scores / (tail_scores.sum() + EPSILON) + hr = -(pr * (pr + EPSILON).log()).sum() + chunk_entropy_tail = (hr / math.log(remaining_tokens)).clamp(0.0, 1.0).unsqueeze(0) + chunk_scores = torch.cat([chunk_scores, chunk_scores_tail]) + chunk_entropy = torch.cat([chunk_entropy, chunk_entropy_tail]) + score_threshold = chunk_scores.median() + entropy_threshold = chunk_entropy.median() # 2. Greedy pass over chunks in decreasing semantic score. - 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(): + high_score_chunks = (chunk_scores >= score_threshold).tolist() + low_entropy_chunks = (chunk_entropy < entropy_threshold).tolist() + keep = torch.zeros(kv_len, dtype=torch.bool, device=scores.device) + for chunk_idx in torch.argsort(chunk_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]: + start, end = bounds[chunk_idx] + n_kept = min(end - start, budget) + if high_score_chunks[chunk_idx] and low_entropy_chunks[chunk_idx]: # 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 + n_kept = min(n_kept, self.low_entropy_chunk_length) + + if n_kept == end - start: + keep[start:end] = True 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 + top_indices = torch.topk(scores[start:end], n_kept).indices + start + keep[top_indices] = True + budget -= n_kept # 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. @@ -146,7 +137,7 @@ def compress( 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 + keep[leftover[torch.topk(scores[leftover], add).indices]] = True # 4. Gather the retained keys and values in positional order. indices = keep.nonzero(as_tuple=False).squeeze(-1).sort()[0] From 21f6ed634361fe95998a0e272a996fade296919e Mon Sep 17 00:00:00 2001 From: Liran Azran Date: Tue, 18 Aug 2026 11:17:18 +0300 Subject: [PATCH 5/6] - Added assert self.chunk_length > self.low_entropy_chunk_length Signed-off-by: Liran Azran --- kvpress/presses/entropy_gated_chunkkv_press.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/kvpress/presses/entropy_gated_chunkkv_press.py b/kvpress/presses/entropy_gated_chunkkv_press.py index f0d33eb81..8e8d389f5 100644 --- a/kvpress/presses/entropy_gated_chunkkv_press.py +++ b/kvpress/presses/entropy_gated_chunkkv_press.py @@ -55,6 +55,9 @@ class EntropyGatedChunkKVPress(ChunkKVPress): def __post_init__(self): super().__post_init__() assert self.chunk_length > 1, "EntropyGatedChunkKVPress requires chunk_length > 1" + assert self.chunk_length > self.low_entropy_chunk_length, ( + "EntropyGatedChunkKVPress requires chunk_length > low_entropy_chunk_length" + ) def compress( self, From c57ddd423747878798ba08e00e649b918f796b2d Mon Sep 17 00:00:00 2001 From: Liran Azran Date: Tue, 18 Aug 2026 12:18:28 +0300 Subject: [PATCH 6/6] Refine EntropyGatedChunkKVPress: address review comments - Collapse the negative-score rebasing to a single out-of-place line, dropping the (chunk_token_scores < 0).any() guard and the chunk_min temporary. Keep it out-of-place: chunk_token_scores is a view into scores, so an in-place -= would mutate scores and corrupt the later top-k ranking. - Consolidate the two __post_init__ asserts into one enforcing chunk_length > low_entropy_chunk_length >= 1, adding the missing lower bound so an important-but-spiky chunk always keeps at least one token. - Add a blank line before the greedy-pass section and drop the unused S / H_tilde notation from the section-1 comment. Signed-off-by: Liran Azran --- kvpress/presses/entropy_gated_chunkkv_press.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/kvpress/presses/entropy_gated_chunkkv_press.py b/kvpress/presses/entropy_gated_chunkkv_press.py index 8e8d389f5..218a026c7 100644 --- a/kvpress/presses/entropy_gated_chunkkv_press.py +++ b/kvpress/presses/entropy_gated_chunkkv_press.py @@ -54,9 +54,8 @@ class EntropyGatedChunkKVPress(ChunkKVPress): def __post_init__(self): super().__post_init__() - assert self.chunk_length > 1, "EntropyGatedChunkKVPress requires chunk_length > 1" - assert self.chunk_length > self.low_entropy_chunk_length, ( - "EntropyGatedChunkKVPress requires chunk_length > low_entropy_chunk_length" + assert self.chunk_length > self.low_entropy_chunk_length >= 1, ( + "EntropyGatedChunkKVPress requires chunk_length > low_entropy_chunk_length >= 1" ) def compress( @@ -81,7 +80,7 @@ def compress( budget = max(1, int(kv_len * (1 - self.press.compression_ratio))) - # 1. Per-chunk semantic score S and normalized entropy H_tilde. + # 1. Per-chunk score and entropy. n_chunks = math.ceil(kv_len / chunk_len) bounds = [(i * chunk_len, min(i * chunk_len + chunk_len, kv_len)) for i in range(n_chunks)] n_complete = kv_len // chunk_len @@ -89,9 +88,7 @@ def compress( chunk_token_scores = scores[: n_complete * chunk_len].view(n_complete, chunk_len) chunk_scores = chunk_token_scores.mean(dim=1) - if (chunk_token_scores < 0).any(): - chunk_min = chunk_token_scores.amin(dim=1, keepdim=True) - chunk_token_scores = chunk_token_scores - chunk_min.clamp(max=0.0) + chunk_token_scores = chunk_token_scores - chunk_token_scores.amin(dim=1, keepdim=True).clamp(max=0.0) p = chunk_token_scores / (chunk_token_scores.sum(dim=1, keepdim=True) + EPSILON) h = -(p * (p + EPSILON).log()).sum(dim=1) chunk_entropy = (h / math.log(chunk_len)).clamp(0.0, 1.0) @@ -113,6 +110,7 @@ def compress( score_threshold = chunk_scores.median() entropy_threshold = chunk_entropy.median() + # 2. Greedy pass over chunks in decreasing semantic score. high_score_chunks = (chunk_scores >= score_threshold).tolist() low_entropy_chunks = (chunk_entropy < entropy_threshold).tolist()