From 635d923c4db0bc65378f486aa3220c8ba16cfeda Mon Sep 17 00:00:00 2001 From: Fabio Massimo Ercoli Date: Tue, 4 Aug 2026 18:36:08 +0200 Subject: [PATCH 1/2] Implement filtering press (#257) Signed-off-by: Fabio Massimo Ercoli --- README.md | 2 + kvpress/__init__.py | 4 + kvpress/padded_tensor.py | 90 ++++++++ kvpress/presses/filtering_press.py | 111 ++++++++++ kvpress/presses/keydiff_press.py | 10 +- kvpress/presses/uniform_filtering_press.py | 75 +++++++ tests/test_filtering_press.py | 236 +++++++++++++++++++++ tests/test_padded_tensor.py | 121 +++++++++++ tests/test_uniform_filtering_press.py | 215 +++++++++++++++++++ 9 files changed, 863 insertions(+), 1 deletion(-) create mode 100644 kvpress/padded_tensor.py create mode 100644 kvpress/presses/filtering_press.py create mode 100644 kvpress/presses/uniform_filtering_press.py create mode 100644 tests/test_filtering_press.py create mode 100644 tests/test_padded_tensor.py create mode 100644 tests/test_uniform_filtering_press.py diff --git a/README.md b/README.md index 03a7988bf..e40bb5813 100644 --- a/README.md +++ b/README.md @@ -144,6 +144,8 @@ Finally we provide wrapper presses that can be combined with other presses: - `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) - `DecodingPress` ([source](kvpress/presses/decoding_press.py)): allow for compression during decoding, see decoding section in this README. - `CompressionRatioDecodingPress` ([source](kvpress/presses/compression_ratio_decoding_press.py)): compress during decoding to keep a fixed fraction of all tokens seen so far. +- `FilteringPress` ([source](kvpress/presses/filtering_press.py)): filter tokens during decoding by making per-head online keep/skip decisions, compatible with append-only cache architectures. +- `UniformFilteringPress` ([source](kvpress/presses/uniform_filtering_press.py)): simpler variant of FilteringPress that makes uniform (all-or-nothing) keep/skip decisions via majority vote across heads. - `PrefillDecodingPress` ([source](kvpress/presses/prefill_decoding_press.py)): allow to compress both during prefilling and during decoding. - `DMSPress` ([source](kvpress/presses/dms_press.py), [paper](https://arxiv.org/abs/2506.05345)): evict keys and values with scores below a given threshold of any `ScorerPress` instead of relying on top-k scores. Support both prefilling and decoding (if decoding=True), but only supports dense-prefill and not sparse-prefill. ⚠️ Does not include the trained evictors from the DMS paper. - `CAMPress` ([source](kvpress/presses/cam_press.py), [paper](https://openreview.net/forum?id=LCTmppB165)): A decoding press that merges the kv cache of evicted tokens into keep tokens to preserve information. diff --git a/kvpress/__init__.py b/kvpress/__init__.py index 62d383b3f..a72807ba3 100644 --- a/kvpress/__init__.py +++ b/kvpress/__init__.py @@ -20,6 +20,7 @@ from kvpress.presses.dms_press import DMSPress from kvpress.presses.duo_attention_press import DuoAttentionPress from kvpress.presses.expected_attention_press import ExpectedAttentionPress +from kvpress.presses.filtering_press import FilteringPress from kvpress.presses.expected_attention_with_stats import ExpectedAttentionStatsPress from kvpress.presses.fastkvzip_press import FastKVzipPress from kvpress.presses.finch_press import FinchPress @@ -46,6 +47,7 @@ from kvpress.presses.streaming_llm_press import StreamingLLMPress from kvpress.presses.think_press import ThinKPress from kvpress.presses.tova_press import TOVAPress +from kvpress.presses.uniform_filtering_press import UniformFilteringPress # Patch the attention functions to support head-wise compression patch_attention_functions() @@ -91,8 +93,10 @@ "KVzapPress", "DMSPress", "FastKVzipPress", + "FilteringPress", "KVComposePress", "MergingPress", "CapPress", "LUKVPress", + "UniformFilteringPress", ] diff --git a/kvpress/padded_tensor.py b/kvpress/padded_tensor.py new file mode 100644 index 000000000..f71f7ac5b --- /dev/null +++ b/kvpress/padded_tensor.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch + +FILL_VALUE = 0.0 + + +class PaddedTensor: + """A 4D tensor (batch, heads, seq, head_dim) with ragged seq dimension. + + Valid data is packed at the front of dim 2 (positions 0..length-1 per head). + Padding positions may contain stale data; call fill_padding() or to_dense() + to materialise FILL_VALUE there. + + Designed for KV cache storage in the filtering press. + """ + + def __init__(self, data: torch.Tensor, lengths: torch.Tensor): + assert data.ndim == 4 # (batch, heads, seq, head_dim) + assert lengths.shape == data.shape[:2] # (batch, heads) + self.data = data + self.lengths = lengths + + @property + def max_length(self) -> int: + return self.data.shape[2] + + @property + def device(self) -> torch.device: + return self.data.device + + @property + def dtype(self) -> torch.dtype: + return self.data.dtype + + def valid_mask(self, include_last=False) -> torch.Tensor: + """Boolean mask of valid (non-padding) positions: shape (batch, heads, seq).""" + indices = torch.arange(self.max_length, device=self.device) + mask = indices < self.lengths.unsqueeze(-1) + if include_last and self.max_length > 0: + mask[:, :, -1] = True + return mask + + def fill_padding(self, fill_value: float = FILL_VALUE) -> None: + """Fill positions beyond valid lengths in-place.""" + mask = ~self.valid_mask().unsqueeze(-1).expand_as(self.data) + self.data[mask] = fill_value + + def accept_last(self, accepted: torch.Tensor) -> None: + """Incorporate the last position into the valid prefix for accepted heads. + + For heads where there is a gap between the valid prefix and the last position + (lengths < max_length - 1), copies data from the last position to + position lengths[head] (first slot after the valid prefix). + Increments lengths for all accepted heads where lengths < max_length. + + accepted: boolean (batch, heads) + """ + last_pos = self.max_length - 1 + needs_swap = accepted & (self.lengths < last_pos) + if needs_swap.any(): + b_idx, h_idx = needs_swap.nonzero(as_tuple=True) + target_pos = self.lengths[b_idx, h_idx] + self.data[b_idx, h_idx, target_pos] = self.data[b_idx, h_idx, last_pos] + + needs_increment = accepted & (self.lengths <= last_pos) + self.lengths = self.lengths + needs_increment.long() + + def remove_last(self, head_bitset: torch.Tensor) -> None: + """Decrement lengths for selected heads. + + head_bitset: boolean (batch, heads) + """ + self.lengths = (self.lengths - head_bitset.long()).clamp_(min=0) + + def shrink(self) -> None: + """Trim backing tensor to the maximum valid length across all heads.""" + max_valid = int(self.lengths.max().item()) if self.lengths.numel() > 0 else 0 + if max_valid < self.max_length: + self.data = self.data[:, :, :max_valid, :].contiguous() + + def clone(self) -> PaddedTensor: + """Create an independent deep copy.""" + return PaddedTensor(self.data.clone(), self.lengths.clone()) + + def __repr__(self) -> str: + return f"PaddedTensor(shape={list(self.data.shape)}, lengths={self.lengths})" diff --git a/kvpress/presses/filtering_press.py b/kvpress/presses/filtering_press.py new file mode 100644 index 000000000..0b04f5405 --- /dev/null +++ b/kvpress/presses/filtering_press.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +import torch +from torch import nn + +from kvpress.padded_tensor import PaddedTensor +from kvpress.presses.decoding_press import DecodingPress + + +@dataclass +class FilteringPress(DecodingPress): + """ + A decoding press that filters tokens during decoding by making online keep/skip decisions. + + Instead of retroactive eviction (scoring all tokens and removing the lowest-scored), + this press decides for each new decode token whether to keep it in the cache. + Only the newest token can be removed — existing cache entries are never modified. + + This makes the press compatible with append-only cache architectures (e.g. vLLM's + paged KV cache). During prefill, this press is a no-op — filtering only applies + to the decode phase, where tokens arrive one at a time and the cache is append-only. + + The decision is made per head: each head independently scores all tokens + (including the new one) using the wrapped ScorerPress and checks whether the + new token's score is above the eviction threshold at the target compression + ratio. Rejected heads mark the position as padding; accepted heads that find + an earlier padding slot move the token there to keep valid tokens packed. + When all heads have padding at the last position, it is removed to shrink + the cache. + + This press requires logical ``position_ids`` to be passed through the model + forward call. + + Parameters + ---------- + base_press : ScorerPress + The scorer press used to compute importance scores for tokens. + target_compression_ratio : float, default=0.5 + Target fraction of tokens to filter out during decoding. + compression_interval : int, default=1 + Number of decoding steps between filtering decisions. + fill_padding : bool, default=True + Zero out padding positions after filtering. Per-head rejection creates + positions that are padding for some heads but valid for others; zeroing + prevents stale data from affecting attention. Disabling is faster but + may degrade quality. + hidden_states_buffer_size : int, default=256 + Maximum number of hidden states to keep before compression. + """ + + target_compression_ratio: float = 0.5 + compression_interval: int = 1 + fill_padding: bool = True + target_size: int = field(default=1, init=False) + + def __post_init__(self): + super().__post_init__() + assert 0 <= self.target_compression_ratio < 1, "target_compression_ratio must be between 0 and 1" + self._lengths = {} + + 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]: + total_tokens_seen = int(kwargs["position_ids"].max().item()) + 1 + n_kept = max(1, int(total_tokens_seen * (1 - self.target_compression_ratio))) + n_kept = min(n_kept, keys.shape[2]) + + layer_idx = getattr(module, "layer_idx", 0) + if layer_idx in self._lengths: + lengths = self._lengths[layer_idx] + else: + lengths = torch.full(keys.shape[:2], keys.shape[2] - 1, dtype=torch.long, device=keys.device) + + kt = PaddedTensor(keys.clone(), lengths.clone()) + vt = PaddedTensor(values.clone(), lengths.clone()) + valid_mask = kt.valid_mask(include_last=True) + + scores = self.base_press.score( + module, hidden_states, kt.data, vt.data, attentions, {**kwargs, "valid_mask": valid_mask} + ) + scores[~valid_mask] = float("-inf") + + threshold = scores.topk(n_kept, dim=-1, sorted=True).values[:, :, -1] + rejected = scores[:, :, -1] < threshold + + if rejected.all(): + return keys[:, :, :-1, :].contiguous(), values[:, :, :-1, :].contiguous() + + kt.accept_last(~rejected) + vt.accept_last(~rejected) + if self.fill_padding: + kt.fill_padding() + vt.fill_padding() + kt.shrink() + vt.shrink() + + self._lengths[layer_idx] = kt.lengths.clone() + return kt.data, vt.data + + def reset(self): + super().reset() + self._lengths = {} diff --git a/kvpress/presses/keydiff_press.py b/kvpress/presses/keydiff_press.py index b2bba83fa..fd606b986 100644 --- a/kvpress/presses/keydiff_press.py +++ b/kvpress/presses/keydiff_press.py @@ -42,5 +42,13 @@ def score( attentions: torch.Tensor, kwargs, ) -> torch.Tensor: - anchor = F.normalize(keys, p=2, dim=-1).mean(dim=2, keepdim=True) + normalized = F.normalize(keys, p=2, dim=-1) + valid_mask = kwargs.get("valid_mask") + + if valid_mask is None: + anchor = normalized.mean(dim=2, keepdim=True) + else: + mask = valid_mask.unsqueeze(-1) + anchor = (normalized * mask).sum(dim=2, keepdim=True) / mask.sum(dim=2, keepdim=True).clamp(min=1) + return -F.cosine_similarity(keys, anchor, dim=-1) diff --git a/kvpress/presses/uniform_filtering_press.py b/kvpress/presses/uniform_filtering_press.py new file mode 100644 index 000000000..45417a957 --- /dev/null +++ b/kvpress/presses/uniform_filtering_press.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +import torch +from torch import nn + +from kvpress.presses.decoding_press import DecodingPress + + +@dataclass +class UniformFilteringPress(DecodingPress): + """ + A simpler variant of FilteringPress that makes uniform (all-or-nothing) keep/skip + decisions per token instead of per-head decisions. + + Like FilteringPress, this press decides for each new decode token whether to keep + it in the cache, making it compatible with append-only cache architectures + (e.g. vLLM's paged KV cache). During prefill it is a no-op. + + The key difference from FilteringPress: each head independently checks whether + the new token would survive retroactive compression, then a majority vote across + heads produces a single keep/remove decision. If a majority of heads would evict + the token, it is removed from the cache entirely. This avoids per-head ragged + lengths and does not require PaddedTensor. + + Parameters + ---------- + base_press : ScorerPress + The scorer press used to compute importance scores for tokens. + target_compression_ratio : float, default=0.5 + Target fraction of tokens to filter out during decoding. + compression_interval : int, default=1 + Number of decoding steps between filtering decisions. + hidden_states_buffer_size : int, default=256 + Maximum number of hidden states to keep before compression. + """ + + target_compression_ratio: float = 0.5 + compression_interval: int = 1 + target_size: int = field(default=1, init=False) + + def __post_init__(self): + super().__post_init__() + assert 0 <= self.target_compression_ratio < 1, "target_compression_ratio must be between 0 and 1" + + 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]: + k_len = keys.shape[2] + n_kept = max(1, int(k_len * (1 - self.target_compression_ratio))) + + if n_kept >= k_len: + return keys, values + + scores = self.base_press.score(module, hidden_states, keys, values, attentions, kwargs) + + new_token_scores = scores[:, :, -1] + threshold = scores.topk(n_kept, dim=-1).values[:, :, -1] + survives_per_head = new_token_scores >= threshold + + keep = survives_per_head.float().mean(dim=-1) >= 0.5 + + if not keep.any(): + keys = keys[:, :, :-1, :].contiguous() + values = values[:, :, :-1, :].contiguous() + + return keys, values diff --git a/tests/test_filtering_press.py b/tests/test_filtering_press.py new file mode 100644 index 000000000..3474e83d2 --- /dev/null +++ b/tests/test_filtering_press.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Tests for FilteringPress — online per-token keep/skip decisions during decoding. +""" + +from dataclasses import dataclass + +import pytest +import torch +from transformers import DynamicCache, pipeline + +from kvpress import ( + FilteringPress, + KeyDiffPress, + KnormPress, + PrefillDecodingPress, + StreamingLLMPress, + TOVAPress, +) +from kvpress.presses.scorer_press import ScorerPress + + +@dataclass +class FixedScorePress(ScorerPress): + fixed_scores: torch.Tensor = None + + def score(self, module, hidden_states, keys, values, attentions, kwargs): + return self.fixed_scores + + +@pytest.fixture(scope="module") +def pipe(): + return pipeline("kv-press-text-generation", model="MaxJeblick/llama2-0b-unit-test", device_map="auto") + + +CONTEXT = "The quick brown fox jumps over the lazy dog. " * 10 +QUESTION = "What animal jumps over the dog?" + + +def test_filtering_press_reduces_cache(pipe): + """FilteringPress should produce a smaller cache than no compression.""" + model = pipe.model + tokenizer = pipe.tokenizer + device = model.device + + input_ids = tokenizer.encode(CONTEXT, return_tensors="pt").to(device) + + cache_baseline = DynamicCache() + with torch.no_grad(): + model.generate(input_ids, past_key_values=cache_baseline, max_new_tokens=20, do_sample=False) + baseline_len = cache_baseline.get_seq_length() + + press = FilteringPress(base_press=KnormPress(), target_compression_ratio=0.9) + cache_filtered = DynamicCache() + with torch.no_grad(), press(model): + model.generate(input_ids, past_key_values=cache_filtered, max_new_tokens=20, do_sample=False) + filtered_len = cache_filtered.get_seq_length() + + assert filtered_len < baseline_len, ( + f"filtered cache ({filtered_len}) should be smaller than baseline ({baseline_len})" + ) + + +def test_filtering_press_no_op_at_zero_ratio(pipe): + """target_compression_ratio=0 should not filter any tokens.""" + cache_baseline = DynamicCache() + pipe(CONTEXT, question=QUESTION, cache=cache_baseline, max_new_tokens=20) + + press = FilteringPress(base_press=KnormPress(), target_compression_ratio=0.0) + cache_filtered = DynamicCache() + pipe(CONTEXT, question=QUESTION, press=press, cache=cache_filtered, max_new_tokens=20) + + for layer_idx in range(len(cache_baseline.layers)): + assert cache_baseline.layers[layer_idx].keys.shape[2] == cache_filtered.layers[layer_idx].keys.shape[2] + + +def test_filtering_press_with_prefill_decoding(pipe): + """FilteringPress should work as decoding_press inside PrefillDecodingPress.""" + combined_press = PrefillDecodingPress( + prefilling_press=KeyDiffPress(compression_ratio=0.5), + decoding_press=FilteringPress(base_press=KeyDiffPress(), target_compression_ratio=0.5), + ) + + cache = DynamicCache() + result = pipe(CONTEXT, question=QUESTION, press=combined_press, cache=cache, max_new_tokens=15) + + assert len(result["answer"]) > 0, "No answer generated" + + +@pytest.mark.parametrize("scorer_cls", [KnormPress, KeyDiffPress, TOVAPress, StreamingLLMPress]) +def test_filtering_press_with_different_scorers(pipe, scorer_cls): + """FilteringPress should work with any ScorerPress.""" + press = FilteringPress(base_press=scorer_cls(), target_compression_ratio=0.5) + + cache = DynamicCache() + result = pipe(CONTEXT, question=QUESTION, press=press, cache=cache, max_new_tokens=15) + + assert len(result["answer"]) > 0, f"No answer generated with {scorer_cls.__name__}" + + +def test_filtering_press_higher_ratio_filters_more(pipe): + """Higher compression ratio should produce a smaller cache.""" + model = pipe.model + tokenizer = pipe.tokenizer + device = model.device + + input_ids = tokenizer.encode(CONTEXT, return_tensors="pt").to(device) + + cache_low = DynamicCache() + press_low = FilteringPress(base_press=KnormPress(), target_compression_ratio=0.3) + with torch.no_grad(), press_low(model): + model.generate(input_ids, past_key_values=cache_low, max_new_tokens=20, do_sample=False) + + cache_high = DynamicCache() + press_high = FilteringPress(base_press=KnormPress(), target_compression_ratio=0.7) + with torch.no_grad(), press_high(model): + model.generate(input_ids, past_key_values=cache_high, max_new_tokens=20, do_sample=False) + + low_len = cache_low.get_seq_length() + high_len = cache_high.get_seq_length() + assert high_len <= low_len, ( + f"higher ratio cache ({high_len}) should be <= lower ratio cache ({low_len})" + ) + + +def test_filtering_press_reuse_across_sequences(pipe): + """Reusing a FilteringPress across sequences should not crash.""" + press = FilteringPress(base_press=KnormPress(), target_compression_ratio=0.5) + + model = pipe.model + device = model.device + long_ids = torch.arange(1, 81, dtype=torch.long, device=device).unsqueeze(0) + short_ids = torch.arange(1, 9, dtype=torch.long, device=device).unsqueeze(0) + + with torch.no_grad(), press(model): + model.generate(long_ids, max_new_tokens=6, do_sample=False) + model.generate(short_ids, max_new_tokens=6, do_sample=False) + + +BATCH, N_HEADS, SEQ_LEN, HEAD_DIM = 1, 2, 10, 4 + + +def _make_press(scores, ratio=0.5): + scorer = FixedScorePress() + scorer.fixed_scores = scores + return FilteringPress(base_press=scorer, target_compression_ratio=ratio) + + +def _make_dummy_tensors(seq_len=SEQ_LEN): + keys = torch.randn(BATCH, N_HEADS, seq_len, HEAD_DIM) + values = torch.randn(BATCH, N_HEADS, seq_len, HEAD_DIM) + hidden_states = torch.randn(BATCH, seq_len, HEAD_DIM) + kwargs = {"position_ids": torch.arange(seq_len).unsqueeze(0)} + return keys, values, hidden_states, kwargs + + +def _base_scores(): + """Scores where positions 0-4 are high (5.0) and 5-8 are low (1.0), last token varies.""" + scores = torch.zeros(BATCH, N_HEADS, SEQ_LEN) + scores[:, :, :5] = 5.0 + scores[:, :, 5:9] = 1.0 + return scores + + +def test_compress_all_heads_accept(): + """Token kept at last position when all heads accept.""" + scores = _base_scores() + scores[:, :, -1] = 5.0 + press = _make_press(scores) + keys, values, hidden_states, kwargs = _make_dummy_tensors() + + out_keys, out_values = press.compress(None, hidden_states, keys, values, None, kwargs) + + assert out_keys.shape[2] == SEQ_LEN + assert not torch.isinf(out_keys[:, :, -1, :]).any() + + +def test_compress_all_heads_reject(): + """Cache shrinks when all heads reject the new token.""" + scores = _base_scores() + scores[:, :, -1] = 0.0 + press = _make_press(scores) + keys, values, hidden_states, kwargs = _make_dummy_tensors() + + out_keys, out_values = press.compress(None, hidden_states, keys, values, None, kwargs) + + assert out_keys.shape[2] == SEQ_LEN - 1 + + +def test_compress_one_head_rejects(): + """Shape unchanged when one head rejects; rejected head gets -inf at last position.""" + scores = _base_scores() + scores[:, 0, -1] = 5.0 + scores[:, 1, -1] = 0.0 + press = _make_press(scores) + keys, values, hidden_states, kwargs = _make_dummy_tensors() + + out_keys, out_values = press.compress(None, hidden_states, keys, values, None, kwargs) + + assert out_keys.shape[2] == SEQ_LEN + assert not torch.isinf(out_keys[:, 0, -1, :]).any(), "accepted head should keep valid data" + assert (out_keys[:, 1, -1, :] == 0.0).all(), "rejected head should have padding fill value" + + +def test_compress_accepted_head_fills_gap(): + """Accepted head packs new token into prefix at stored length position.""" + scores = _base_scores() + scores[:, 0, -1] = 5.0 # head 0 accepts (new token at last position) + scores[:, 1, -1] = 0.0 # head 1 rejects + press = _make_press(scores) + keys, values, hidden_states, kwargs = _make_dummy_tensors() + + # Simulate prior state: head 0 has 8 valid tokens, head 1 has 9 + original_new_key = keys[:, 0, -1, :].clone() + press._lengths[0] = torch.tensor([[8, 9]]) + + out_keys, out_values = press.compress(None, hidden_states, keys, values, None, kwargs) + + # Head 0 accepted: new token packed at position 8, lengths=9 + # Head 1 rejected: lengths=9 → both heads at 9 → shrink to 9 + assert out_keys.shape[2] == SEQ_LEN - 1 + assert torch.allclose(out_keys[0, 0, 8, :], original_new_key[0]) + + +def test_compress_filters_even_with_small_cache(): + """Filtering applies even when the cache is smaller than n_kept.""" + small_seq = 3 + scores = torch.tensor([[[5.0, 5.0, 0.0], [5.0, 5.0, 0.0]]]) + press = _make_press(scores, ratio=0.5) + keys, values, hidden_states, kwargs = _make_dummy_tensors(seq_len=small_seq) + + out_keys, out_values = press.compress(None, hidden_states, keys, values, None, kwargs) + + assert out_keys.shape[2] == small_seq - 1 diff --git a/tests/test_padded_tensor.py b/tests/test_padded_tensor.py new file mode 100644 index 000000000..805fef18d --- /dev/null +++ b/tests/test_padded_tensor.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from kvpress.padded_tensor import FILL_VALUE, PaddedTensor + + +def test_accept_last_with_swap(): + data = torch.randn(1, 2, 6, 4) + new_token = data[0, 0, -1, :].clone() + lengths = torch.tensor([[3, 5]]) + pt = PaddedTensor(data, lengths) + + accepted = torch.tensor([[True, True]]) + pt.accept_last(accepted) + + assert pt.lengths.tolist() == [[4, 6]] + torch.testing.assert_close(pt.data[0, 0, 3, :], new_token) + + +def test_accept_last_no_gap(): + data = torch.randn(1, 2, 6, 4) + lengths = torch.tensor([[5, 5]]) + pt = PaddedTensor(data, lengths) + + accepted = torch.tensor([[True, True]]) + pt.accept_last(accepted) + assert pt.lengths.tolist() == [[6, 6]] + + +def test_accept_last_already_full(): + data = torch.randn(1, 2, 6, 4) + lengths = torch.tensor([[6, 6]]) + pt = PaddedTensor(data, lengths) + + accepted = torch.tensor([[True, True]]) + pt.accept_last(accepted) + assert pt.lengths.tolist() == [[6, 6]] + + +def test_accept_last_mixed(): + data = torch.randn(1, 2, 6, 4) + lengths = torch.tensor([[3, 5]]) + pt = PaddedTensor(data, lengths) + + accepted = torch.tensor([[True, False]]) + pt.accept_last(accepted) + assert pt.lengths.tolist() == [[4, 5]] + + +def test_fill_padding(): + data = torch.randn(1, 2, 6, 4) + lengths = torch.tensor([[3, 5]]) + pt = PaddedTensor(data, lengths) + + pt.fill_padding() + + assert not torch.isinf(pt.data[0, 0, :3, :]).any() + assert (pt.data[0, 0, 3:, :] == FILL_VALUE).all() + assert not torch.isinf(pt.data[0, 1, :5, :]).any() + assert (pt.data[0, 1, 5:, :] == FILL_VALUE).all() + + +def test_fill_padding_with_custom_value(): + data = torch.randn(1, 2, 6, 4) + lengths = torch.tensor([[3, 5]]) + pt = PaddedTensor(data, lengths) + + pt.fill_padding(0) + + assert (pt.data[0, 0, 3:, :] == 0).all() + assert (pt.data[0, 1, 5:, :] == 0).all() + + +def test_valid_mask(): + data = torch.randn(1, 2, 6, 4) + lengths = torch.tensor([[3, 6]]) + pt = PaddedTensor(data, lengths) + + mask = pt.valid_mask() + assert mask[0, 0].tolist() == [True, True, True, False, False, False] + assert mask[0, 1].tolist() == [True, True, True, True, True, True] + + +def test_shrink(): + data = torch.randn(1, 2, 8, 4) + lengths = torch.tensor([[4, 6]]) + pt = PaddedTensor(data, lengths) + + pt.shrink() + assert pt.data.shape == (1, 2, 6, 4) + + +def test_remove_last(): + data = torch.randn(1, 2, 6, 4) + lengths = torch.tensor([[5, 6]]) + pt = PaddedTensor(data, lengths) + + pt.remove_last(torch.tensor([[True, False]])) + assert pt.lengths.tolist() == [[4, 6]] + + +def test_valid_mask_include_last(): + data = torch.randn(1, 2, 6, 4) + lengths = torch.tensor([[3, 5]]) + pt = PaddedTensor(data, lengths) + + mask = pt.valid_mask(include_last=True) + assert mask[0, 0].tolist() == [True, True, True, False, False, True] + assert mask[0, 1].tolist() == [True, True, True, True, True, True] + + +def test_clone(): + data = torch.randn(1, 2, 6, 4) + lengths = torch.tensor([[3, 5]]) + pt = PaddedTensor(data, lengths) + + pt2 = pt.clone() + pt2.data[0, 0, 0, 0] = 999.0 + assert pt.data[0, 0, 0, 0] != 999.0 diff --git a/tests/test_uniform_filtering_press.py b/tests/test_uniform_filtering_press.py new file mode 100644 index 000000000..33805171c --- /dev/null +++ b/tests/test_uniform_filtering_press.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Tests for UniformFilteringPress — online per-token keep/skip decisions via majority vote. +""" + +from dataclasses import dataclass + +import pytest +import torch +from transformers import DynamicCache, pipeline + +from kvpress import ( + KeyDiffPress, + KnormPress, + PrefillDecodingPress, + StreamingLLMPress, + TOVAPress, + UniformFilteringPress, +) +from kvpress.presses.scorer_press import ScorerPress + + +@dataclass +class FixedScorePress(ScorerPress): + fixed_scores: torch.Tensor = None + + def score(self, module, hidden_states, keys, values, attentions, kwargs): + return self.fixed_scores + + +@pytest.fixture(scope="module") +def pipe(): + return pipeline("kv-press-text-generation", model="MaxJeblick/llama2-0b-unit-test", device_map="auto") + + +CONTEXT = "The quick brown fox jumps over the lazy dog. " * 10 +QUESTION = "What animal jumps over the dog?" + + +def test_uniform_filtering_press_reduces_cache(pipe): + """UniformFilteringPress should produce a smaller cache than no compression.""" + model = pipe.model + tokenizer = pipe.tokenizer + device = model.device + + input_ids = tokenizer.encode(CONTEXT, return_tensors="pt").to(device) + + cache_baseline = DynamicCache() + with torch.no_grad(): + model.generate(input_ids, past_key_values=cache_baseline, max_new_tokens=20, do_sample=False) + baseline_len = cache_baseline.get_seq_length() + + press = UniformFilteringPress(base_press=KnormPress(), target_compression_ratio=0.9) + cache_filtered = DynamicCache() + with torch.no_grad(), press(model): + model.generate(input_ids, past_key_values=cache_filtered, max_new_tokens=20, do_sample=False) + filtered_len = cache_filtered.get_seq_length() + + assert filtered_len < baseline_len, ( + f"filtered cache ({filtered_len}) should be smaller than baseline ({baseline_len})" + ) + + +def test_uniform_filtering_press_no_op_at_zero_ratio(pipe): + """target_compression_ratio=0 should not filter any tokens.""" + cache_baseline = DynamicCache() + pipe(CONTEXT, question=QUESTION, cache=cache_baseline, max_new_tokens=20) + + press = UniformFilteringPress(base_press=KnormPress(), target_compression_ratio=0.0) + cache_filtered = DynamicCache() + pipe(CONTEXT, question=QUESTION, press=press, cache=cache_filtered, max_new_tokens=20) + + for layer_idx in range(len(cache_baseline.layers)): + assert cache_baseline.layers[layer_idx].keys.shape[2] == cache_filtered.layers[layer_idx].keys.shape[2] + + +def test_uniform_filtering_press_with_prefill_decoding(pipe): + """UniformFilteringPress should work as decoding_press inside PrefillDecodingPress.""" + combined_press = PrefillDecodingPress( + prefilling_press=KeyDiffPress(compression_ratio=0.5), + decoding_press=UniformFilteringPress(base_press=KeyDiffPress(), target_compression_ratio=0.5), + ) + + cache = DynamicCache() + result = pipe(CONTEXT, question=QUESTION, press=combined_press, cache=cache, max_new_tokens=15) + + assert len(result["answer"]) > 0, "No answer generated" + + +@pytest.mark.parametrize("scorer_cls", [KnormPress, KeyDiffPress, TOVAPress, StreamingLLMPress]) +def test_uniform_filtering_press_with_different_scorers(pipe, scorer_cls): + """UniformFilteringPress should work with any ScorerPress.""" + press = UniformFilteringPress(base_press=scorer_cls(), target_compression_ratio=0.5) + + cache = DynamicCache() + result = pipe(CONTEXT, question=QUESTION, press=press, cache=cache, max_new_tokens=15) + + assert len(result["answer"]) > 0, f"No answer generated with {scorer_cls.__name__}" + + +def test_uniform_filtering_press_higher_ratio_filters_more(pipe): + """Higher compression ratio should produce a smaller cache.""" + model = pipe.model + tokenizer = pipe.tokenizer + device = model.device + + input_ids = tokenizer.encode(CONTEXT, return_tensors="pt").to(device) + + cache_low = DynamicCache() + press_low = UniformFilteringPress(base_press=KnormPress(), target_compression_ratio=0.3) + with torch.no_grad(), press_low(model): + model.generate(input_ids, past_key_values=cache_low, max_new_tokens=20, do_sample=False) + + cache_high = DynamicCache() + press_high = UniformFilteringPress(base_press=KnormPress(), target_compression_ratio=0.7) + with torch.no_grad(), press_high(model): + model.generate(input_ids, past_key_values=cache_high, max_new_tokens=20, do_sample=False) + + low_len = cache_low.get_seq_length() + high_len = cache_high.get_seq_length() + assert high_len <= low_len, ( + f"higher ratio cache ({high_len}) should be <= lower ratio cache ({low_len})" + ) + + +def test_uniform_filtering_press_reuse_across_sequences(pipe): + """Reusing a UniformFilteringPress across sequences should not crash.""" + press = UniformFilteringPress(base_press=KnormPress(), target_compression_ratio=0.5) + + model = pipe.model + device = model.device + long_ids = torch.arange(1, 81, dtype=torch.long, device=device).unsqueeze(0) + short_ids = torch.arange(1, 9, dtype=torch.long, device=device).unsqueeze(0) + + with torch.no_grad(), press(model): + model.generate(long_ids, max_new_tokens=6, do_sample=False) + model.generate(short_ids, max_new_tokens=6, do_sample=False) + + +# --- Unit tests --- + +BATCH, N_HEADS, SEQ_LEN, HEAD_DIM = 1, 2, 10, 4 + + +def _make_press(scores, ratio=0.5): + scorer = FixedScorePress() + scorer.fixed_scores = scores + return UniformFilteringPress(base_press=scorer, target_compression_ratio=ratio) + + +def _make_dummy_tensors(seq_len=SEQ_LEN): + keys = torch.randn(BATCH, N_HEADS, seq_len, HEAD_DIM) + values = torch.randn(BATCH, N_HEADS, seq_len, HEAD_DIM) + hidden_states = torch.randn(BATCH, seq_len, HEAD_DIM) + kwargs = {} + return keys, values, hidden_states, kwargs + + +def _base_scores(): + """Scores where positions 0-4 are high (5.0) and 5-8 are low (1.0), last token varies.""" + scores = torch.zeros(BATCH, N_HEADS, SEQ_LEN) + scores[:, :, :5] = 5.0 + scores[:, :, 5:9] = 1.0 + return scores + + +def test_compress_all_heads_accept(): + """Token kept when all heads accept.""" + scores = _base_scores() + scores[:, :, -1] = 5.0 + press = _make_press(scores) + keys, values, hidden_states, kwargs = _make_dummy_tensors() + + out_keys, out_values = press.compress(None, hidden_states, keys, values, None, kwargs) + + assert out_keys.shape[2] == SEQ_LEN + + +def test_compress_all_heads_reject(): + """Cache shrinks when all heads reject the new token.""" + scores = _base_scores() + scores[:, :, -1] = 0.0 + press = _make_press(scores) + keys, values, hidden_states, kwargs = _make_dummy_tensors() + + out_keys, out_values = press.compress(None, hidden_states, keys, values, None, kwargs) + + assert out_keys.shape[2] == SEQ_LEN - 1 + + +def test_compress_tie_keeps_token(): + """With 2 heads, 1 accepting and 1 rejecting gives mean=0.5 >= 0.5 — token is kept.""" + scores = _base_scores() + scores[:, 0, -1] = 5.0 # head 0 accepts + scores[:, 1, -1] = 0.0 # head 1 rejects + press = _make_press(scores) + keys, values, hidden_states, kwargs = _make_dummy_tensors() + + out_keys, out_values = press.compress(None, hidden_states, keys, values, None, kwargs) + + assert out_keys.shape[2] == SEQ_LEN, "tie should keep the token" + + +def test_compress_no_op_when_n_kept_ge_k_len(): + """No filtering when n_kept >= k_len.""" + scores = _base_scores() + scores[:, :, -1] = 0.0 # would normally be rejected + press = _make_press(scores, ratio=0.0) + keys, values, hidden_states, kwargs = _make_dummy_tensors() + + out_keys, out_values = press.compress(None, hidden_states, keys, values, None, kwargs) + + assert out_keys.shape[2] == SEQ_LEN, "should be no-op when n_kept >= k_len" From 48cc0e262c60b793d0054f87a795300fae54526c Mon Sep 17 00:00:00 2001 From: Fabio Massimo Ercoli Date: Fri, 7 Aug 2026 06:21:27 +0200 Subject: [PATCH 2/2] Evaluate filtering press Signed-off-by: Fabio Massimo Ercoli --- evaluation/evaluate.py | 2 ++ evaluation/evaluate_registry.py | 7 +++++++ 2 files changed, 9 insertions(+) diff --git a/evaluation/evaluate.py b/evaluation/evaluate.py index 49e57715e..65e2b1359 100644 --- a/evaluation/evaluate.py +++ b/evaluation/evaluate.py @@ -298,6 +298,8 @@ def _setup_press(self): press.compression_interval = self.config.compression_interval or press.compression_interval press.target_size = self.config.target_size or press.target_size press.hidden_states_buffer_size = self.config.hidden_states_buffer_size or press.hidden_states_buffer_size + if hasattr(press, "target_compression_ratio"): + press.target_compression_ratio = compression_ratio logger.info( f"Set DecodingPress compression_interval to {self.config.compression_interval}, target_size to {self.config.target_size}, hidden_states_buffer_size to {self.config.hidden_states_buffer_size}" ) diff --git a/evaluation/evaluate_registry.py b/evaluation/evaluate_registry.py index 1a8b47f1c..de3ec0b4f 100644 --- a/evaluation/evaluate_registry.py +++ b/evaluation/evaluate_registry.py @@ -19,6 +19,7 @@ ChunkKVPress, CompactorPress, ComposedPress, + CompressionRatioDecodingPress, CriticalAdaKVPress, CriticalKVPress, CURPress, @@ -27,6 +28,7 @@ DuoAttentionPress, ExpectedAttentionPress, FastKVzipPress, + FilteringPress, FinchPress, KeyDiffPress, KnormPress, @@ -44,6 +46,7 @@ StreamingLLMPress, ThinKPress, TOVAPress, + UniformFilteringPress, ) # These dictionaries define the available datasets, scorers, and KVPress methods for evaluation. @@ -129,4 +132,8 @@ "merging_snapkv": MergingPress(SnapKVPress()), "merging_expected_attention": MergingPress(ExpectedAttentionPress(epsilon=1e-2)), "merging_kvzap_mlp": MergingPress(KVzapPress(model_type="mlp")), + "compression_ratio_decoding_keydiff": CompressionRatioDecodingPress(base_press=KeyDiffPress()), + "uniform_filtering_keydiff": UniformFilteringPress(base_press=KeyDiffPress()), + "filtering_zerofill_keydiff": FilteringPress(base_press=KeyDiffPress()), + "filtering_stale_keydiff": FilteringPress(base_press=KeyDiffPress(), fill_padding=False), }