LLama: Optimize KV cache management for multi-round conversations

- Implements prefix-matching logic to truncate stale "ghost" tokens in C++ KV cache
- Prevents attention misalignment and context poisoning during multi-turn interactions
- Reduces memory overhead by reusing matched prefixes efficiently

Signed-off-by: JamePeng <jame_peng@sina.com>

Author: JamePeng

llama_cpp/llama.py

@@ -1077,14 +1077,28 @@ def adapter(token_data_array: llama_cpp.llama_token_data_array):
             longest_prefix = self.longest_token_prefix(self._input_ids, tokens[:-1])
             if longest_prefix > 0:
                 reset = False
+
+                # Physically erase trailing "ghost" tokens from the C++ KV cache
+                # to prevent attention misalignment in multi-round chats.
+                if longest_prefix < self.n_tokens:
+                    if self.verbose:
+                        print(f"Llama.generate: Truncating KV cache size from {self.n_tokens} to {longest_prefix}", file=sys.stderr)
+                    self._ctx.memory_seq_rm(0, longest_prefix, -1)
+
+                # Adjust the tokens array and cursor to reuse the matched cache
                 tokens = tokens[longest_prefix:]
                 self.n_tokens = longest_prefix
+
                 if self.verbose:
                     print(
                         f"Llama.generate: {longest_prefix} prefix-match hit, "
                         f"remaining {len(tokens)} prompt tokens to eval",
                         file=sys.stderr,
                     )
+            else:
+                # No prefix matched. Completely clear the KV cache to prevent context poisoning.
+                self.n_tokens = 0
+                self._ctx.memory_clear(True)
 
         # Reset the model state
         if reset:

llama_cpp/llama_chat_format.py

@@ -3817,9 +3817,6 @@ def __call__(self, **kwargs):
         kwargs['stop'] = [self.MINICPMV_EOS_TOKEN, self.MINICPMV_PAD_TOKEN]
 
         llama = kwargs['llama']
-        llama.reset()
-        llama._ctx.memory_clear(True)
-        llama.n_tokens = 0
 
         if hasattr(llama, 'input_ids'):
             llama.input_ids.fill(0)
@@ -3950,11 +3947,6 @@ def __call__(self, **kwargs):
 
         llama = kwargs['llama']
 
-        # Clear state for multiple runs
-        llama.reset()
-        llama._ctx.memory_clear(True)
-        llama.n_tokens = 0
-
         if hasattr(llama, 'input_ids'):
             llama.input_ids.fill(0)
 
@@ -4053,9 +4045,6 @@ def __call__(self, **kwargs):
         kwargs['stop'] = [self.GLM46V_EOS_TOKEN, "<|user|>", "<|observation|>", "<|code_middle|>"] # Stop token patch
 
         llama = kwargs['llama']
-        llama.reset()
-        llama._ctx.memory_clear(True)
-        llama.n_tokens = 0
 
         if hasattr(llama, 'input_ids'):
             llama.input_ids.fill(0)
@@ -4140,9 +4129,6 @@ def __call__(self, **kwargs):
         kwargs['stop'] = [self.GRANITE_EOS_TOKEN]
 
         llama = kwargs['llama']
-        llama.reset()
-        llama._ctx.memory_clear(True)
-        llama.n_tokens = 0
 
         if hasattr(llama, 'input_ids'):
             llama.input_ids.fill(0)
@@ -4205,9 +4191,6 @@ def __init__(self, image_min_tokens: int = -1, image_max_tokens: int = -1, **kwa
     def __call__(self, **kwargs):
 
         llama = kwargs['llama']
-        llama.reset()
-        llama._ctx.memory_clear(True)
-        llama.n_tokens = 0
 
         if hasattr(llama, 'input_ids'):
             llama.input_ids.fill(0)
@@ -4325,9 +4308,6 @@ def __call__(self, **kwargs):
         kwargs['stop'] = [self.PADDLEOCR_EOS_TOKEN]
 
         llama = kwargs['llama']
-        llama.reset()
-        llama._ctx.memory_clear(True)
-        llama.n_tokens = 0
 
         if hasattr(llama, 'input_ids'):
             llama.input_ids.fill(0)
@@ -4381,11 +4361,6 @@ class Qwen25VLChatHandler(Llava15ChatHandler):
     def __call__(self, **kwargs):
         llama = kwargs['llama']
 
-        # Clear state for multiple runs
-        llama.reset()
-        llama._ctx.memory_clear(True)
-        llama.n_tokens = 0
-
         if hasattr(llama, 'input_ids'):
             llama.input_ids.fill(0)
 
@@ -4525,11 +4500,6 @@ def __call__(self, **kwargs):
 
         llama = kwargs['llama']
 
-        # Clear state for multiple runs
-        llama.reset()
-        llama._ctx.memory_clear(True)
-        llama.n_tokens = 0
-
         if hasattr(llama, 'input_ids'):
             llama.input_ids.fill(0)