From 39748e9c8301944128eb2e681bfb074e9a2de860 Mon Sep 17 00:00:00 2001 From: Jiawei Lin Date: Fri, 11 Sep 2026 14:43:00 +0200 Subject: [PATCH] Add EchoPress Approximate KVzip scores through virtual context reconstruction during prefill. Reconstruct the first chunk exactly and calibrate later virtual scores per request, layer, and KV head. Signed-off-by: Jiawei Lin --- README.md | 1 + evaluation/evaluate_registry.py | 2 + kvpress/__init__.py | 2 + kvpress/presses/echo_press.py | 335 ++++++++++++++++++++++++++ kvpress/presses/echo_press_kernels.py | 198 +++++++++++++++ tests/default_presses.py | 17 ++ tests/presses/test_echo_press.py | 109 +++++++++ 7 files changed, 664 insertions(+) create mode 100644 kvpress/presses/echo_press.py create mode 100644 kvpress/presses/echo_press_kernels.py create mode 100644 tests/presses/test_echo_press.py diff --git a/README.md b/README.md index 1b425f04..0ae3ab45 100644 --- a/README.md +++ b/README.md @@ -130,6 +130,7 @@ Some presses rely on a different logic: - `KVzipPress` ([source](kvpress/presses/kvzip_press.py), [paper](https://arxiv.org/abs/2505.23416)): identify redundant KV pairs through context reconstruction. Achieve near-lossless compression at the cost of multiple forward passes. - `KVgradPress` ([source](kvpress/presses/kvgrad_press.py), [paper](https://openreview.net/forum?id=cg1wTCJjjk)): similar to `KVzipPress`, but scores KV pairs by an input × gradient attribution of their effect on the last hidden states, at the cost of one backward pass per chunk. - `RestoreKVPress` ([source](kvpress/presses/restorekv_press.py), [paper](https://arxiv.org/abs/2608.01247)): extends the KV cache with 8 learned restore tokens encoded by a LoRA module before budget-matched KVzip pruning. +- `EchoPress` ([source](kvpress/presses/echo_press.py), [paper](https://arxiv.org/abs/2610.00412)): approximate KVzip scores by re-rotating prefill queries and keys to the positions of a virtual repetition. - `KVComposePress` ([source](kvpress/presses/kvcompose_press.py), [paper](https://arxiv.org/abs/2509.05165)): attention-guided eviction, aligning per-head selections into composite tokens to preserve cache structure. > [!NOTE] diff --git a/evaluation/evaluate_registry.py b/evaluation/evaluate_registry.py index 9552e9b0..16ba6a18 100644 --- a/evaluation/evaluate_registry.py +++ b/evaluation/evaluate_registry.py @@ -26,6 +26,7 @@ DMSPress, DropKVPress, DuoAttentionPress, + EchoPress, EntropyGatedChunkKVPress, ExpectedAttentionPress, FastKVzipPress, @@ -102,6 +103,7 @@ "kvgrad": KVgradPress(), "kvzip": KVzipPress(), "kvzip_plus": KVzipPress(kvzip_plus_normalization=True), + "echo": EchoPress(), # Exact first chunk with per-request, per-head virtual-to-exact calibration. "kvzap_linear": DMSPress(press=KVzapPress(model_type="linear")), "kvzap_mlp": DMSPress(press=KVzapPress(model_type="mlp")), "kvzap_mlp_head": KVzapPress(model_type="mlp"), diff --git a/kvpress/__init__.py b/kvpress/__init__.py index 40b7cb58..e3c34001 100644 --- a/kvpress/__init__.py +++ b/kvpress/__init__.py @@ -20,6 +20,7 @@ from kvpress.presses.dms_press import DMSPress from kvpress.presses.dropkv_press import DropKVPress from kvpress.presses.duo_attention_press import DuoAttentionPress +from kvpress.presses.echo_press import EchoPress from kvpress.presses.entropy_gated_chunkkv_press import EntropyGatedChunkKVPress from kvpress.presses.expected_attention_press import ExpectedAttentionPress from kvpress.presses.expected_attention_with_stats import ExpectedAttentionStatsPress @@ -87,6 +88,7 @@ "KeyDiffPress", "KVgradPress", "KVzipPress", + "EchoPress", "ExpectedAttentionStatsPress", "CAMPress", "DecodingPress", diff --git a/kvpress/presses/echo_press.py b/kvpress/presses/echo_press.py new file mode 100644 index 00000000..a01124fd --- /dev/null +++ b/kvpress/presses/echo_press.py @@ -0,0 +1,335 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import math +from contextlib import contextmanager +from dataclasses import dataclass, field +from functools import partial +from types import MethodType +from typing import Generator, List + +import torch +from transformers import AutoTokenizer, Gemma3PreTrainedModel, PreTrainedModel, PreTrainedTokenizer +from transformers.models.llama.modeling_llama import rotate_half +from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention + +from kvpress.presses.base_press import SUPPORTED_MODELS +from kvpress.presses.echo_press_kernels import HAS_TRITON, virtual_scores_triton +from kvpress.presses.kvzip_press import KVzipPress +from kvpress.utils import extract_keys_and_values + +logger = logging.getLogger(__name__) + +QUERY_BLOCK = 1024 # Limit the materialized attention block size. + + +def _rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """x: [b, h, t, d]; cos/sin: [b, t, d]""" + return x * cos.unsqueeze(1) + rotate_half(x) * sin.unsqueeze(1) + + +def _scores_from_logits(s_sink, s_cached, s_copy, causal): + """Return each cached key's maximum attention over the query block.""" + s_cached = s_cached.float() + s_copy = s_copy.float().masked_fill(causal, float("-inf")) + row_max = torch.maximum(s_cached.amax(-1, keepdim=True), s_copy.amax(-1, keepdim=True)) + if s_sink is not None: + s_sink = s_sink.float() + row_max = torch.maximum(row_max, s_sink.amax(-1, keepdim=True)) + p_cached = (s_cached - row_max).exp() + denom = p_cached.sum(-1, keepdim=True) + (s_copy - row_max).exp().sum(-1, keepdim=True) + if s_sink is not None: + denom = denom + (s_sink - row_max).exp().sum(-1, keepdim=True) + return (p_cached / denom).amax(dim=(0, 2, 3)) + + +_SCORERS: dict = {} + + +def _get_scorer(compile_scoring: bool): + """Compile the scoring step once, falling back to eager mode.""" + if not compile_scoring: + return _scores_from_logits + if "compiled" not in _SCORERS: + try: + _SCORERS["compiled"] = torch.compile(_scores_from_logits, dynamic=True) + except Exception as e: # noqa: BLE001 + logger.warning(f"torch.compile unavailable ({e}); using the eager scoring step") + _SCORERS["compiled"] = _scores_from_logits + return _SCORERS["compiled"] + + +@dataclass +class EchoPress(KVzipPress): + """ + Approximate KVzip scores during prefill using virtual context repetition. + + RoPE composition lets EchoPress score a repeated chunk from its prefill queries and + cached keys, avoiding reconstruction passes for later chunks. EchoPress reconstructs + the first chunk exactly and uses its paired virtual and exact scores to calibrate later + virtual scores independently for every request, layer, and KV head. Eviction behavior + is inherited from KVzipPress. + + Only batch size 1 and a single full-context prefill are supported. + + Parameters + ---------- + compression_ratio : float, default=0.0 + Fraction of key-value pairs to remove during compression. + layerwise : bool, default=False + Whether to enable uniform compression ratios across layers (see KVzipPress). + n_sink : int, default=4 + Number of initial tokens to preserve as attention sinks. + chunk_size : int, default=2048 + Total length of the first exact pass and size of later virtual chunks. + exact_first_chunk : bool, default=True + Re-score the first chunk with one reconstruction pass. Disable only to reproduce + the legacy all-virtual variant. + guard_tokens : int, default=0 + Number of tokens to retain at each chunk boundary. + score_calibration : bool, default=True + Apply a per-request score map fitted on the first chunk. + calibration_direction : str, default="virtual_to_exact" + Direction of the per-sample quantile map: "exact_to_virtual" maps the + exact first chunk onto the virtual scale, while "virtual_to_exact" maps + later virtual chunks onto the exact first-chunk scale. + calibration_scope : str, default="head" + Quantile-map granularity: "global", "layer", or "head". + scoring_backend : str, default="auto" + One of "auto", "triton", "torch", or "torch_eager". + """ + + chunk_size: int = 2048 + exact_first_chunk: bool = True + guard_tokens: int = 0 + score_calibration: bool = True + calibration_direction: str = "virtual_to_exact" + calibration_scope: str = "head" + scoring_backend: str = "auto" + _q: dict = field(default_factory=dict, repr=False) # Prefill queries by layer. + _v0: dict = field(default_factory=dict, repr=False) # First-chunk virtual scores by layer. + _chunks: list = field(default_factory=list, repr=False) # (start, end, prompt length) + _pairs: list = field(default_factory=list, repr=False) # (chunk IDs, repeated IDs) + + def __post_init__(self): + assert 0 <= self.compression_ratio < 1, "Compression ratio must be between 0 and 1" + assert self.scoring_backend in ("auto", "triton", "torch", "torch_eager"), "unknown scoring_backend" + assert self.calibration_direction in ( + "exact_to_virtual", + "virtual_to_exact", + ), "unknown calibration_direction" + assert self.calibration_scope in ("global", "layer", "head"), "unknown calibration_scope" + self._reset_internal_parameters() + + def _init_prompt_ids(self, tokenizer: PreTrainedTokenizer): + """Find the chat template's prefix and suffix.""" + if tokenizer.chat_template is None: + prefix_text, suffix_text = "", "\n" + else: + dummy_context = "dummy context" + separator = "\n" + "#" * len(dummy_context) + temp_context = tokenizer.apply_chat_template( + [{"role": "user", "content": dummy_context + separator}], + add_generation_prompt=True, + tokenize=False, + enable_thinking=False, + ) + context, suffix_text = temp_context.split(separator) + prefix_text = context.split(dummy_context)[0] + self.prefix_length = tokenizer.encode(prefix_text, return_tensors="pt", add_special_tokens=False).shape[-1] + self._suffix_ids = tokenizer.encode(suffix_text, return_tensors="pt", add_special_tokens=False) + + def prepare( + self, + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + prev_postfix_size=8, + ) -> List[tuple[torch.Tensor, torch.Tensor]]: + """Build chunks, keeping the first repeated sequence at `chunk_size`.""" + ctx_ids = self._context_ids[:, self.prefix_length :].to("cpu") + self.score_val = torch.zeros( + (model.config.num_hidden_layers, 1, model.config.num_key_value_heads, self.context_length), + dtype=model.dtype, + device=model.device, + ) + self.score_val[..., : self.n_sink] = 1.0 + + prompts = [ + "\n\nRepeat the previous context exactly.", + "\n\nRepeat the part of the previous context exactly, starting with", + ] + q_ids = [tokenizer.encode(p, return_tensors="pt", add_special_tokens=False) for p in prompts] + first_size = max(1, self.chunk_size - q_ids[0].shape[1] - self._suffix_ids.shape[1]) + chunks = [ctx_ids[:, :first_size]] + self._chunk_fn(ctx_ids[:, first_size:], self.chunk_size) + chunks = [c for c in chunks if c.shape[1] > 0] + pairs = [] + for i, a_ids in enumerate(chunks): + if i == 0: + prompt_ids = q_ids[0] + else: + prompt_ids = torch.cat([q_ids[1], chunks[i - 1][:, -prev_postfix_size:]], dim=1) + pairs.append((a_ids, torch.cat([prompt_ids, self._suffix_ids, a_ids], dim=1))) + return pairs + + @contextmanager + def __call__(self, model: PreTrainedModel) -> Generator: + if not isinstance(model, SUPPORTED_MODELS): + logger.warning(f"Model {type(model)} not tested, supported models: {SUPPORTED_MODELS}") + if isinstance(model, Gemma3PreTrainedModel): + raise ValueError("EchoPress is not supported for Gemma3ForCausalLM") + if not all(hasattr(layer.self_attn, "q_proj") for layer in model.model.layers): + raise NotImplementedError("EchoPress requires attention modules with a q_proj") + + self.post_init_from_model(model) + tokenizer = AutoTokenizer.from_pretrained(model.config.name_or_path) + self._init_prompt_ids(tokenizer) + self._inv_freq = model.model.rotary_emb.inv_freq + + original_forward = model.model.forward + + def wrapped_forward(model_self, *args, **kwargs): + # Prepare chunks before the layer hooks run. + self._context_ids = kwargs["input_ids"] + self._cache = kwargs["past_key_values"] + self.context_length = self._context_ids.shape[1] + self._pairs = self.prepare(model, tokenizer) + start, self._chunks = self.prefix_length, [] + for chunk_ids, repeat_ids in self._pairs: + self._chunks.append((start, start + chunk_ids.shape[1], repeat_ids.shape[1] - chunk_ids.shape[1])) + start += chunk_ids.shape[1] + return original_forward(*args, **kwargs) + + model.model.forward = MethodType(wrapped_forward, model.model) + handles = [] + if self.compression_ratio > 0: + for layer_idx, layer in enumerate(model.model.layers): + handles.append(layer.self_attn.q_proj.register_forward_hook(partial(self._q_hook, layer_idx=layer_idx))) + handles.append(layer.self_attn.register_forward_hook(self._prefill_hook, with_kwargs=True)) + try: + try: + yield + finally: + model.model.forward = original_forward + for handle in handles: + handle.remove() + if self.compression_ratio > 0 and self._context_ids is not None: + with torch.no_grad(): + self._finalize(model) + self.compress_post(model) + finally: + self._q.clear() + self._v0.clear() + self._chunks, self._pairs = [], [] + self._reset_internal_parameters() + + def _q_hook(self, module, args, output, layer_idx: int): + """Save projected queries until the attention hook runs.""" + self._q[layer_idx] = output + + def _prefill_hook(self, module, args, kwargs, output): + """Compute each chunk's virtual scores for one layer.""" + layer_idx = int(module.layer_idx) + q = self._q.pop(layer_idx) + assert q.shape[0] == 1, "EchoPress only supports batch size 1" + assert q.shape[1] == self.context_length, "EchoPress requires a single prefill forward" + q = q.view(1, q.shape[1], module.config.num_attention_heads, module.head_dim).transpose(1, 2) + if isinstance(module, Qwen3Attention): + q = module.q_norm(q) + cos, sin = kwargs["position_embeddings"] + q = _rope(q, cos, sin) # [1, H, n, d] + keys, _ = extract_keys_and_values(self._cache, layer_idx) # [1, Hkv, n, d] + with torch.no_grad(): + for i, (start, end, n_prompt) in enumerate(self._chunks): + scores = self._chunk_scores(module, q, keys, start, end, n_prompt) + self.score_val[layer_idx, 0, :, start:end] = scores.to(self.score_val.dtype) + if i == 0: + self._v0[layer_idx] = scores + + def _rotate(self, x: torch.Tensor, delta: int) -> torch.Tensor: + """Apply a RoPE rotation for a constant position offset.""" + freqs = delta * self._inv_freq.float() + emb = torch.cat([freqs, freqs]) + return x * emb.cos() + rotate_half(x) * emb.sin() + + def _chunk_scores(self, module, q: torch.Tensor, keys: torch.Tensor, start: int, end: int, n_prompt: int): + """Return virtual KVzip scores for the chunk `[start, end)`.""" + m = end - start + sink = min(self.n_sink, start) + num_heads_kv = keys.shape[1] + num_groups = q.shape[1] // num_heads_kv + head_dim = module.head_dim + delta = self.context_length + n_prompt - start + # Match KVzip's bf16 logits; normalization remains fp32. + q1 = q[:, :, start:end].view(1, num_heads_kv, num_groups, m, head_dim) + q2 = self._rotate(q1.float(), delta).to(q1.dtype) + backend = self.scoring_backend + if backend == "auto": + backend = "triton" if HAS_TRITON and q.is_cuda and head_dim in (64, 128, 256) else "torch" + if backend == "triton": + return virtual_scores_triton(q1[0], q2[0], keys[0, :, start:end], keys[0, :, :sink]) + k_sink = keys[:, :, :sink].unsqueeze(2).transpose(-2, -1) # [1, Hkv, 1, d, sink] + k_chunk = keys[:, :, start:end].unsqueeze(2).transpose(-2, -1) # [1, Hkv, 1, d, m] + q1 = q1 * (1.0 / math.sqrt(head_dim)) + q2 = q2 * (1.0 / math.sqrt(head_dim)) + scorer = _get_scorer(backend == "torch") + scores = torch.zeros(num_heads_kv, m, device=q.device) + cols = torch.arange(m, device=q.device)[None, :] + for j0 in range(0, m, QUERY_BLOCK): + j1 = min(j0 + QUERY_BLOCK, m) + causal = cols > torch.arange(j0, j1, device=q.device)[:, None] + s_sink = torch.matmul(q2[:, :, :, j0:j1], k_sink) if sink > 0 else None + s_cached = torch.matmul(q2[:, :, :, j0:j1], k_chunk) + s_copy = torch.matmul(q1[:, :, :, j0:j1], k_chunk) + scores = torch.maximum(scores, scorer(s_sink, s_cached, s_copy, causal)) + return scores + + @staticmethod + def _quantile_map(source: torch.Tensor, target: torch.Tensor, values: torch.Tensor) -> torch.Tensor: + """Quantile-map each row of `values` from `source` to `target`.""" + vs, ts = source.sort(-1).values, target.sort(-1).values + m0 = vs.shape[-1] + x = torch.minimum(torch.maximum(values, vs[:, :1]), vs[:, -1:]) + idx = torch.searchsorted(vs, x).clamp(1, m0 - 1) + x0, x1 = vs.gather(1, idx - 1), vs.gather(1, idx) + y0, y1 = ts.gather(1, idx - 1), ts.gather(1, idx) + w = ((x - x0) / (x1 - x0).clamp_min(1e-12)).clamp(0, 1) + return y0 + w * (y1 - y0) + + def _apply_per_sample_calibration(self, start: int, end: int): + """Fit on paired first-chunk scores and apply the requested map direction.""" + exact = self.score_val[:, 0, :, start:end].float() + virtual = torch.stack([self._v0[layer_idx] for layer_idx in range(exact.shape[0])]) + n_layers, n_kv, _ = exact.shape + rows = {"global": 1, "layer": n_layers, "head": n_layers * n_kv}[self.calibration_scope] + source, target = (exact, virtual) if self.calibration_direction == "exact_to_virtual" else (virtual, exact) + if self.calibration_direction == "exact_to_virtual": + values = exact + value_slice = slice(start, end) + else: + value_start = self._chunks[1][0] + values = self.score_val[:, 0, :, value_start:].float() + value_slice = slice(value_start, None) + mapped = self._quantile_map(source.reshape(rows, -1), target.reshape(rows, -1), values.reshape(rows, -1)) + self.score_val[:, 0, :, value_slice] = mapped.view_as(values).to(self.score_val.dtype) + + def _finalize(self, model: PreTrainedModel): + """Apply exact first-chunk scores, calibration, and boundary guards.""" + start, end, _ = self._chunks[0] + if self.exact_first_chunk: + _, repeat_ids = self._pairs[0] + self.start_idx, self.end_idx = start, end + hooks = [ + layer.self_attn.register_forward_hook(self.forward_hook, with_kwargs=True) + for layer in model.model.layers + ] + try: + model(input_ids=repeat_ids.to(model.device), past_key_values=self._cache, logits_to_keep=1) + finally: + for hook in hooks: + hook.remove() + if self.exact_first_chunk and self.score_calibration and len(self._chunks) > 1: + self._apply_per_sample_calibration(start, end) + if self.guard_tokens > 0: + for start, _, _ in self._chunks: + self.score_val[..., start : start + self.guard_tokens] = 1.0 diff --git a/kvpress/presses/echo_press_kernels.py b/kvpress/presses/echo_press_kernels.py new file mode 100644 index 00000000..3f8a5495 --- /dev/null +++ b/kvpress/presses/echo_press_kernels.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Triton kernels for EchoPress virtual scoring. + +The first pass computes log-sum-exp over sink, cached, and causal-copy keys. +The second recomputes cached-key logits and takes the maximum probability per key. +""" +import math + +import torch + +try: + import triton + import triton.language as tl + + HAS_TRITON = True +except ImportError: # pragma: no cover + HAS_TRITON = False + + +if HAS_TRITON: + + @triton.jit + def _lse_kernel( + Q1, + Q2, + K, + KS, + LSE, + m, + sink, + scale, + stride_qh, + stride_qg, + stride_qm, + stride_kh, + stride_km, + stride_ksh, + stride_ksm, + stride_lh, + stride_lg, + HAS_SINK: tl.constexpr, + BQ: tl.constexpr, + BK: tl.constexpr, + BS: tl.constexpr, + D: tl.constexpr, + ): + pid_q = tl.program_id(0) + h = tl.program_id(1) + g = tl.program_id(2) + offs_q = pid_q * BQ + tl.arange(0, BQ) + offs_d = tl.arange(0, D) + q_mask = offs_q < m + q_ptrs = h * stride_qh + g * stride_qg + offs_q[:, None] * stride_qm + offs_d[None, :] + q1 = tl.load(Q1 + q_ptrs, mask=q_mask[:, None], other=0.0) + q2 = tl.load(Q2 + q_ptrs, mask=q_mask[:, None], other=0.0) + m_i = tl.full([BQ], float("-inf"), tl.float32) + l_i = tl.zeros([BQ], tl.float32) + if HAS_SINK: + offs_s = tl.arange(0, BS) + s_mask = offs_s < sink + ks = tl.load( + KS + h * stride_ksh + offs_s[:, None] * stride_ksm + offs_d[None, :], mask=s_mask[:, None], other=0.0 + ) + s = tl.dot(q2, tl.trans(ks)) * scale + s = tl.where(s_mask[None, :], s, float("-inf")) + m_new = tl.maximum(m_i, tl.max(s, 1)) + l_i = l_i * tl.exp(m_i - m_new) + tl.sum(tl.exp(s - m_new[:, None]), 1) + m_i = m_new + for k0 in range(0, m, BK): + offs_k = k0 + tl.arange(0, BK) + k_mask = offs_k < m + k = tl.load( + K + h * stride_kh + offs_k[:, None] * stride_km + offs_d[None, :], mask=k_mask[:, None], other=0.0 + ) + kt = tl.trans(k) + s2 = tl.dot(q2, kt) * scale + s2 = tl.where(k_mask[None, :], s2, float("-inf")) + s1 = tl.dot(q1, kt) * scale + s1 = tl.where(k_mask[None, :] & (offs_k[None, :] <= offs_q[:, None]), s1, float("-inf")) + m_new = tl.maximum(m_i, tl.maximum(tl.max(s2, 1), tl.max(s1, 1))) + alpha = tl.exp(m_i - m_new) + l_i = l_i * alpha + tl.sum(tl.exp(s2 - m_new[:, None]), 1) + tl.sum(tl.exp(s1 - m_new[:, None]), 1) + m_i = m_new + lse = m_i + tl.log(l_i) + tl.store(LSE + h * stride_lh + g * stride_lg + offs_q, lse, mask=q_mask) + + @triton.jit + def _colmax_kernel( + Q2, + K, + LSE, + OUT, + m, + scale, + stride_qh, + stride_qg, + stride_qm, + stride_kh, + stride_km, + stride_lh, + stride_lg, + stride_oh, + BQ: tl.constexpr, + BK: tl.constexpr, + D: tl.constexpr, + ): + pid_k = tl.program_id(0) + h = tl.program_id(1) + g = tl.program_id(2) + offs_k = pid_k * BK + tl.arange(0, BK) + offs_d = tl.arange(0, D) + k_mask = offs_k < m + k = tl.load(K + h * stride_kh + offs_k[:, None] * stride_km + offs_d[None, :], mask=k_mask[:, None], other=0.0) + kt = tl.trans(k) + acc = tl.full([BK], float("-inf"), tl.float32) + for q0 in range(0, m, BQ): + offs_q = q0 + tl.arange(0, BQ) + q_mask = offs_q < m + q2 = tl.load( + Q2 + h * stride_qh + g * stride_qg + offs_q[:, None] * stride_qm + offs_d[None, :], + mask=q_mask[:, None], + other=0.0, + ) + lse = tl.load(LSE + h * stride_lh + g * stride_lg + offs_q, mask=q_mask, other=float("inf")) + v = tl.dot(q2, kt) * scale - lse[:, None] + v = tl.where(q_mask[:, None], v, float("-inf")) + acc = tl.maximum(acc, tl.max(v, 0)) + tl.atomic_max(OUT + h * stride_oh + offs_k, acc, mask=k_mask) + + +def virtual_scores_triton( + q1: torch.Tensor, + q2: torch.Tensor, + k_chunk: torch.Tensor, + k_sink: torch.Tensor, + block_q: int = 64, + block_k: int = 64, +) -> torch.Tensor: + """Return fp32 virtual scores with shape `[Hkv, m]`.""" + assert HAS_TRITON, "triton is not available" + Hkv, G, m, d = q1.shape + sink = k_sink.shape[1] + assert d in (64, 128, 256), "head_dim must be 64, 128 or 256" + assert q1.stride(-1) == 1 and q2.stride(-1) == 1 and k_chunk.stride(-1) == 1 + scale = 1.0 / math.sqrt(d) + lse = torch.empty(Hkv, G, m, dtype=torch.float32, device=q1.device) + out = torch.full((Hkv, m), float("-inf"), dtype=torch.float32, device=q1.device) + bs = max(16, triton.next_power_of_2(sink)) if sink > 0 else 16 + grid = (triton.cdiv(m, block_q), Hkv, G) + _lse_kernel[grid]( + q1, + q2, + k_chunk, + k_sink if sink > 0 else k_chunk, + lse, + m, + sink, + scale, + q1.stride(0), + q1.stride(1), + q1.stride(2), + k_chunk.stride(0), + k_chunk.stride(1), + k_sink.stride(0) if sink > 0 else 0, + k_sink.stride(1) if sink > 0 else 0, + lse.stride(0), + lse.stride(1), + HAS_SINK=sink > 0, + BQ=block_q, + BK=block_k, + BS=bs, + D=d, + num_warps=4, + ) + grid = (triton.cdiv(m, block_k), Hkv, G) + _colmax_kernel[grid]( + q2, + k_chunk, + lse, + out, + m, + scale, + q2.stride(0), + q2.stride(1), + q2.stride(2), + k_chunk.stride(0), + k_chunk.stride(1), + lse.stride(0), + lse.stride(1), + out.stride(0), + BQ=block_q, + BK=block_k, + D=d, + num_warps=4, + ) + return out.exp() diff --git a/tests/default_presses.py b/tests/default_presses.py index 98cc6409..f466e356 100644 --- a/tests/default_presses.py +++ b/tests/default_presses.py @@ -10,6 +10,7 @@ CURPress, DropKVPress, DuoAttentionPress, + EchoPress, ExpectedAttentionPress, ExpectedAttentionStatsPress, FastKVzipPress, @@ -160,6 +161,22 @@ def post_init_from_model(self, model): {"compression_ratio": 0.8, "chunk_size": 64, "layerwise": True}, ], }, + { + "cls": EchoPress, + "kwargs": [ + {"compression_ratio": 0.5, "chunk_size": 64}, + {"compression_ratio": 0.8, "layerwise": True, "guard_tokens": 2}, + {"compression_ratio": 0.5, "chunk_size": 64, "score_calibration": False}, + {"compression_ratio": 0.5, "chunk_size": 64, "calibration_scope": "head"}, + { + "compression_ratio": 0.5, + "chunk_size": 64, + "calibration_scope": "head", + "calibration_direction": "virtual_to_exact", + }, + {"compression_ratio": 0.5, "chunk_size": 64, "scoring_backend": "torch"}, + ], + }, {"cls": TestRestoreKVPress, "kwargs": [{"compression_ratio": 0.2}, {"compression_ratio": 0.8}]}, {"cls": TestFastKVzipPress, "kwargs": [{"compression_ratio": 0.2}, {"compression_ratio": 0.8}]}, {"cls": CURPress, "kwargs": [{"compression_ratio": 0.2}, {"compression_ratio": 0.8}]}, diff --git a/tests/presses/test_echo_press.py b/tests/presses/test_echo_press.py new file mode 100644 index 00000000..044d071d --- /dev/null +++ b/tests/presses/test_echo_press.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from transformers import DynamicCache + +from kvpress import EchoPress +from tests.fixtures import unit_test_model # noqa: F401 + + +def test_default_is_per_request_per_head_virtual_to_exact(): + press = EchoPress(compression_ratio=0.5) + assert press.exact_first_chunk + assert press.score_calibration + assert press.calibration_direction == "virtual_to_exact" + assert press.calibration_scope == "head" + + +@pytest.mark.parametrize( + ("direction", "expected"), + [ + ("exact_to_virtual", [1.0, 2.0, 3.0, 1.5, 2.5, 3.0]), + ("virtual_to_exact", [10.0, 20.0, 30.0, 15.0, 25.0, 30.0]), + ], +) +def test_per_sample_calibration_direction(direction, expected): + press = EchoPress(calibration_direction=direction) + press.score_val = torch.tensor([[[[10.0, 20.0, 30.0, 1.5, 2.5, 3.0]]]]) + press._v0 = {0: torch.tensor([[1.0, 2.0, 3.0]])} + press._chunks = [(0, 3, 0), (3, 6, 0)] + + press._apply_per_sample_calibration(0, 3) + + torch.testing.assert_close(press.score_val[0, 0, 0], torch.tensor(expected)) + + +def test_online_calibration_separates_heads_and_clips_endpoints(): + press = EchoPress() + scales = torch.tensor([[10.0, 100.0], [1000.0, 10000.0]])[..., None] + virtual = torch.tensor([3.0, 1.0, 2.0]).expand(2, 2, 3) + exact = virtual * scales + later = torch.tensor([0.0, 1.5, 4.0]).expand(2, 2, 3) + press.score_val = torch.cat((exact, later), dim=-1).unsqueeze(1) + press._v0 = dict(enumerate(virtual)) + press._chunks = [(0, 3, 0), (3, 6, 0)] + + press._apply_per_sample_calibration(0, 3) + + torch.testing.assert_close(press.score_val[:, 0, :, :3], exact, rtol=0, atol=0) + expected = torch.tensor([1.0, 1.5, 3.0]) * scales + torch.testing.assert_close(press.score_val[:, 0, :, 3:], expected) + + +@pytest.mark.parametrize( + ("exact_first_chunk", "score_calibration", "chunk_size"), + [(True, True, 64), (True, False, 64), (False, True, 64), (True, True, 512)], +) +def test_online_scoring_across_requests( + unit_test_model, monkeypatch, exact_first_chunk, score_calibration, chunk_size # noqa: F811 +): + model = unit_test_model + press = EchoPress( + compression_ratio=0.5, + chunk_size=chunk_size, + exact_first_chunk=exact_first_chunk, + score_calibration=score_calibration, + ) + calibration_anchors = [] + calibrate = press._apply_per_sample_calibration + + def capture_calibration(start, end): + exact = press.score_val[..., start:end].clone() + calibration_anchors.append(press._v0[0].clone()) + calibrate(start, end) + torch.testing.assert_close(press.score_val[..., start:end], exact, rtol=0, atol=0) + assert torch.isfinite(press.score_val).all() + + monkeypatch.setattr(press, "_apply_per_sample_calibration", capture_calibration) + generator = torch.Generator(device=model.device).manual_seed(42) + for request_idx, n_tokens in enumerate((128, 300), start=1): + ids = torch.randint(3, model.config.vocab_size, (1, n_tokens), device=model.device, generator=generator) + forward_lengths = [] + + def record_forward(module, args, kwargs): + forward_lengths.append(kwargs["input_ids"].shape[1]) + + hook = model.register_forward_pre_hook(record_forward, with_kwargs=True) + cache = DynamicCache() + try: + with torch.no_grad(), press(model): + model(input_ids=ids, past_key_values=cache, logits_to_keep=1) + finally: + hook.remove() + + assert forward_lengths[0] == n_tokens + assert len(forward_lengths) == 1 + int(exact_first_chunk) + if exact_first_chunk: + assert forward_lengths[1] <= chunk_size + should_calibrate = exact_first_chunk and score_calibration and chunk_size == 64 + assert len(calibration_anchors) == request_idx * int(should_calibrate) + n_masked = sum(len(layer.self_attn.masked_key_indices[0]) for layer in model.model.layers) + n_total = model.config.num_hidden_layers * model.config.num_key_value_heads * n_tokens + assert n_masked == int(n_total * press.compression_ratio) + assert cache.get_seq_length() == n_tokens # KVzip-style eviction uses attention masks. + assert not press._q and not press._v0 + + if calibration_anchors: + assert not torch.equal(calibration_anchors[0], calibration_anchors[1])