Implement filtering press (#257) - #258
Conversation
Signed-off-by: Fabio Massimo Ercoli <fabiomassimo.ercoli@gmail.com>
Signed-off-by: Fabio Massimo Ercoli <fabiomassimo.ercoli@gmail.com>
|
Hi @fax4ever ! Thanks for opening this PR 🙂 General commentI think I need some clarification here to understand whether this is in scope for KVPress. KVPress is currently a research-friendly library for comparing KV cache compression methods. We have not tried to provide direct compatibility with inference engines such as vLLM. I think this can fit the repo if it is framed as an experimental proxy for the question:
But then I think we should be explicit about what this is not:
What do you think of this framing ? EvaluationAbout the proposed evaluation on RULER. RULER is mainly a long-context/prefill compression benchmark. These presses are decoding-only and are no-ops during prefill. That does not seem like a meaningful comparison against standard KVPress methods that compress the context during prefill. For decoding, please consider using other datasets. Also, if you have other results or tests that seem meaningful to you, please feel free to share here, so we can use them for future reference. Having broader evaluation of this press would also be helpful. CodeI left some comments in the code, but something important: there are two new presses in this PR, uniform and non uniform. For the non uniform (so per head filtering), I have two observations:
|
| 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 |
There was a problem hiding this comment.
nit: maybe we can keep the order nicer here and move this one after expected attention (alphab order)
| 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 |
There was a problem hiding this comment.
DecodingPress.forward_hook still does cache_layer.keys = keys . Please reword to something like: "The decision model is compatible with append-only caches - each step's decision depends only on the newest token - but this reference implementation still writes back a full tensor. Integrating with a real paged-cache backend (e.g. vLLM) is out of scope for this PR."
| 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 |
There was a problem hiding this comment.
Are we sure this one is correct and could actually become compatible with vLLM ? I mean, does vLLM allow per-head block allocation or do all heads for a token live in one block. ? Just want to get some context about how vLLM works here :)
|
|
||
| kt.accept_last(~rejected) | ||
| vt.accept_last(~rejected) | ||
| if self.fill_padding: |
There was a problem hiding this comment.
kvpress already has infrastructure for head-wise masking without probability leakage: kvpress/attention_patch.py patches every attention function to substitute a fake key at positions listed in module.masked_key_indices, for example AdaKVPress uses this pattern.
Zeroing padded rows still lets attention assign non-zero softmax mass to them, so fill_padding=True silently degrades quality and fill_padding=False is worse.
| attentions: torch.Tensor, | ||
| kwargs: dict, | ||
| ) -> tuple[torch.Tensor, torch.Tensor]: | ||
| total_tokens_seen = int(kwargs["position_ids"].max().item()) + 1 |
There was a problem hiding this comment.
FilteringPress uses total_tokens_seen but UniformFilteringPress uses k_len for the same purpose. Should we align them, or document why they intentionally differ ?
| kwargs, | ||
| ) -> torch.Tensor: | ||
| anchor = F.normalize(keys, p=2, dim=-1).mean(dim=2, keepdim=True) | ||
| normalized = F.normalize(keys, p=2, dim=-1) |
There was a problem hiding this comment.
What about other scorer presses ? They will silently produce biased scores under FilteringPress ?
| @@ -0,0 +1,111 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
This non uniform FilteringPress could probably be implemented through the existing masked_key_indices mechanism instead of introducing PaddedTensor and per-head ragged cache lengths ? KVPress already has an attention patch that supports head-wise masking through module.masked_key_indices, and several presses use that path (leaving a comment later here)
| @@ -0,0 +1,75 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
If we use masked_key_indices as suggested in the other comment we can cleanup this file and make this PR slimmer
2 participants
PR description
Description of your PR. Fixes #257
Checklist
Before submitting a PR, please make sure:
Tests are working (
make test)Code is formatted correctly (
make style, on errors try fix withmake format)Copyright header is included
All commits are signed-off using
git commit -s(new press)
mypress_press.pyis in thepressesdirectory(new press)
MyPressis in__init__.py(new press)
README.mdis updated with a 1 liner about the new press in the Available presses section(new press) New press is in the
default_presseslist intests/default_presses.py(new press) A docstring is provided that follows the same structure as the existing ones