run backbone model only for prefill

Author: giulio98

kvpress/pipeline.py

@@ -195,11 +195,11 @@ def _forward(
             cache = DynamicCache()
 
         with press(self.model) if press is not None else contextlib.nullcontext():
-            self.model(
+            # We run the model without the lm head for pre-filling.
+            self.model.model(
                 input_ids=context_ids,
                 past_key_values=cache,
                 output_attentions=self.output_attentions(press),
-                num_logits_to_keep=1,
             )
 
         logger.debug(f"Context Length: {context_length}")