Fix NonCausalAttnPress key-padding mask value 🤖🤖🤖

kvpress/presses/non_causal_attention_press.py

@@ -88,7 +88,7 @@ def non_causal_chunked_attn(q: torch.Tensor, k: torch.Tensor, chunk_size: int) -
         # (B, H, num_chunks, chunk_size, chunk_size)
         dots = torch.matmul(q_chunks, k_chunks.transpose(-2, -1))
         dots[:, :, -1].masked_fill_(query_mask.unsqueeze(-1), 0)
-        dots[:, :, -1].masked_fill_(key_mask.unsqueeze(-2), -1e-9)
+        dots[:, :, -1].masked_fill_(key_mask.unsqueeze(-2), torch.finfo(dots.dtype).min)
         attn = torch.softmax(dots.to(torch.float32), dim=-1)
         # sum over query and trim padding
         return attn.sum(dim=-2).view(B, H, S_pad)[..., :S]

tests/presses/test_compactor_press.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import pytest
 import torch
 
 from kvpress import CompactorPress, LeverageScorePress, NonCausalAttnPress
@@ -35,3 +36,21 @@ def test_non_causal_attn_press(unit_test_model):  # noqa: F811
         with press(unit_test_model):
             input_ids = torch.arange(10, 40).to(unit_test_model.device)
             unit_test_model(input_ids.unsqueeze(0), use_cache=True)
+
+
+def test_non_causal_chunked_attn_does_not_leak_mass_to_padding_keys():
+    # S=2 with chunk_size=4 puts the whole (single) chunk in the ragged case: 2 real
+    # query/key positions plus 2 padded ones. Every one of the 4 query rows in that
+    # chunk (real or padded) puts its full softmax unit of mass on the 2 real key
+    # columns when padding is masked out correctly, so the returned per-key scores
+    # (summed over all 4 queries) must add up to ~chunk_size regardless of the actual
+    # q/k values. A masking value too small to dominate the logits lets softmax leak
+    # part of that mass onto the phantom zero-vector padding keys instead, which are
+    # then dropped, so the real keys' scores come back under-counted.
+    torch.manual_seed(0)
+    B, H, S, d, chunk_size = 1, 1, 2, 4, 4
+    q = torch.randn(B, H, S, d)
+    k = torch.randn(B, H, S, d)
+
+    scores = NonCausalAttnPress.non_causal_chunked_attn(q, k, chunk_size)
+    assert scores.sum().item() == pytest.approx(chunk_size, abs=1e-3)