From 7520bdc1aa19be6f1db3659d8c9a5dfd8a301651 Mon Sep 17 00:00:00 2001 From: Max Jeblick Date: Fri, 7 Nov 2025 14:58:20 +0100 Subject: [PATCH] wip on finding the bug Signed-off-by: Max Jeblick --- kvpress/presses/kvzip_press.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/kvpress/presses/kvzip_press.py b/kvpress/presses/kvzip_press.py index 0f13e68a..f0ba82c9 100644 --- a/kvpress/presses/kvzip_press.py +++ b/kvpress/presses/kvzip_press.py @@ -116,10 +116,7 @@ def __call__(self, model: PreTrainedModel) -> Generator: def wrapped_forward(model_self, *args, **kwargs): self._context_ids = kwargs["input_ids"] - assert ( - "past_key_value" in kwargs or "past_key_values" in kwargs - ), f"KVzipPress requires 'past_key_value' or 'past_key_values' during prefilling. Got {kwargs.keys()}" - self._cache = kwargs.get("past_key_values", None) or kwargs.get("past_key_value", None) + self._cache = kwargs["past_key_values"] return original_forward(*args, **kwargs) model.model.forward = MethodType(wrapped_forward, model.model)