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.

Author: Liranitz

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)

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

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)

tests/presses/test_entropy_gated_chunkkv_press.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",