From 00d588708a126d81919d473d94f21c95c5acf54f Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 4 Nov 2026 14:54:46 +0100 Subject: [PATCH] =?UTF-8?q?Keep=20SnapKV=20window=20scores=20on=20device?= =?UTF-8?q?=20=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: 0z5a --- kvpress/presses/snapkv_press.py | 5 ++- tests/presses/test_snapkv_window.py | 60 +++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 2 deletions(-) create mode 100644 tests/presses/test_snapkv_window.py diff --git a/kvpress/presses/snapkv_press.py b/kvpress/presses/snapkv_press.py index 61cd85055..f26e64750 100644 --- a/kvpress/presses/snapkv_press.py +++ b/kvpress/presses/snapkv_press.py @@ -97,7 +97,8 @@ def score( scores = scores.view(bsz, num_key_value_heads, num_key_value_groups, k_len - self.window_size) scores = scores.mean(2) - # Add back the observation window. Use max score to make sure the window is not pruned. - scores = F.pad(scores, (0, self.window_size), value=scores.max().item() + 1) + # Keep the observation window without copying its sentinel to the host. + window_scores = (scores.max() + 1).expand(*scores.shape[:-1], self.window_size) + scores = torch.cat((scores, window_scores), dim=-1) return scores diff --git a/tests/presses/test_snapkv_window.py b/tests/presses/test_snapkv_window.py new file mode 100644 index 000000000..2780f3091 --- /dev/null +++ b/tests/presses/test_snapkv_window.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch + +from kvpress.presses.snapkv_press import SnapKVPress + + +def make_inputs(device, dtype): + # Constant attention per query head makes the expected GQA averages exact. + batch, heads, kv_heads, length, window, head_dim = 2, 4, 2, 16, 4, 8 + head_values = torch.tensor( + [[1 / 8, 2 / 8, 3 / 8, 4 / 8], [2 / 8, 4 / 8, 6 / 8, 8 / 8]], + device=device, + dtype=dtype, + ) + attention = head_values[:, :, None, None].expand(batch, heads, window, length).contiguous() + module = SimpleNamespace(config=SimpleNamespace(num_attention_heads=heads), head_dim=head_dim) + hidden = torch.zeros(batch, length, head_dim, device=device, dtype=dtype) + keys = torch.zeros(batch, kv_heads, length, head_dim, device=device, dtype=dtype) + return module, hidden, keys, attention + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable"), + ), + ], +) +def test_window_scores_preserve_global_sentinel_and_head_averages(device, dtype): + module, hidden, keys, attention = make_inputs(device, dtype) + press = SnapKVPress(compression_ratio=0.5, window_size=4, kernel_size=1) + + scores = press.score(module, hidden, keys, keys, attention, {}) + + group_values = torch.tensor([[3 / 16, 7 / 16], [3 / 8, 7 / 8]], device=device, dtype=dtype) + expected = group_values[:, :, None].expand(2, 2, 16).clone() + expected[:, :, -4:] = 15 / 8 + torch.testing.assert_close(scores, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +def test_scoring_does_not_synchronize_cuda_to_read_window_sentinel(): + module, hidden, keys, attention = make_inputs("cuda", torch.float32) + press = SnapKVPress(compression_ratio=0.5, window_size=4, kernel_size=1) + previous = torch.cuda.get_sync_debug_mode() + try: + torch.cuda.set_sync_debug_mode("error") + scores = press.score(module, hidden, keys, keys, attention, {}) + finally: + torch.cuda.set_sync_debug_mode(previous) + assert scores.shape == (2, 2, 16)