Fix ChunkKV selection for batched inputs 🤖🤖🤖

kvpress/presses/chunkkv_press.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from dataclasses import dataclass
@@ -103,20 +103,20 @@ def compress(
         n_chunks_kept = max(1, int((num_complete_chunks + (remaining_tokens > 0)) * (1 - self.press.compression_ratio)))
         top_chunks = chunk_scores.topk(n_chunks_kept, dim=-1)
 
-        # 4. Create indices for selected chunks
-        indices = []
-        for chunk_idx in top_chunks.indices[0]:
-            if chunk_idx < num_complete_chunks:
-                # For complete chunks
+        # 4. Create indices for the chunks selected by each batch element.
+        batch_indices = []
+        for selected_chunks in top_chunks.indices:
+            selected_indices = []
+            for chunk_idx in selected_chunks:
                 start_idx = chunk_idx * self.chunk_length
-                chunk_indices = torch.arange(start_idx, start_idx + self.chunk_length, device=keys.device)
-            else:
-                # For the remaining partial chunk
-                chunk_indices = torch.arange(num_complete_chunks * self.chunk_length, kv_len, device=keys.device)
-            indices.append(chunk_indices)
-
-        indices = torch.cat(indices).sort()[0]
-        indices = indices.view(1, 1, -1, 1).expand(keys.shape[0], keys.shape[1], -1, module.head_dim)
+                end_idx = min(start_idx + self.chunk_length, kv_len)
+                selected_indices.append(torch.arange(start_idx, end_idx, device=keys.device))
+            batch_indices.append(torch.cat(selected_indices).sort()[0])
+
+        if any(indices.numel() != batch_indices[0].numel() for indices in batch_indices):
+            raise ValueError("ChunkKVPress selected unequal cache lengths across the batch")
+        indices = torch.stack(batch_indices).view(keys.shape[0], 1, -1, 1)
+        indices = indices.expand(-1, keys.shape[1], -1, module.head_dim)
 
         # 5. Use gather to collect selected keys and values
         keys = keys.gather(2, indices).contiguous()

tests/presses/test_chunkkv_batch.py

@@ -0,0 +1,50 @@
+# 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 import ChunkKVPress
+from kvpress.presses.scorer_press import ScorerPress
+
+
+class FixedScorer(ScorerPress):
+    def score(self, module, hidden_states, keys, values, attentions, kwargs):
+        return kwargs["scores"]
+
+
+def compress(scores, chunk_length=2):
+    batch, _, length = scores.shape
+    keys = torch.arange(batch * length, dtype=torch.float32).reshape(batch, 1, length, 1)
+    values = keys + 100
+    press = ChunkKVPress(FixedScorer(compression_ratio=0.5), chunk_length=chunk_length)
+    return press.compress(
+        SimpleNamespace(head_dim=1), torch.zeros(batch, length, 1), keys, values, None, {"scores": scores}
+    )
+
+
+def test_chunk_selection_is_independent_per_batch_element():
+    scores = torch.tensor([[[9.0, 9.0, 1.0, 1.0]], [[1.0, 1.0, 9.0, 9.0]]])
+    keys, values = compress(scores)
+    torch.testing.assert_close(keys[:, 0, :, 0], torch.tensor([[0.0, 1.0], [6.0, 7.0]]))
+    torch.testing.assert_close(values, keys + 100)
+
+
+def test_partial_chunks_with_equal_lengths_are_selected_per_sample():
+    scores = torch.tensor([[[9.0, 9.0, 1.0, 1.0, 0.0]], [[1.0, 1.0, 9.0, 9.0, 0.0]]])
+    keys, _ = compress(scores)
+    torch.testing.assert_close(keys[:, 0, :, 0], torch.tensor([[0.0, 1.0], [7.0, 8.0]]))
+
+
+def test_partial_chunks_with_unequal_lengths_raise():
+    scores = torch.tensor([[[9.0, 9.0, 1.0, 1.0, 0.0]], [[1.0, 1.0, 0.0, 0.0, 9.0]]])
+    with pytest.raises(ValueError, match="unequal cache lengths"):
+        compress(scores)
+
+
+def test_single_sample_can_select_partial_chunk():
+    scores = torch.tensor([[[1.0, 1.0, 0.0, 0.0, 9.0]]])
+    keys, _ = compress(scores)
+    torch.testing.assert_close(keys[:, 0, :, 0], torch.tensor([[4.0]]))