Add EntropyGatedChunkKVPress

Co-authored-by: Liran Azran <liran.azr90@gmail.com>
Signed-off-by: Shahar Ben-Ishay <shahar.benishay@gmail.com>

Author: ShaharBenIshay

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)

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

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

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

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]