Add EntropyGatedChunkKVPress

[ShaharBenIshay]

PR description

Adds EntropyGatedChunkKVPress (EG-ChunkKV), as proposed and approved in #261.

ChunkKVPress scores each chunk by its aggregate token importance and then keeps or drops that
chunk as a whole. Because the score is a magnitude statistic, it cannot distinguish a coherent
chunk, whose importance is spread across its tokens, from a spiky one, where a single needle
token holds nearly all the mass and the remaining chunk_length - 1 tokens are filler. Both can
score identically, and for the spiky chunk, keeping it whole spends chunk_length cache slots to
preserve one useful token — which, under a fixed budget, evicts a chunk that would otherwise be kept.

EG-ChunkKV adds a second, shape statistic computed from the same per-token scores ChunkKV already
has: the normalized within-chunk Shannon entropy H̃_i ∈ [0, 1], where 1 means coherent and 0 means
concentrated in a few tokens. Important-but-spiky chunks are reduced to their top-rescue_size
tokens rather than kept whole, and the freed budget is spent on further chunks. The number of
retained tokens is exactly max(1, floor((1 - compression_ratio) * kv_len)), identical to
ChunkKVPress, so the two are budget-matched by construction.

Selection procedure

1. S_i: the mean of the head-summed, non-negative token scores in chunk i. This is the same
   statistic ChunkKVPress ranks chunks by, so both methods order chunks identically.
2. H̃_i: the normalized Shannon entropy of those same token scores within the chunk.
3. The nuance: instead of taking a top-k over chunks up front, we fix the same token budget L
   that ChunkKV uses and walk chunks greedily in decreasing S_i, spending budget as we go:

important_i = S_i >= median(S)        spiky_i = H̃_i < τ        (τ defaults to median(H̃))

if important_i and spiky_i:   keep the chunk's top-`rescue_size` tokens    ← rescue the needle
elif chunk fits in budget:    keep the whole chunk                         ← coherent, keep whole
else:                         keep the top-`budget` tokens                 ← boundary truncation

If the rescues leave budget unspent, a final fill adds the highest-scoring not-yet-retained tokens,
so exactly L positions are kept and the retained indices are sorted back into original order.

Both gating masks are computed vectorized and moved to CPU once, so the sequential budget walk reads
no GPU scalars per iteration — without that, the per-chunk synchronization made this measurably
slower than plain ChunkKV.

Results

Llama-3.1-8B-Instruct, SnapKV as the inner scorer(for a fair comparison vs chunkkv, although we are even stronger with expected attention), chunk_length=10, rescue_size=4. cr is the
fraction of the KV cache removed. LongBench is the 16-task average; LOOGLE is ROUGE-L × 100. Higher
is better. ChunkKV is run at the same chunk_length=10 so the comparison isolates the entropy gate.

Benchmark	Method	cr=0.70	cr=0.80	cr=0.90
LongBench	no compression	45.85	45.85	45.85
LongBench	ChunkKV	40.91	38.24	32.91
LongBench	EG-ChunkKV	41.34	38.52	34.26
LOOGLE	ChunkKV	26.09	24.89	22.83
LOOGLE	EG-ChunkKV	25.98	25.31	23.29

The gain holds at every ratio and grows as the budget tightens, peaking at +1.35 at
cr=0.90. LOOGLE shows the same direction under aggressive compression (+0.42 and +0.46) but is
slightly behind at cr=0.70 (−0.11), which is consistent with the mechanism: the wasted budget only
becomes costly once slots are genuinely scarce.

Notes for review

• No preprint yet. As mentioned in Entropy Gated ChunkKV (EG-ChunkKV) #261, we have work toward one but nothing released, so the
  docstring cites only the ChunkKV paper.
• default_presses checkbox is intentionally unticked. That list is instantiated as
  cls(compression_ratio=...), but this press requires press= and exposes compression_ratio as a
  delegating property, so it cannot be constructed that way. ChunkKVPress is absent from the list
  for the same reason. Instead the press is added to the wrapper_press matrix in
  tests/presses/test_presses.py, following the ChunkKVPress precedent, which exercises it against
  every press in default_presses.
• chunk_length defaults to 10, not 20 as in ChunkKVPress. Finer granularity gives the gate
  more chunks to reallocate budget between — the table above shows the mechanism is markedly stronger
  at 10 — and it is the configuration all reported numbers use (however we did measure our method at c=20 and it will be presented in our preprint in depth).

Checklist

Before submitting a PR, please make sure:

• Tests are working (make test)

• Code is formatted correctly (make style, on errors try fix with make format)

• Copyright header is included

• All commits are signed-off using git commit -s

• (new press) mypress_press.py is in the presses directory

• (new press) MyPress is in __init__.py

• (new press) README.md is updated with a 1 liner about the new press in the Available presses section

• (new press) New press is in the default_presses list in tests/default_presses.py

• (new press) A docstring is provided that follows the same structure as the existing ones

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Liranitz]

Some further results. All runs used the same model and same hyperparams (cr = fraction of KV cache
removed).

RULER-4096, mean over the 13 subtasks:

Method	cr=0.85	cr=0.90	cr=0.95
ChunkKV	45.03	31.60	21.49
EG-ChunkKV	46.12	36.45	21.99

The gate is not tied to one scorer. Swapping the inner ScorerPress (RULER-4096, 13-subtask
mean, delta = EG-ChunkKV - ChunkKV):

inner scorer	cr=0.85	cr=0.90	cr=0.95
expected_attention	+7.51	+5.56	+4.46
snapkv	+1.09	+4.85	+0.50
tova	-1.57	+3.94	+1.76

Thanks!

[SimJeg]

/ok to test 7797071

[SimJeg]

Hi, thanks for the PR! You will find a first batch of comments to improve the readability of this new press

[SimJeg]

Hi @ShaharBenIshay @Liranitz, thanks for the first updates. Tell me when you think you're done with this first batch of comments.

[ShaharBenIshay]

Thank you @SimJeg for the meaningful comments (and for your time), we have resolved all of them. If, after reviewing the revised version, there are any other comments, we would happily fix and resolve them too.

[SimJeg]

Hi, thanks for the updates ! I left a few final comments

[SimJeg]

/ok to test c57ddd4

[SimJeg]

LGMT ! I'm running the CI/CD and if it passes I will merge. Feel free to open a small PR once you get a preprint out

[ShaharBenIshay]

@SimJeg, much appreciated the feedback.
Thank you!