From 0c9a50a855a992c12ac69a6a689ed4fabb5414a5 Mon Sep 17 00:00:00 2001 From: Yaocheng Tong Date: Wed, 4 Nov 2026 14:13:51 -0500 Subject: [PATCH 1/2] fix: limit logits processors to the sampled prefix --- llama_cpp/llama.py | 5 ++++- tests/test_llama.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 14e2f8500..7ae91ae19 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -739,7 +739,9 @@ def apply_func(token_data_array: llama_cpp.llama_token_data_array_p): ), ) for logit_processor in logits_processor: - recarray.logit[:] = logit_processor(self._input_ids, recarray.logit) + recarray.logit[:] = logit_processor( + self._input_ids[: self._sample_prefix_len], recarray.logit + ) sampler.add_custom(apply_func) @@ -843,6 +845,7 @@ def sample( grammar=grammar, ) + self._sample_prefix_len = idx + 1 if idx is not None else self.n_tokens ridx = idx - self.n_tokens if idx is not None else -1 assert self.ctx is not None diff --git a/tests/test_llama.py b/tests/test_llama.py index c1c16e30c..1d6f08a4c 100644 --- a/tests/test_llama.py +++ b/tests/test_llama.py @@ -11,6 +11,7 @@ import llama_cpp import llama_cpp._internals as internals +from llama_cpp.llama_speculative import LlamaPromptLookupDecoding MODEL = "./vendor/llama.cpp/models/ggml-vocab-llama-spm.gguf" @@ -662,3 +663,35 @@ def test_real_llama_embeddings(llama_cpp_embedding_model_path): repeated_embeddings = model.embed(list(reversed(prompts))) assert len(repeated_embeddings) == len(prompts) assert all(len(repeated) == len(embedding) for repeated in repeated_embeddings) + + +def test_logits_processor_sees_only_accepted_tokens_with_draft_model( + llama_cpp_transformer_model_path, +): + model = llama_cpp.Llama( + llama_cpp_transformer_model_path, + n_ctx=512, + draft_model=LlamaPromptLookupDecoding(num_pred_tokens=4), + verbose=False, + ) + prompt = model.tokenize(b"Once upon a time, there was a little girl. " * 4) + generated = [] + prefixes_match = [] + draft_suffix_seen = [] + + def record_prefix(input_ids, logits): + prefixes_match.append(list(input_ids) == prompt + generated) + draft_suffix_seen.append(model.n_tokens > len(input_ids)) + return logits + + for token in model.generate( + prompt, + temp=0.0, + logits_processor=llama_cpp.LogitsProcessorList([record_prefix]), + ): + generated.append(token) + if len(generated) == 32: + break + + assert prefixes_match == [True] * 32 + assert any(draft_suffix_seen) From 15b46b7f7354a88b196f96e09e91baa38b653918 Mon Sep 17 00:00:00 2001 From: Yaocheng Tong Date: Wed, 4 Nov 2026 14:14:16 -0500 Subject: [PATCH 2/2] docs: add changelog entry for #2401 --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729..85145ac57 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +- fix: limit logits processors to the sampled prefix by @YaochengTong in #2401 + ## [0.3.36] - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098