[ljwljwljwljw]

PR description

Add EchoPress, a training-free approximation to KVzip that reuses prefill queries and cached keys to score virtual context repetitions. Paper: https://arxiv.org/abs/2610.00412

EchoPress virtually reconstructs context by rotating queries during prefill. It performs a single KVzip-style reconstruction pass on the first chunk, capped at 2K tokens, then uses the exact KVzip scores from that chunk to calibrate the virtual scores of subsequent chunks.

Evaluation

EchoPress matches KVZip's accuracy on both RULER-4K and LongBench. Benchmark details can be found in our paper.

Code changes

• PyTorch scoring and an optional Triton backend.
• Public export and an echo entry in the evaluation registry.
• README documentation, shared test configurations, and the unit test for EchoPress.

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.

[SimJeg]

Hi @ljwljwljwljw, thanks for your PR ! We don't accept custom kernels in kvpress. Could you remove the triton kernel and only use the torch backend ? Feel free to let a comment in the docstring to link to your fork with the kernel

[SimJeg]

@ljwljwljwljw I've let a first series of comments. The goal of kvpress is to provide short and easy to read paper reproduction. I believe the code could be very largely compacted if you remove all the unused variants and don't duplicate code from KVzipPress.

[ljwljwljwljw]

Thanks for your review! I updated the code to addresses all the comments above:

• Triton kernel removed. The press now only has the PyTorch path. The kernel lives in my fork (https://github.com/ljwljwljwljw/kvpress/tree/echo-press-triton) and is linked from the docstring.
• Parameters simplified. EchoPress now has only parameters compression_ratio, layerwise, n_sink and chunk_size from KVzipPress.
• get_query_states is used in the prefill hook instead of the q_proj hook.
• __post_init__ calls the parent and asserts kvzip_plus_normalization is False.
• Registry comment removed.
• Duplicated code removed. prepare, the chat-template logic and the eviction are inherited. Only _perform_kvzip_compression is overridden to replay the first chunk and calibrate the other chunks.

Since EchoPress runs the parent __post_init__, it now logs the KVzipPress warning about the 2-3x prefilling overhead, which does not actually apply in EchoPress. Should I modify KVzipPress to add an option for logging the warnning, and turn it off in EchoPress?

[SimJeg]

Thanks for the updates, the code is now much more compact and readable I believe. I will review again. For the warning, feel free to re-write the post init entirely:

        assert 0 <= self.compression_ratio < 1, "Compression ratio must be between 0 and 1"
        assert self.chunk_size > 0, "Chunk size must be positive"
        assert not self.kvzip_plus_normalization, "EchoPress does not support the KVzip+ normalization"
        self._reset_internal_parameters()

[SimJeg]

/ok to test 98fee38

[ljwljwljwljw]

Thanks for the updates, the code is now much more compact and readable I believe. I will review again. For the warning, feel free to re-write the post init entirely:

        assert 0 <= self.compression_ratio < 1, "Compression ratio must be between 0 and 1"
        assert self.chunk_size > 0, "Chunk size must be positive"
        assert not self.kvzip_plus_normalization, "EchoPress does not support the KVzip+ normalization"
        self._reset_internal_parameters()

Updated. Thanks!