Add RestoreKV press 🤖🤖🤖

[higokri]

PR description

RestoreKV (https://arxiv.org/abs/2608.01247) is a budget-matched, single-pass
learned restoration plug-in on top of KVzip. Before eviction, 8 restore tokens
attend to the full KV cache in one LoRA-adapted pass, producing a
context-conditioned restore cache that fills part of the budget while the base
evictor fills the rest (budget-matched). Only the restore-token embeddings and
rank-8 LoRA (~0.4% of params) are trained (~2 h), by self-distillation from the
full-cache teacher.
RestoreKV_plus = RestoreKV+ (KVzip+ / WoV-norm scoring).

Additive only (no upstream lines removed):

• kvpress/presses/restorekv_press.py + tests/presses/test_restorekv_press.py
• generic append-KV hooks in base_press.py / pipeline.py (default no-op for
  existing presses) and a _skip_scoring_hooks guard in kvzip_press.py (KVzip
  behavior unchanged) — needed because RestoreKV appends KV entries before eviction
• RestoreKV / RestoreKV_plus in the evaluation registry
• README entry and a TestRestoreKVPress mock in tests/default_presses.py

Results

RULER-4096 average (13 tasks). compression_ratio = fraction of KV removed;
RestoreKV / RestoreKV+ columns.

compression_ratio	Llama-3.1-8B	Qwen3-8B
0.25	95.53 / 95.58	95.25 / 95.07
0.50	95.56 / 95.59	95.20 / 95.13
0.75	95.33 / 95.48	95.22 / 95.15
0.875	94.62 / 94.85	94.24 / 94.61
0.9375	85.78 / 88.21	81.93 / 86.38

Checkpoints are hosted on the HF Hub and loaded automatically — pass an
<org>/<repo>/<file> reference as checkpoint_path
(e.g. higokri/RestoreKV/llama3.1-8b_restorekv.pt) and the press fetches it via
hf_hub_download: https://huggingface.co/higokri/RestoreKV

Checklist

• Tests are working (make test) — full tests/ suite run locally: 624 passed, 53 skipped, 0 failed (skips are flash-attn / resource-gated cases)
• Code is formatted correctly (make style) — black, isort, flake8, mypy --check-untyped-defs, and the SPDX-header check pass on the changed files
• Copyright header is included
• All commits are signed-off using git commit -s
• (new press) restorekv_press.py is in the presses directory
• (new press) RestoreKVPress 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 @higokri,

Thanks for your contribution ! Important refactoring will be needed before we can merge this PR (e.g. don't edit the BasePress or pipeline, make code more compact by for instance use _append_restore_tokens in compress_post etc.). If I understand well inference works as follows:

1. LoRA disabled - Compute the the full cache $C$ on the sequence of hidden states $(h_1, …, h_L)$
2. LoRA enabled - Extend the sequence with $n=8$ learnt vectors $(h_1, …, h_L, e_1, …, e_n)$ and get extended cache $(C, C_{res})$, where $C_{res}$ has length $n$
3. Prune the cache $C$ using KVzip, and concatenate $C_{res}$ to it

I will provide more detailed feedback in the coming days

[SimJeg]

I left a few first comments. Goal is to limit the changes to other presses

[higokri]

Thanks for the detailed review! I pushed a refactor that keeps everything inside
RestoreKVPress — BasePress and KVzipPress are now untouched (0 changes):

• Removed the _perform_kvzip_compression override; _append_restore_tokens now
  runs at the start of compress_post. KVzip already calls compress_post at the
  end of its scoring pass, so the order is unchanged.
• Moved the scoring-skip into a forward_hook override on RestoreKVPress (your
  snippet), so KVzipPress is untouched.
• Removed get_generation_context_length from BasePress/pipeline/press and used
  your suggested if isinstance(press, RestoreKVPress): context_length = cache.get_seq_length() instead.

Net: base_press.py / kvzip_press.py have zero changes; pipeline.py is +1 import,
+1 isinstance. Full make test and make style pass locally.

Your inference summary is exactly right. One detail worth adding: to ensure a fair budget comparison with the other KVpress baselines, the method is budget-matched. When pruning $C$ into $C'$, we evict $n$ extra entries, so the total number of retained entries ($|C'| + n$, where $n$ is the length of $C_{\text{res}}$) equals that of a plain KVzip run at the same compression_ratio. In other words, the restore tokens come at no additional budget cost.

Happy to adjust the pipeline isinstance if you find a cleaner approach. I've also submitted results to the Hugging Face leaderboard Space, in case they're useful for reference. Thanks again!

[SimJeg]

Hi @higokri, thanks for this first batch of updates. A second batch of updates to make the code more readable would be to use the PEFT library (and add peft>=0.20.0 to pyproject.toml). It would allow to go from ~280 loc to ~110 loc but it implies for you to re-publish the checkpoints in the right format. Would you be willing to do this ?

Please find below a proposed implementation (⚠️ not tested, output form Cursor). It also:

• renamed _prepare_model to post_init_from_model (standard in kvpress)
• implies to add self.post_init_from_model(model) at the beginning of the __call__ function in KVzipPress as in the BasePress
• removed __post_init__
• renamed a few variables and remove _ prefix as it's not used in kvpress, removed budget_matched variable (as it's always true) and

# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import logging
from dataclasses import dataclass, field

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from torch import nn
from transformers import PreTrainedModel

from kvpress.presses.kvzip_press import KVzipPress

logger = logging.getLogger(__name__)


@dataclass
class RestoreKVPress(KVzipPress):
    """RestoreKV (https://arxiv.org/abs/2608.01247): learned restoration on top of KVzip.

    Before eviction, n=8 restore tokens attend to the full KV cache in one
    LoRA-adapted pass, producing a context-conditioned restore cache that fills
    n*L*H slots of the same budget while the base evictor fills the rest
    (budget-matched). LoRA is active only for this one-time pass. Only the restore
    embeddings and LoRA (~0.4%) are trained, by self-distillation from the
    full-cache teacher.
    """

    restore_embeddings: torch.Tensor | None = field(init=False, default=None, repr=False)
    restore_model_name: str | None = field(init=False, default=None)
    encoding_restore_tokens: bool = field(init=False, default=False, repr=False)

    @property
    def num_restore_tokens(self) -> int:
        if self.restore_embeddings is None:
            return 0
        return self.restore_embeddings.shape[0]

    def post_init_from_model(self, model: PreTrainedModel):
        model_name = model.config.name_or_path.split("/")[-1]
        variant = "_plus" if self.kvzip_plus_normalization else ""
        adapter_name = f"restorekv{variant}"
        restore_model_name = f"higokri/RestoreKV-{model_name}{variant}"
        if restore_model_name == self.restore_model_name:
            return

        weights_path = hf_hub_download(restore_model_name, "adapter_model.safetensors")
        adapter_state = load_file(weights_path)
        self.restore_embeddings = adapter_state.pop("restore_embeddings").to(model.device, dtype=model.dtype)
        if adapter_name not in getattr(model, "peft_config", {}):
            model.load_adapter(
                restore_model_name,
                adapter_name=adapter_name,
                adapter_state_dict=adapter_state,
            )
        model.disable_adapters()
        self.restore_model_name = restore_model_name
        logger.info("Loaded %s with %d restore tokens", restore_model_name, self.num_restore_tokens)

    def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dict, output: list):
        # Skip KVzip scoring while the restore tokens are being encoded (see append_restore_tokens).
        if self.encoding_restore_tokens:
            return output
        return super().forward_hook(module, input, kwargs, output)

    def append_restore_tokens(self, model: PreTrainedModel):
        """Single LoRA-adapted restore pass over the full cache (before eviction);
        appends the n restore tokens' K/V in place. Scoring hooks are skipped."""
        if self._cache is None or self.restore_embeddings is None:
            raise RuntimeError("Restore append requires a populated context cache")
        before = self._cache.get_seq_length()
        restore_embeds = self.restore_embeddings.unsqueeze(0)
        cache_position = torch.arange(
            before,
            before + self.num_restore_tokens,
            device=restore_embeds.device,
        )
        self.encoding_restore_tokens = True
        model.set_adapter("restorekv_plus" if self.kvzip_plus_normalization else "restorekv")
        model.enable_adapters()
        try:
            with torch.inference_mode():
                model.model(
                    inputs_embeds=restore_embeds,
                    past_key_values=self._cache,
                    position_ids=cache_position.unsqueeze(0),
                    cache_position=cache_position,
                    use_cache=True,
                )
        finally:
            model.disable_adapters()
            self.encoding_restore_tokens = False

        after = self._cache.get_seq_length()
        if after - before != self.num_restore_tokens:
            raise RuntimeError(
                "Restore append length mismatch: " f"expected {self.num_restore_tokens}, observed {after - before}"
            )

    def compress_post(self, model: PreTrainedModel):
        # KVzip has finished scoring the context; append the restore tokens to the
        # full cache (budget-matched), then let KVzip evict the lowest-scored pairs.
        self.append_restore_tokens(model)
        requested_ratio = self.compression_ratio
        if self.context_length > 0:
            restore_overhead = self.num_restore_tokens / self.context_length
            self.compression_ratio = min(1.0, requested_ratio + restore_overhead)
        try:
            super().compress_post(model)
        finally:
            self.compression_ratio = requested_ratio

[higokri]

Hi @SimJeg, thanks for the proposed implementation! This is a great suggestion. I'll rework the press with PEFT following your snippet, re-publish the checkpoints in the adapter format, and let you know once everything is tested and pushed.

[higokri]

Follow-up on the PEFT rework: restorekv_press.py is down from ~280 to ~110 loc, covering all your points:

• ✅ PEFT adapters, with peft>=0.20.0 added to pyproject.toml
• ✅ _prepare_model renamed to post_init_from_model (plus the self.post_init_from_model(model) call at the start of KVzipPress.__call__)
• ✅ removed __post_init__ and budget_matched, and dropped the _ prefixes
• ✅ checkpoints re-published as PEFT adapters at higokri/RestoreKV-<model>[_plus]

The refactored code runs end to end. I'm currently re-running the RULER benchmark to confirm the scores match the previous version, and will report back once it finishes.

One design question before I push: where to store the 8 restore-token embeddings. Both options work, so it's purely a matter of cleanliness:

• Separate file (what I did): keep adapter_model.safetensors as a clean, standard PEFT adapter, with the embeddings in a separate restore_embeddings.safetensors (66 KB) in the same repo.
• Single file: bundle the embeddings into the adapter file and pop them at load time. This also works, but load_adapter then emits an unexpected keys: restore_embeddings warning (cosmetic; the LoRA still loads correctly), and the adapter is no longer a standard PEFT adapter.

Do you prefer the separate file or everything in one? Either works for me.

[SimJeg]

@higokri I think you did not push. Also you can remove updates in evaluate.py too.
For files it's up to you, what matters to use is readability of the codebase.

[higokri]

@SimJeg Pushed (evaluate.py reverted to upstream too). Verified the outputs match the previous implementation. Thanks for the detailed guidance throughout!

[SimJeg]

/ok to test 0a56c43

[SimJeg]

Here is a final batch of review. Once fixed we're ready for a merge :) thanks for your reactivity. Please notify us if you need to update the predictions for the leadeboard.

Also you can simply push on your branch instead of push force so that we can keep commit history of this PR

[higokri]

Applied the final batch, pushed as a new commit (no force-push) to keep the history:

• README + docstring updates
• removed the guard / length checks
• inlined the embeddings
• switched to self.context_length
• dropped the standalone test
• grouped the adapters in an HF collection

I'll also clean up the old higokri/RestoreKV repo soon.

For the leaderboard: our RULER results are submitted at https://huggingface.co/spaces/nvidia/kvpress-leaderboard/discussions/18. Could you take a look and update the leaderboard with those when you get a chance? Thanks!

[SimJeg]

/ok to test ce3cbc4

[SimJeg]

Thanks for your contribution @higokri. LB has been updated too