From 40237e6bb5fe37518379fedd822538f1d9b94b1a Mon Sep 17 00:00:00 2001
From: cuteNBD <206630735+bingdongni@users.noreply.github.com>
Date: Sat, 7 Nov 2026 10:28:37 +0800
Subject: [PATCH] =?UTF-8?q?Fix=20singleton=20chunks=20in=20ChunkPress=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: cuteNBD <206630735+bingdongni@users.noreply.github.com>
---
 kvpress/presses/chunk_press.py        |  5 +++
 tests/presses/test_chunk_singleton.py | 45 +++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)
 create mode 100644 tests/presses/test_chunk_singleton.py

diff --git a/kvpress/presses/chunk_press.py b/kvpress/presses/chunk_press.py
index 06609dd09..725412185 100644
--- a/kvpress/presses/chunk_press.py
+++ b/kvpress/presses/chunk_press.py
@@ -65,6 +65,11 @@ def compress(
         kv_len = keys.shape[2]
         indices = []
         for i in range(0, kv_len, self.chunk_length):
+            if min(self.chunk_length, kv_len - i) == 1:
+                # There is no selection to make when a chunk has only one token.
+                indices.append(torch.full((*keys.shape[:2], 1), i, device=keys.device, dtype=torch.long))
+                continue
+
             # Slice the rotary embeddings so that presses recomputing queries (e.g. SnapKVPress)
             # use the positions of the chunk and not the positions at the end of the full sequence
             chunk_kwargs = kwargs
diff --git a/tests/presses/test_chunk_singleton.py b/tests/presses/test_chunk_singleton.py
new file mode 100644
index 000000000..4badbac7d
--- /dev/null
+++ b/tests/presses/test_chunk_singleton.py
@@ -0,0 +1,45 @@
+# 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, LlamaConfig, LlamaForCausalLM
+
+from kvpress import ChunkPress, KnormPress, SnapKVPress, TOVAPress
+
+
+@pytest.mark.parametrize(
+    "device",
+    ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"))],
+)
+@pytest.mark.parametrize("seq_len,chunk_length", [(1, 4), (5, 4), (9, 4), (5, 1)])
+@pytest.mark.parametrize("method", ["tova", "snapkv", "knorm"])
+def test_singleton_chunks_are_kept(device, seq_len, chunk_length, method):
+    config = LlamaConfig(
+        vocab_size=64,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        max_position_embeddings=128,
+        _attn_implementation="sdpa",
+    )
+    model = LlamaForCausalLM(config).eval().to(device)
+    inputs = torch.arange(2 * seq_len, device=device).reshape(2, seq_len) % config.vocab_size
+    presses = {
+        "tova": TOVAPress(compression_ratio=0.5),
+        "snapkv": SnapKVPress(compression_ratio=0.5, window_size=1, kernel_size=3),
+        "knorm": KnormPress(compression_ratio=0.5),
+    }
+
+    with torch.inference_mode():
+        reference = model(inputs, past_key_values=DynamicCache()).past_key_values
+        with ChunkPress(press=presses[method], chunk_length=chunk_length)(model):
+            compressed = model(inputs, past_key_values=DynamicCache()).past_key_values
+
+    expected_length = sum(max(1, int(min(chunk_length, seq_len - i) * 0.5)) for i in range(0, seq_len, chunk_length))
+    assert compressed.layers[0].keys.shape[2] == expected_length
+    assert compressed.layers[0].values.shape[2] == expected_length
+    torch.testing.assert_close(compressed.layers[0].keys[:, :, -1], reference.layers[0].keys[:, :, -1])
+    torch.testing.assert_close(compressed.layers[0].values[:, :, -1], reference.layers[0].values[:, :, -1])