Implement filtering press (#257)

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.

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}"
             )

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

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

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})"

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 = {}

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)

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