From 45177906dde2c1cd452f00c33e524df856c29f61 Mon Sep 17 00:00:00 2001 From: tangziyao <672208690@qq.com> Date: Mon, 15 Jun 2026 00:36:26 +0800 Subject: [PATCH 1/3] =?UTF-8?q?Add=20LUKVPress=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: tangziyao <672208690@qq.com> --- README.md | 1 + evaluation/evaluate_registry.py | 2 + kvpress/__init__.py | 2 + kvpress/presses/lukv_press.py | 189 ++++++++++++++++++++++++++++++++ tests/default_presses.py | 14 +++ 5 files changed, 208 insertions(+) create mode 100644 kvpress/presses/lukv_press.py diff --git a/README.md b/README.md index 2a2cd9a5..665e4d3b 100644 --- a/README.md +++ b/README.md @@ -134,6 +134,7 @@ Some presses rely on a different logic: Finally we provide wrapper presses that can be combined with other presses: - `AdaKVPress` ([source](kvpress/presses/adakv_press.py), [paper](https://arxiv.org/abs/2407.11550)): prune bottom scores of any `ScorerPress` but across all heads, achieving head-wise compressions +- `LUKVPress` ([source](kvpress/presses/lukv_press.py), [paper](https://arxiv.org/abs/2602.08585)): applies layer/head budget curves to a `ScorerPress` - `PerLayerCompressionPress` ([source](kvpress/presses/per_layer_compression_press.py)): compress each layer with a different compression ratio (experimental) - `ComposedPress` ([source](kvpress/presses/composed_press.py)): compose multiple presses together by chaining their forward hooks - `KeyRerotationPress` ([source](kvpress/presses/key_rerotation_press.py)): rerotate pruned keys to have continuous RoPE embeddings diff --git a/evaluation/evaluate_registry.py b/evaluation/evaluate_registry.py index 629f93eb..1a8b47f1 100644 --- a/evaluation/evaluate_registry.py +++ b/evaluation/evaluate_registry.py @@ -34,6 +34,7 @@ KVzapPress, KVzipPress, LagKVPress, + LUKVPress, MergingPress, ObservedAttentionPress, PyramidKVPress, @@ -98,6 +99,7 @@ "kvzap_mlp_head": KVzapPress(model_type="mlp"), "kvzap_mlp_layer": AdaKVPress(KVzapPress(model_type="mlp")), "lagkv": LagKVPress(), + "lukv": LUKVPress(ExpectedAttentionPress(epsilon=2e-2), sink=4, window=1), "knorm": KnormPress(), "observed_attention": ObservedAttentionPress(), "pyramidkv": PyramidKVPress(), diff --git a/kvpress/__init__.py b/kvpress/__init__.py index c89fbee2..35735833 100644 --- a/kvpress/__init__.py +++ b/kvpress/__init__.py @@ -30,6 +30,7 @@ from kvpress.presses.kvzip_press import KVzipPress from kvpress.presses.lagkv_press import LagKVPress from kvpress.presses.leverage_press import LeverageScorePress +from kvpress.presses.lukv_press import LUKVPress from kvpress.presses.merging_press import MergingPress from kvpress.presses.non_causal_attention_press import NonCausalAttnPress from kvpress.presses.observed_attention_press import ObservedAttentionPress @@ -91,4 +92,5 @@ "KVComposePress", "MergingPress", "CapPress", + "LUKVPress", ] diff --git a/kvpress/presses/lukv_press.py b/kvpress/presses/lukv_press.py new file mode 100644 index 00000000..06e6e12a --- /dev/null +++ b/kvpress/presses/lukv_press.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field +from io import BytesIO +from typing import Optional + +import numpy as np +import requests # type: ignore[import-untyped] +import torch +from cachetools import LRUCache, cached # type: ignore[import-untyped] +from torch import nn +from transformers import PreTrainedModel + +from kvpress.presses.base_press import BasePress +from kvpress.presses.expected_attention_press import ExpectedAttentionPress +from kvpress.presses.scorer_press import ScorerPress + +LLAMA_3_1_8B_EA_CURVE_URL = ( + "https://raw.githubusercontent.com/baidu-baige/LU-KV/main/" + "evaluation/curve_data/llama-3.1-8b/ea_0.02_sink4_win1_llama_avg_ratio.npy" +) + +BUDGET_CURVE_URLS = { + # ExpectedAttentionPress(epsilon=2e-2), sink=4, window=1 + ("meta-llama/Llama-3.1-8B", "ExpectedAttentionPress"): LLAMA_3_1_8B_EA_CURVE_URL, + ("meta-llama/Llama-3.1-8B-Instruct", "ExpectedAttentionPress"): LLAMA_3_1_8B_EA_CURVE_URL, + ("meta-llama/Meta-Llama-3.1-8B-Instruct", "ExpectedAttentionPress"): LLAMA_3_1_8B_EA_CURVE_URL, +} + +cache = LRUCache(maxsize=128) + + +@cached(cache, key=lambda url: url) +def load_budget_curve(url: str) -> np.ndarray: + response = requests.get(url) + response.raise_for_status() + return np.load(BytesIO(response.content), allow_pickle=False) + + +@dataclass +class LUKVPress(BasePress): + """ + LU-KV: head-wise budget allocation around a score-based press. + + LU-KV wraps a ``ScorerPress`` and uses a pre-computed budget curve to allocate + different token budgets to each attention layer and KV head. The default + configuration is ``LUKVPress(ExpectedAttentionPress(epsilon=2e-2), sink=4, window=1)``. + Based on Predicting Future Utility: Global Combinatorial Optimization for + Task-Agnostic KV Cache Eviction (https://arxiv.org/abs/2602.08585). + + Budget curves are model- and scorer-specific. To add a new curve, upload the + ``.npy`` file with shape ``[99, num_layers, num_kv_heads]`` and add an entry + to ``BUDGET_CURVE_URLS`` keyed by ``(model.config.name_or_path, ScorerPressName)``. + The default curve is for ``ExpectedAttentionPress(epsilon=2e-2), sink=4, window=1`` + on Llama-3.1-8B. + + Parameters + ---------- + press : ScorerPress, default=ExpectedAttentionPress(epsilon=2e-2) + The scoring method used to rank cached tokens within each KV head. + compression_ratio : float, default=0.0 + Fraction of KV pairs to remove globally. + sink : int, default=4 + Number of initial tokens to protect from eviction. + window : int, default=1 + Number of most recent tokens to protect from eviction. + """ + + press: ScorerPress = field(default_factory=lambda: ExpectedAttentionPress(epsilon=2e-2)) + compression_ratio: float = 0.0 + sink: int = 4 + window: int = 1 + + _budget_curves: Optional[np.ndarray] = field(init=False, repr=False, default=None) + + def __post_init__(self): + assert isinstance(self.press, ScorerPress), "LUKVPress requires a ScorerPress as input" + assert 0 <= self.compression_ratio < 1, "Compression ratio must be between 0 and 1" + assert self.sink >= 0, "sink must be non-negative" + assert self.window >= 0, "window must be non-negative" + + def post_init_from_model(self, model: PreTrainedModel): + self.press.post_init_from_model(model) + if self._budget_curves is None: + self._budget_curves = self._load_budget_curves(model) + + def _load_budget_curves(self, model: PreTrainedModel) -> np.ndarray: + model_name = model.config.name_or_path + press_name = type(self.press).__name__ + url = BUDGET_CURVE_URLS.get((model_name, press_name)) + if url is None: + raise KeyError( + f"No LU-KV budget curve found for model={model_name!r} and press={press_name!r}. " + f"Available curves: {list(BUDGET_CURVE_URLS.keys())}." + ) + + budget_curves = load_budget_curve(url) + num_layers, num_key_value_heads = self._model_curve_shape(model) + expected_shape = (99, num_layers, num_key_value_heads) + if budget_curves.shape != expected_shape: + raise ValueError(f"LU-KV budget curve must have shape {expected_shape}, got {budget_curves.shape}.") + + return budget_curves + + def _model_curve_shape(self, model: PreTrainedModel) -> tuple[int, int]: + num_layers = getattr(model.config, "num_hidden_layers", None) + num_key_value_heads = getattr(model.config, "num_key_value_heads", None) + if num_key_value_heads is None: + num_key_value_heads = getattr(model.config, "num_attention_heads", None) + if num_layers is None or num_key_value_heads is None: + raise ValueError("LU-KV requires num_hidden_layers and num_key_value_heads or num_attention_heads.") + return int(num_layers), int(num_key_value_heads) + + 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]: + if self.compression_ratio <= 0: + return keys, values + if self._budget_curves is None: + raise ValueError("LU-KV budget curves are not loaded. Use LUKVPress as a model context manager first.") + assert module.config._attn_implementation != "eager", "eager mode not supported" + + bsz, num_key_value_heads, seq_len, _ = keys.shape + scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs) + + protected_score = scores.max().item() + 1 + if self.sink > 0: + safe_sink = min(self.sink, seq_len) + scores[..., :safe_sink] = protected_score + if self.window > 0: + window_start = max(0, seq_len - self.window) + scores[..., window_start:] = protected_score + + target_idx = int(round(self.compression_ratio * 100)) - 1 + target_idx = max(0, min(98, target_idx)) + layer_idx = module.layer_idx + try: + local_prune_ratios = torch.as_tensor( + self._budget_curves[target_idx, layer_idx], + device=keys.device, + dtype=torch.float32, + ) + except IndexError as exc: + raise ValueError( + f"LU-KV budget curve does not contain target index {target_idx} and layer {layer_idx}." + ) from exc + + if local_prune_ratios.shape[0] != num_key_value_heads: + raise ValueError( + "LU-KV budget curve KV-head count does not match current keys: " + f"curve has {local_prune_ratios.shape[0]}, keys have {num_key_value_heads}." + ) + + head_keep_rates = (1.0 - local_prune_ratios).clamp(min=0.0, max=1.0) + ideal_keep_counts = head_keep_rates * seq_len + total_keep_target = int(torch.round(ideal_keep_counts.sum()).item()) + total_keep_target = max(num_key_value_heads, min(num_key_value_heads * seq_len, total_keep_target)) + base_keep_counts = torch.floor(ideal_keep_counts).long() + remainder = total_keep_target - int(base_keep_counts.sum().item()) + + if remainder > 0: + fractional_parts = ideal_keep_counts - base_keep_counts + top_k_indices = torch.topk(fractional_parts, k=min(remainder, num_key_value_heads)).indices + base_keep_counts[top_k_indices] += 1 + + final_keep_per_head = base_keep_counts.clamp_(min=1, max=seq_len) + num_to_prune_per_head = seq_len - final_keep_per_head + + if torch.all(num_to_prune_per_head <= 0): + module.masked_key_indices = None + return keys, values + + sorted_indices = torch.argsort(scores, dim=-1, descending=True, stable=True) + rank = torch.arange(seq_len, device=scores.device).view(1, 1, seq_len).expand_as(sorted_indices) + keep_mask = rank < final_keep_per_head.view(1, num_key_value_heads, 1) + prune_mask = ~keep_mask + + batch_indices, head_indices, rank_indices = torch.where(prune_mask) + pruned_seq_indices = sorted_indices[batch_indices, head_indices, rank_indices] + module.masked_key_indices = (batch_indices, head_indices, pruned_seq_indices) # type: ignore[assignment] + + return keys, values diff --git a/tests/default_presses.py b/tests/default_presses.py index 61f9eeed..fcc8d203 100644 --- a/tests/default_presses.py +++ b/tests/default_presses.py @@ -18,6 +18,7 @@ KVzipPress, LagKVPress, LeverageScorePress, + LUKVPress, NonCausalAttnPress, PyramidKVPress, QFilterPress, @@ -68,6 +69,18 @@ def post_init_from_model(self, model): self.gates.append(module) +class TestLUKVPress(LUKVPress): + """Test version of LUKVPress that creates a mock budget curve instead of downloading one.""" + + def post_init_from_model(self, model): + self.press.post_init_from_model(model) + if self._budget_curves is None: + n_layers = model.config.num_hidden_layers + n_heads = model.config.num_key_value_heads + prune_ratios = np.arange(1, 100, dtype=np.float32).reshape(99, 1, 1) / 100 + self._budget_curves = np.broadcast_to(prune_ratios, (99, n_layers, n_heads)).copy() + + # contains all presses to be tested # kwargs should be ordered easy to hard compression default_presses = [ @@ -153,4 +166,5 @@ def post_init_from_model(self, model): ], }, {"cls": CapPress, "kwargs": [{"compression_ratio": 0.5}, {"compression_ratio": 0.8}]}, + {"cls": TestLUKVPress, "kwargs": [{"compression_ratio": 0.5}, {"compression_ratio": 0.8}]}, ] From 7f783364ab6f8634b0665242971212a17f9558ee Mon Sep 17 00:00:00 2001 From: tangziyao <672208690@qq.com> Date: Wed, 1 Jul 2026 17:12:57 +0800 Subject: [PATCH 2/3] =?UTF-8?q?Document=20LUKV=20curve=20generation=20?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: tangziyao <672208690@qq.com> --- kvpress/presses/lukv_press.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/kvpress/presses/lukv_press.py b/kvpress/presses/lukv_press.py index 06e6e12a..a21f7969 100644 --- a/kvpress/presses/lukv_press.py +++ b/kvpress/presses/lukv_press.py @@ -55,6 +55,15 @@ class LUKVPress(BasePress): The default curve is for ``ExpectedAttentionPress(epsilon=2e-2), sink=4, window=1`` on Llama-3.1-8B. + New budget curves can be generated with the official LU-KV repository: + https://github.com/baidu-baige/LU-KV. In that repository, run + ``cd evaluation && bash curve_data/generate_curve.sh`` after setting + ``MODEL_TYPE``, ``MODEL_PATH``, ``DATASET_PATH``, ``METHOD_CONFIGS``, and + ``LAYERWISE_FLAG`` in ``evaluation/curve_data/generate_curve.sh``. The script + first extracts per-context scorer statistics with ``step1_.py`` and + then computes the averaged layer/head budget curve with + ``step2_compute_curve.py``. + Parameters ---------- press : ScorerPress, default=ExpectedAttentionPress(epsilon=2e-2) @@ -140,7 +149,7 @@ def compress( target_idx = int(round(self.compression_ratio * 100)) - 1 target_idx = max(0, min(98, target_idx)) - layer_idx = module.layer_idx + layer_idx = int(getattr(module, "layer_idx")) try: local_prune_ratios = torch.as_tensor( self._budget_curves[target_idx, layer_idx], From f4572382b1da6d7615eac83cf91a8e329ebe2609 Mon Sep 17 00:00:00 2001 From: tangziyao <672208690@qq.com> Date: Wed, 1 Jul 2026 18:40:10 +0800 Subject: [PATCH 3/3] =?UTF-8?q?Keep=20one=20LUKV=20curve=20mapping=20?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: tangziyao <672208690@qq.com> --- kvpress/presses/lukv_press.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/kvpress/presses/lukv_press.py b/kvpress/presses/lukv_press.py index a21f7969..5dec73ae 100644 --- a/kvpress/presses/lukv_press.py +++ b/kvpress/presses/lukv_press.py @@ -23,9 +23,7 @@ BUDGET_CURVE_URLS = { # ExpectedAttentionPress(epsilon=2e-2), sink=4, window=1 - ("meta-llama/Llama-3.1-8B", "ExpectedAttentionPress"): LLAMA_3_1_8B_EA_CURVE_URL, ("meta-llama/Llama-3.1-8B-Instruct", "ExpectedAttentionPress"): LLAMA_3_1_8B_EA_CURVE_URL, - ("meta-llama/Meta-Llama-3.1-8B-Instruct", "ExpectedAttentionPress"): LLAMA_3_1_8B_EA_CURVE_URL, } cache = LRUCache(maxsize=128)