From d719c10ba0b19aa7a92d82821b1b3ba87e850d9f Mon Sep 17 00:00:00 2001 From: Amir Fathi Date: Mon, 28 Sep 2026 22:50:36 +0000 Subject: [PATCH] Fix NonCausalAttnPress key-padding mask value MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit non_causal_chunked_attn masked padded key columns in the ragged last chunk with -1e-9, which is negligible next to typical dot-product logits and barely changes the softmax. Real queries in that chunk leaked attention mass onto the phantom zero-vector padding keys instead of concentrating it on the real keys, which then get trimmed away, so the per-key importance score used for eviction came back systematically under-counted whenever the sequence length was not an exact multiple of chunk_size. Mask with torch.finfo(dots.dtype).min instead, matching the pattern kvzip_press.py already uses for the same kind of key masking. Fixes #295 🤖🤖🤖 Signed-off-by: Amir Fathi --- kvpress/presses/non_causal_attention_press.py | 2 +- tests/presses/test_compactor_press.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/kvpress/presses/non_causal_attention_press.py b/kvpress/presses/non_causal_attention_press.py index 81c5415a9..4f17528f5 100644 --- a/kvpress/presses/non_causal_attention_press.py +++ b/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] diff --git a/tests/presses/test_compactor_press.py b/tests/presses/test_compactor_press.py index 3b261df7a..56738b1f8 100644 --- a/tests/presses/test_compactor_press.py +++ b/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)