From c15aa5d6c425b486d2522259c8e3e842ac81f9f8 Mon Sep 17 00:00:00 2001 From: Nate Date: Mon, 12 Oct 2026 17:46:42 -0400 Subject: [PATCH] Fix causal attention masks for uneven layer caches Signed-off-by: Nate --- kvpress/attention_patch.py | 71 ++++++++++-- tests/test_attention_patch.py | 204 +++++++++++++++++++++++++++++++++- tests/test_uneven_cache.py | 109 ++++++++++++++++++ 3 files changed, 375 insertions(+), 9 deletions(-) create mode 100644 tests/test_uneven_cache.py diff --git a/kvpress/attention_patch.py b/kvpress/attention_patch.py index 270ad9e6..07be86b0 100644 --- a/kvpress/attention_patch.py +++ b/kvpress/attention_patch.py @@ -1,9 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import importlib +from functools import wraps + import torch from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +from kvpress.presses.base_press import SUPPORTED_MODELS + def search_hyperplane(X, max_iter: int = 1000): """ @@ -48,6 +53,33 @@ def search_hyperplane(X, max_iter: int = 1000): raise ValueError("Could not find fake keys such that for every query q, exp() = 0") +def _pad_causal_attention_mask(attention_mask: torch.Tensor, query_length: int, padding_length: int) -> torch.Tensor: + """Extend only an unpadded causal mask, whose additional past columns are known to be visible.""" + error = ( + "Cannot extend this attention mask for a longer layer cache. " + "Only unpadded causal masks can be extended; provide a layer-specific attention mask." + ) + mask_length = attention_mask.shape[-1] + if attention_mask.shape[-2] != query_length or mask_length < query_length: + raise ValueError(error) + + key_positions = torch.arange(mask_length, device=attention_mask.device) + query_positions = torch.arange(query_length, device=attention_mask.device).unsqueeze(-1) + allowed = key_positions <= query_positions + mask_length - query_length + if attention_mask.dtype == torch.bool: + valid = attention_mask == allowed + elif attention_mask.dtype.is_floating_point: + blocked = torch.isneginf(attention_mask) | (attention_mask == torch.finfo(attention_mask.dtype).min) + valid = torch.where(allowed, attention_mask == 0, blocked) + else: + raise ValueError(error) + if not valid.all(): + raise ValueError(error) + + visible = True if attention_mask.dtype == torch.bool else 0.0 + return torch.nn.functional.pad(attention_mask, (padding_length, 0), value=visible) + + def attention_patch(func): """ Decorator to update the keys before the attention computation at the indices provided in module.masked_key_indices @@ -66,7 +98,8 @@ def attention_patch(func): The wrapped attention function that supports head-wise key masking. """ - def wrapper(module, query, key, value, attention_mask, dropout, **kwargs): + @wraps(func) + def wrapper(module, query, key, value, attention_mask, dropout=0.0, **kwargs): if query.shape[2] == key.shape[2]: # Prefilling module.masked_key_indices = None @@ -86,23 +119,38 @@ def wrapper(module, query, key, value, attention_mask, dropout, **kwargs): batch_indices, head_indices, seq_indices = module.masked_key_indices key[batch_indices, head_indices, seq_indices] = k[batch_indices, head_indices] + if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 4 and attention_mask.shape[-1] != 1: + # The shared mask uses the first layer's cache length. New query tokens + # are at the right edge of every layer, regardless of its context budget. + difference = key.shape[-2] - attention_mask.shape[-1] + if difference < 0: + attention_mask = attention_mask[..., -key.shape[-2] :] + elif difference > 0: + if kwargs.get("sliding_window") is not None or getattr(module, "sliding_window", None) is not None: + raise ValueError( + "Cannot extend a sliding-window attention mask for a longer layer cache; " + "provide a layer-specific attention mask." + ) + attention_mask = _pad_causal_attention_mask(attention_mask, query.shape[-2], difference) + # see https://github.com/NVIDIA/kvpress/pull/115#issuecomment-3183785597 # cu_seq_lens_k are only in kwargs if model.generate is used. if "cu_seq_lens_k" in kwargs: kwargs["cu_seq_lens_k"][-1] = key.shape[-2] - return func(module, query, key, value, attention_mask, dropout, **kwargs) + return func(module, query, key, value, attention_mask, dropout=dropout, **kwargs) return wrapper def patch_attention_functions(): """ - Apply attention patching to all transformer attention functions. + Apply attention patching to registered backends and supported models' eager functions. This function automatically patches all attention functions registered in transformers' ALL_ATTENTION_FUNCTIONS to support head-wise key masking. - It enables KVPress compression methods that require head-specific masking - (like AdaKV) to work correctly during text generation. + It also patches the model-specific eager functions, which are not registered there. + This enables head-specific masking (like AdaKV) and unequal per-layer caches + (like PyramidKV) to work correctly during text generation. The patching is applied globally and affects all transformer models loaded after this function is called. It's automatically called when importing @@ -110,9 +158,16 @@ def patch_attention_functions(): Notes ----- - This function modifies the global attention functions in the transformers - library. The modifications do not affect models that don't use head-wise compression (i.e. don't have - module.masked_key_indices). + This function modifies the global attention functions in the transformers library. + Attention masks with the correct key length, singleton key dimensions, and + two-dimensional padding masks are left unchanged. Four-dimensional masks are + right-aligned to each layer's keys. Longer caches can extend unpadded causal masks + with visible past-context columns. Restricted masks, including sliding-window + masks, require a layer-specific mask when the layer cache is longer. """ for name, func in ALL_ATTENTION_FUNCTIONS.items(): ALL_ATTENTION_FUNCTIONS[name] = attention_patch(func) + for model in SUPPORTED_MODELS: + modeling_module = importlib.import_module(model.__module__) + eager_attention = getattr(modeling_module, "eager_attention_forward") + setattr(modeling_module, "eager_attention_forward", attention_patch(eager_attention)) diff --git a/tests/test_attention_patch.py b/tests/test_attention_patch.py index 1e3dccdb..7f3696d9 100644 --- a/tests/test_attention_patch.py +++ b/tests/test_attention_patch.py @@ -1,10 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import importlib + import pytest import torch +from torch import nn -from kvpress.attention_patch import search_hyperplane +from kvpress.attention_patch import attention_patch, search_hyperplane +from kvpress.presses.base_press import SUPPORTED_MODELS def attention_weights(X, K, attention_scaling): @@ -50,3 +54,201 @@ def test_search_hyperplane_raises_for_non_separable_queries(): X = torch.cat([X, -X], dim=1) # no hyperplane has both q and -q on its positive side with pytest.raises(ValueError, match="Could not find fake keys"): search_hyperplane(X, max_iter=10) + + +@pytest.mark.parametrize("key_length", [4, 7, 10, 13]) +@pytest.mark.parametrize("query_length", [1, 3]) +@pytest.mark.parametrize("dtype", [torch.bool, torch.float32]) +def test_attention_patch_aligns_causal_mask_to_layer_cache(key_length, query_length, dtype): + """Unequal layer caches must retain the causal mask for the new query tokens.""" + torch.manual_seed(0) + query = torch.randn(1, 1, query_length, 4) + key = torch.randn(1, 1, key_length, 4) + value = torch.randn_like(key) + mask_length = 10 + allowed = torch.arange(mask_length) <= torch.arange(query_length).unsqueeze(-1) + mask_length - query_length + mask = allowed.unsqueeze(0).unsqueeze(0) + if dtype != torch.bool: + mask = torch.zeros_like(mask, dtype=dtype).masked_fill(~mask, torch.finfo(dtype).min) + original_mask = mask.clone() + + def attention(module, query, key, value, attention_mask, dropout): + return torch.nn.functional.scaled_dot_product_attention(query, key, value, attention_mask, dropout) + + actual = attention_patch(attention)(nn.Module(), query, key, value, mask, 0.0) + expected_allowed = torch.arange(key_length) <= torch.arange(query_length).unsqueeze(-1) + key_length - query_length + expected_mask = expected_allowed.unsqueeze(0).unsqueeze(0) + if dtype != torch.bool: + expected_mask = torch.zeros_like(expected_mask, dtype=dtype).masked_fill(~expected_mask, torch.finfo(dtype).min) + expected = attention(nn.Module(), query, key, value, expected_mask, 0.0) + + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(mask, original_mask) + + +@pytest.mark.parametrize("mask", [None, torch.tensor([[1, 1, 0, 0]])]) +def test_attention_patch_preserves_flash_attention_padding_mask(mask): + """Flash attention's two-dimensional padding mask is not a causal score matrix.""" + query = torch.zeros(1, 1, 1, 4) + key = torch.zeros(1, 1, 2, 4) + + def attention(module, query, key, value, attention_mask, dropout): + return attention_mask + + assert attention_patch(attention)(nn.Module(), query, key, key, mask, 0.0) is mask + + +@pytest.mark.parametrize("dtype", [torch.bool, torch.float32]) +@pytest.mark.parametrize("query_length", [1, 4]) +def test_attention_patch_preserves_broadcast_mask(dtype, query_length): + """A singleton key dimension must keep broadcasting, including fully masked queries.""" + query = torch.zeros(1, 1, query_length, 4) + key = torch.zeros(1, 1, 4, 4) + value = torch.arange(16, dtype=torch.float32).reshape_as(key) + mask = torch.ones(1, 1, query_length, 1, dtype=torch.bool) + mask[..., 0, :] = False + if dtype != torch.bool: + mask = torch.zeros_like(mask, dtype=dtype).masked_fill(~mask, -torch.inf) + + def attention(module, query, key, value, attention_mask, dropout): + assert attention_mask is mask + return torch.nn.functional.scaled_dot_product_attention(query, key, value, attention_mask, dropout) + + actual = attention_patch(attention)(nn.Module(), query, key, value, mask) + expected = attention(nn.Module(), query, key, value, mask, 0.0) + torch.testing.assert_close(actual, expected) + assert actual[..., 0, :].count_nonzero() == 0 + + +@pytest.mark.parametrize("dtype", [torch.bool, torch.float32]) +@pytest.mark.parametrize("restriction", ["padding", "sliding", "future", "query_broadcast"]) +def test_attention_patch_rejects_extending_restricted_masks(dtype, restriction): + """Unknown masking constraints must not be replaced by visible context columns.""" + query = torch.zeros(1, 1, 2, 4) + key = torch.zeros(1, 1, 8, 4) + key_positions = torch.arange(5) + query_positions = torch.arange(2).unsqueeze(-1) + 3 + mask = (key_positions <= query_positions).unsqueeze(0).unsqueeze(0) + if restriction == "padding": + mask[..., 0] = False + elif restriction == "sliding": + mask &= key_positions > query_positions - 3 + elif restriction == "future": + mask[..., 0, -1] = True + else: + mask = mask[..., -1:, :] + if dtype != torch.bool: + mask = torch.zeros_like(mask, dtype=dtype).masked_fill(~mask, torch.finfo(dtype).min) + original_mask = mask.clone() + + def attention(module, query, key, value, attention_mask, dropout): + pytest.fail("Restricted masks must be rejected before calling the backend") + + with pytest.raises(ValueError, match="layer-specific"): + attention_patch(attention)(nn.Module(), query, key, key, mask) + torch.testing.assert_close(mask, original_mask) + + +@pytest.mark.parametrize("dtype", [torch.bool, torch.float32]) +@pytest.mark.parametrize("window_source", ["module", "keyword"]) +def test_attention_patch_rejects_extending_sliding_window_masks(dtype, window_source): + """Even a window covering the entire shared mask can exclude keys in a longer layer.""" + query = torch.zeros(1, 1, 3, 4) + key = torch.zeros(1, 1, 13, 4) + allowed = torch.arange(10) <= torch.arange(3).unsqueeze(-1) + 7 + mask = allowed.unsqueeze(0).unsqueeze(0) + if dtype != torch.bool: + mask = torch.zeros_like(mask, dtype=dtype).masked_fill(~mask, torch.finfo(dtype).min) + module = nn.Module() + kwargs = {} + if window_source == "module": + module.sliding_window = 10 + else: + kwargs["sliding_window"] = 10 + + def attention(module, query, key, value, attention_mask, dropout, **kwargs): + pytest.fail("Sliding-window masks must be rejected before calling the backend") + + with pytest.raises(ValueError, match="layer-specific"): + attention_patch(attention)(module, query, key, key, mask, **kwargs) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +@pytest.mark.parametrize("use_infinity", [False, True]) +def test_attention_patch_extends_additive_causal_mask(dtype, use_infinity): + """Both finite minimum values and negative infinity represent blocked causal positions.""" + query = torch.zeros(1, 1, 2, 4) + key = torch.zeros(1, 1, 8, 4) + allowed = torch.arange(5) <= torch.arange(2).unsqueeze(-1) + 3 + blocked = -torch.inf if use_infinity else torch.finfo(dtype).min + mask = torch.zeros(1, 1, 2, 5, dtype=dtype).masked_fill(~allowed, blocked) + + def attention(module, query, key, value, attention_mask, dropout): + return attention_mask + + actual = attention_patch(attention)(nn.Module(), query, key, key, mask) + expected_allowed = torch.arange(8) <= torch.arange(2).unsqueeze(-1) + 6 + expected = torch.zeros(1, 1, 2, 8, dtype=dtype).masked_fill(~expected_allowed, blocked) + assert actual.dtype == dtype + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("dtype", [torch.bool, torch.float32]) +def test_attention_patch_preserves_layer_specific_restricted_mask(dtype): + """Restricted masks with the correct key length must remain usable without modification.""" + query = torch.zeros(1, 1, 2, 4) + key = torch.zeros(1, 1, 5, 4) + positions = torch.arange(5) + query_positions = torch.arange(2).unsqueeze(-1) + 3 + mask = ((positions <= query_positions) & (positions > query_positions - 3))[None, None] + if dtype != torch.bool: + mask = torch.zeros_like(mask, dtype=dtype).masked_fill(~mask, torch.finfo(dtype).min) + module = nn.Module() + module.sliding_window = 3 + + def attention(module, query, key, value, attention_mask, dropout): + return attention_mask + + assert attention_patch(attention)(module, query, key, key, mask) is mask + + +def test_attention_patch_rejects_extending_attention_bias(): + """Finite attention biases cannot be inferred for additional context keys.""" + query = torch.zeros(1, 1, 2, 4) + key = torch.zeros(1, 1, 8, 4) + allowed = torch.arange(5) <= torch.arange(2).unsqueeze(-1) + 3 + mask = torch.zeros(1, 1, 2, 5).masked_fill(~allowed, -torch.inf) + mask[..., 0] = 0.5 + + def attention(module, query, key, value, attention_mask, dropout): + pytest.fail("Attention biases must be rejected before calling the backend") + + with pytest.raises(ValueError, match="layer-specific"): + attention_patch(attention)(nn.Module(), query, key, key, mask) + + +@pytest.mark.parametrize("model", SUPPORTED_MODELS) +@pytest.mark.parametrize("key_length", [7, 13]) +def test_supported_models_eager_attention_preserves_model_specific_behavior(model, key_length): + """Wrap each original eager backend, including its scaling and soft-capping options.""" + function = importlib.import_module(model.__module__).eager_attention_forward + original = function.__wrapped__ + module = nn.Module().eval() + module.num_key_value_groups = 2 + module.head_dim = 4 + torch.manual_seed(0) + query = torch.randn(1, 2, 3, 4) + key = torch.randn(1, 1, key_length, 4) + value = torch.randn_like(key) + allowed = torch.arange(10) <= torch.arange(3).unsqueeze(-1) + 7 + mask = torch.zeros(1, 1, 3, 10).masked_fill(~allowed, torch.finfo(torch.float32).min) + original_mask = mask.clone() + + actual_output, actual_weights = function(module, query, key, value, mask, scaling=0.25, softcap=1.0) + expected_allowed = torch.arange(key_length) <= torch.arange(3).unsqueeze(-1) + key_length - 3 + expected_mask = torch.zeros(1, 1, 3, key_length).masked_fill(~expected_allowed, torch.finfo(torch.float32).min) + expected_output, expected_weights = original(module, query, key, value, expected_mask, scaling=0.25, softcap=1.0) + + torch.testing.assert_close(actual_output, expected_output) + torch.testing.assert_close(actual_weights, expected_weights) + torch.testing.assert_close(mask, original_mask) diff --git a/tests/test_uneven_cache.py b/tests/test_uneven_cache.py new file mode 100644 index 00000000..bd1118c2 --- /dev/null +++ b/tests/test_uneven_cache.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from copy import deepcopy + +import pytest +import torch +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from transformers import DynamicCache, LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast + +from kvpress import KnormPress, PerLayerCompressionPress, PyramidKVPress, SnapKVPress +from kvpress.pipeline import KVPressTextGenerationPipeline + + +@pytest.fixture(params=["sdpa", "eager"]) +def small_model(request): + """Build a CPU model locally without downloading a checkpoint or tokenizer.""" + torch.manual_seed(0) + config = LlamaConfig( + vocab_size=32, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=2, + attn_implementation=request.param, + eos_token_id=31, + pad_token_id=0, + ) + return LlamaForCausalLM(config).eval() + + +@pytest.fixture(params=["pyramid", "per_layer", "increasing", "mixed", "uniform"]) +def press(request): + """Exercise unequal layer budgets in both directions and a uniform-cache control.""" + if request.param == "pyramid": + return PyramidKVPress(compression_ratio=0.5, window_size=4, beta=2) + if request.param == "per_layer": + return PerLayerCompressionPress(press=KnormPress(), compression_ratios=[0.25, 0.5, 0.75]) + if request.param == "increasing": + return PerLayerCompressionPress(press=KnormPress(), compression_ratios=[0.75, 0.5, 0.25]) + if request.param == "mixed": + return PerLayerCompressionPress(press=KnormPress(), compression_ratios=[0.5, 0.25, 0.75]) + return SnapKVPress(compression_ratio=0.5, window_size=4) + + +def context_ids(): + """Use enough context tokens to avoid PyramidKV's uniform-budget fallback.""" + return torch.arange(128).remainder(30).add(1).unsqueeze(0) + + +@torch.no_grad() +def test_pipeline_with_uneven_cache(small_model, press): + """Question and answer generation must preserve every layer's compressed context.""" + vocabulary = {"[UNK]": 0, **{str(i): i for i in range(1, 32)}} + tokenizer_object = Tokenizer(WordLevel(vocabulary, unk_token="[UNK]")) + tokenizer_object.pre_tokenizer = Whitespace() + tokenizer = PreTrainedTokenizerFast(tokenizer_object=tokenizer_object, unk_token="[UNK]") + pipeline = KVPressTextGenerationPipeline(model=small_model, tokenizer=tokenizer, device="cpu") + cache = DynamicCache() + inputs = context_ids() + with press(small_model): + small_model.model(input_ids=inputs, past_key_values=cache) + expected_lengths = [cache.get_seq_length(i) for i in range(3)] + + cache = DynamicCache() + answers = pipeline._forward( + {"context_ids": inputs, "questions_ids": [torch.tensor([[2, 3, 4, 5]]), torch.tensor([[6, 7, 8]])]}, + press=press, + cache=cache, + max_new_tokens=2, + ) + + assert len(answers) == 2 + assert all(isinstance(answer, str) for answer in answers) + assert [cache.get_seq_length(i) for i in range(3)] == expected_lengths + if isinstance(press, (PyramidKVPress, PerLayerCompressionPress)): + assert len(set(expected_lengths)) == 3 + else: + assert len(set(expected_lengths)) == 1 + + +@torch.no_grad() +def test_generate_with_uneven_cache(small_model, press): + """The standard Transformers generation API must also accept unequal layer budgets.""" + inputs = context_ids() + with press(small_model): + output = small_model.generate(inputs, max_new_tokens=3, do_sample=False) + assert output.shape[1] > inputs.shape[1] + + +@torch.no_grad() +def test_uneven_cache_preserves_causality(small_model, press): + """A multi-token question must match token-by-token decoding from the same cache.""" + cache = DynamicCache() + inputs = context_ids() + with press(small_model): + small_model.model(input_ids=inputs, past_key_values=cache) + query = torch.tensor([[2, 3, 4, 5]]) + positions = torch.arange(inputs.shape[1], inputs.shape[1] + query.shape[1]).unsqueeze(0) + batched_logits = small_model(query, past_key_values=deepcopy(cache), position_ids=positions).logits + sequential_logits = [] + for i in range(query.shape[1]): + output = small_model(query[:, i : i + 1], past_key_values=cache, position_ids=positions[:, i : i + 1]) + sequential_logits.append(output.logits) + + torch.testing.assert_close(batched_logits, torch.cat(sequential_logits, dim=1), atol=1e-6, rtol=1e-5)