fix: handle Qwen 3.5 hybrid prefix reuse

Author: abetlen

llama_cpp/_internals.py

@@ -291,10 +291,10 @@ def kv_cache_clear(self):
         assert self.memory is not None, "Memory is not initialized"
         llama_cpp.llama_memory_clear(self.memory, True)
 
-    def kv_cache_seq_rm(self, seq_id: int, p0: int, p1: int):
+    def kv_cache_seq_rm(self, seq_id: int, p0: int, p1: int) -> bool:
         assert self.memory is not None, "Memory is not initialized"
         seq_id = seq_id if seq_id >= 0 else 0
-        llama_cpp.llama_memory_seq_rm(self.memory, seq_id, p0, p1)
+        return llama_cpp.llama_memory_seq_rm(self.memory, seq_id, p0, p1)
 
     def kv_cache_seq_cp(self, seq_id_src: int, seq_id_dst: int, p0: int, p1: int):
         assert self.memory is not None, "Memory is not initialized"

llama_cpp/llama.py

@@ -891,13 +891,20 @@ def generate(
                 else:
                     break
             if longest_prefix > 0:
-                reset = False
-                tokens = tokens[longest_prefix:]
-                self.n_tokens = longest_prefix
-                if self.verbose:
+                if self._ctx.kv_cache_seq_rm(-1, longest_prefix, -1):
+                    reset = False
+                    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,
+                        )
+                elif self.verbose:
                     print(
-                        f"Llama.generate: {longest_prefix} prefix-match hit, "
-                        f"remaining {len(tokens)} prompt tokens to eval",
+                        f"Llama.generate: {longest_prefix} prefix-match found "
+                        f"but partial kv removal not supported, re-evaluating full prompt",
                         file=sys.stderr,
                     )
 

tests/test_llama.py

@@ -1,5 +1,7 @@
 import ctypes
 import multiprocessing
+from types import SimpleNamespace
+from unittest.mock import Mock
 
 import numpy as np
 from scipy.special import log_softmax
@@ -15,6 +17,10 @@
 MODEL = "./vendor/llama.cpp/models/ggml-vocab-llama-spm.gguf"
 
 
+class EvalCalled(Exception):
+    pass
+
+
 def test_llama_cpp_version():
     assert llama_cpp.__version__
 
@@ -232,3 +238,75 @@ def test_real_llama_embeddings(llama_cpp_model_path):
     )
     # Smoke test for now
     model.embed("Hello World")
+
+
+def test_kv_cache_seq_rm_returns_bool(monkeypatch):
+    context = internals.LlamaContext.__new__(internals.LlamaContext)
+    context.memory = object()
+    calls = []
+
+    def fake_llama_memory_seq_rm(memory, seq_id, p0, p1):
+        calls.append((memory, seq_id, p0, p1))
+        return True
+
+    monkeypatch.setattr(llama_cpp, "llama_memory_seq_rm", fake_llama_memory_seq_rm)
+
+    assert context.kv_cache_seq_rm(-1, 4, -1) is True
+    assert calls == [(context.memory, 0, 4, -1)]
+
+
+def make_test_llama(kv_cache_seq_rm_return):
+    llama = llama_cpp.Llama.__new__(llama_cpp.Llama)
+    llama.n_tokens = 3
+    llama.n_batch = 8
+    llama._n_ctx = 32
+    llama._n_vocab = 8
+    llama._logits_all = False
+    llama._seed = 1337
+    llama.last_n_tokens_size = 64
+    llama.verbose = False
+    llama.input_ids = np.array([1, 2, 3, 0, 0, 0], dtype=np.intc)
+    llama.scores = np.zeros((6, 8), dtype=np.single)
+    llama._ctx = SimpleNamespace(
+        kv_cache_seq_rm=Mock(return_value=kv_cache_seq_rm_return)
+    )
+    llama._sampler = None
+    llama.eval_tokens_seen = None
+    llama.reset_calls = 0
+
+    def reset():
+        llama.reset_calls += 1
+        llama.n_tokens = 0
+
+    def eval_tokens(tokens):
+        llama.eval_tokens_seen = list(tokens)
+        raise EvalCalled
+
+    llama.reset = reset
+    llama.eval = eval_tokens
+    llama._init_sampler = lambda **kwargs: object()
+    return llama
+
+
+def test_generate_reuses_prefix_when_partial_removal_supported():
+    llama = make_test_llama(True)
+
+    with pytest.raises(EvalCalled):
+        next(llama.generate([1, 2, 3, 4]))
+
+    llama._ctx.kv_cache_seq_rm.assert_called_once_with(-1, 3, -1)
+    assert llama.reset_calls == 0
+    assert llama.n_tokens == 3
+    assert llama.eval_tokens_seen == [4]
+
+
+def test_generate_falls_back_to_reset_when_partial_removal_rejected():
+    llama = make_test_llama(False)
+
+    with pytest.raises(EvalCalled):
+        next(llama.generate([1, 2, 3, 4]))
+
+    llama._ctx.kv_cache_seq_rm.assert_called_once_with(-1, 3, -1)
+    assert llama.reset_calls == 1
+    assert llama.n_tokens == 0
+    assert llama.eval_tokens_seen == [1, 2, 3, 4]