fix: hide unverified draft tokens from logits processors

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- fix: hide unverified draft tokens from logits processors
+
 ## [0.3.36]
 
 - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098

llama_cpp/llama.py

@@ -567,6 +567,10 @@ def free_lora_adapter():
                 )
 
         self._sampler = None
+        # Absolute input_ids position whose logits are being sampled.
+        # The logits-processor callback cannot see sample()'s idx argument.
+        # None keeps the previous behavior: the full evaluated sequence.
+        self._processor_sample_idx: Optional[int] = None
 
         # Cache recurrent/hybrid model detection to avoid repeated FFI calls
         self._is_recurrent = llama_cpp.llama_model_is_recurrent(self._model.model)
@@ -738,8 +742,19 @@ def apply_func(token_data_array: llama_cpp.llama_token_data_array_p):
                         data_soa_address
                     ),
                 )
+                input_ids = self._input_ids
+                prefix_len = self._logits_processor_prefix_len(
+                    self._processor_sample_idx
+                )
+                # Draft and generator-injected tokens are written into the
+                # buffer before they are accepted. Passing that suffix makes
+                # position-dependent processors disagree with greedy decoding.
+                # When nothing is unverified the slice is the whole buffer, so
+                # hand the processor that same array and avoid an extra view.
+                if prefix_len != len(input_ids):
+                    input_ids = input_ids[:prefix_len]
                 for logit_processor in logits_processor:
-                    recarray.logit[:] = logit_processor(self._input_ids, recarray.logit)
+                    recarray.logit[:] = logit_processor(input_ids, recarray.logit)
 
             sampler.add_custom(apply_func)
 
@@ -789,6 +804,25 @@ def apply_func(token_data_array: llama_cpp.llama_token_data_array_p):
                 sampler.add_dist(self._seed)
         return sampler
 
+    def _logits_processor_prefix_len(self, idx: Optional[int]) -> int:
+        """How many leading tokens a logits processor may see.
+
+        ``idx`` is the absolute ``input_ids`` position whose logits are being
+        sampled, the same convention as ``Llama.sample``. The processor prefix
+        is the prompt plus every token up to and including that position.
+        Tokens stored after it are still unverified, so they are excluded.
+
+        ``idx is None`` samples the last evaluated logit and the prefix is the
+        whole evaluated sequence. A negative ``idx`` is not a documented
+        position and keeps that full sequence too.
+        """
+        n_tokens = self.n_tokens
+        if idx is None or idx < 0:
+            return n_tokens
+        if idx + 1 >= n_tokens:
+            return n_tokens
+        return idx + 1
+
     def sample(
         self,
         top_k: int = 40,
@@ -843,6 +877,9 @@ def sample(
                 grammar=grammar,
             )
 
+        # Read by the custom sampler installed in ``_init_sampler``.
+        self._processor_sample_idx = idx
+
         ridx = idx - self.n_tokens if idx is not None else -1
 
         assert self.ctx is not None

tests/test_speculative_logits_processor.py

@@ -0,0 +1,290 @@
+"""Regression tests for logits processors during speculative decoding.
+
+``Llama.generate`` evaluates an accepted token together with the draft block,
+so ``input_ids`` already holds unverified draft tokens when the next position
+is sampled. Processors must see only the prompt plus tokens accepted so far.
+That is the same prefix greedy decoding passes when no draft model is used.
+
+The tests below drive ``generate`` with a stub sampler and do not need a GGUF.
+Manual check against the report in abetlen/llama-cpp-python#2393 (CPU, a few
+seconds on ``stories260K.gguf``):
+
+```bash
+pip install huggingface_hub
+python - <<'PY'
+import numpy as np
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama, LogitsProcessorList
+from llama_cpp.llama_speculative import LlamaPromptLookupDecoding
+
+path = hf_hub_download("ggml-org/models", "tinyllamas/stories260K.gguf")
+PROMPT = b"Once upon a time, there was a little girl. " * 4 + b"Once upon a time,"
+
+def run(draft, use_processor=True, n_new=48):
+    llm = Llama(
+        model_path=path,
+        n_ctx=512,
+        logits_all=True,
+        verbose=False,
+        draft_model=LlamaPromptLookupDecoding(num_pred_tokens=4) if draft else None,
+    )
+    prompt, out, calls, wrong_prefix = llm.tokenize(PROMPT), [], 0, 0
+
+    def parity(input_ids, scores):
+        nonlocal calls, wrong_prefix
+        calls += 1
+        wrong_prefix += list(input_ids) != prompt + out
+        if len(input_ids) % 2 == 0:
+            scores[0::2] = -np.inf
+        return scores
+
+    processor = LogitsProcessorList([parity]) if use_processor else None
+    for tok in llm.generate(prompt, temp=0.0, logits_processor=processor):
+        out.append(tok)
+        if len(out) == n_new:
+            break
+    violations = sum(
+        t % 2 == 0 and (len(prompt) + i) % 2 == 0 for i, t in enumerate(out)
+    )
+    return out, calls, wrong_prefix, violations
+
+assert run(False, use_processor=False)[0] == run(True, use_processor=False)[0]
+for draft in (False, True):
+    out, calls, wrong_prefix, violations = run(draft)
+    assert calls == 48
+    assert wrong_prefix == 0
+    assert violations == 0
+assert run(False)[0] == run(True)[0]
+print("prompt-lookup logits processor prefixes match greedy decoding")
+PY
+```
+"""
+
+import contextlib
+import ctypes
+
+import numpy as np
+
+import llama_cpp
+import llama_cpp.llama as llama_module
+from llama_cpp.llama_speculative import LlamaPromptLookupDecoding
+
+
+def _bare_llama() -> llama_cpp.Llama:
+    llm = llama_cpp.Llama.__new__(llama_cpp.Llama)
+    # ``__del__`` closes this stack. A real model sets it in ``__init__``.
+    llm._stack = contextlib.ExitStack()
+    return llm
+
+
+def test_logits_processor_prefix_len_hides_only_unverified_suffix():
+    llm = _bare_llama()
+    llm.n_tokens = 10
+
+    assert llm._logits_processor_prefix_len(None) == 10
+    assert llm._logits_processor_prefix_len(-1) == 10
+    # Last evaluated position: same prefix greedy decoding already passes.
+    assert llm._logits_processor_prefix_len(9) == 10
+    assert llm._logits_processor_prefix_len(4) == 5
+    # Past the buffer: do not read uninitialized tokens.
+    assert llm._logits_processor_prefix_len(100) == 10
+
+
+class _Context:
+    def __init__(self):
+        self.ctx = object()
+
+    def kv_cache_seq_rm(self, *args, **kwargs):
+        return True
+
+
+class _ScriptedSampler:
+    """Stand-in for ``LlamaSampler`` that runs the real processor callback."""
+
+    def __init__(self):
+        self.apply_func = None
+        self._candidates = None
+
+    def add_custom(self, apply_func):
+        self.apply_func = apply_func
+
+    def add_penalties(self, **kwargs):
+        return None
+
+    def add_greedy(self):
+        return None
+
+    def sample(self, ctx, idx):
+        del ctx, idx
+        size = 8
+        data = (llama_cpp.llama_token_data * size)()
+        for token_id in range(size):
+            data[token_id].id = token_id
+            # Prefer the highest id until a processor masks logits.
+            data[token_id].logit = float(token_id)
+            data[token_id].p = 0.0
+        array = llama_cpp.llama_token_data_array()
+        array.data = data
+        array.size = size
+        array.selected = -1
+        array.sorted = False
+        # Keep the buffer alive across the callback.
+        self._candidates = data
+        if self.apply_func is not None:
+            self.apply_func(ctypes.pointer(array))
+        logits = np.array([data[i].logit for i in range(size)], dtype=np.float32)
+        return int(np.argmax(logits))
+
+
+def _install_sampler(monkeypatch):
+    monkeypatch.setattr(llama_module.internals, "LlamaSampler", _ScriptedSampler)
+
+
+def _llama(draft):
+    llm = _bare_llama()
+    n_ctx = 64
+    llm.n_tokens = 0
+    llm._requires_eval = True
+    llm._n_ctx = n_ctx
+    llm._n_vocab = 8
+    llm.last_n_tokens_size = 64
+    llm._is_recurrent = False
+    llm._is_hybrid = False
+    llm._sampler = None
+    llm._processor_sample_idx = None
+    llm.input_ids = np.zeros((n_ctx,), dtype=np.intc)
+    llm.scores = np.zeros((n_ctx, llm._n_vocab), dtype=np.single)
+    llm._ctx = _Context()
+    llm.draft_model = (
+        LlamaPromptLookupDecoding(max_ngram_size=2, num_pred_tokens=2)
+        if draft
+        else None
+    )
+
+    def eval_tokens(tokens):
+        batch = list(tokens)
+        n_past = llm.n_tokens
+        llm.input_ids[n_past : n_past + len(batch)] = batch
+        llm.n_tokens += len(batch)
+        llm._requires_eval = False
+
+    llm.eval = eval_tokens
+    return llm
+
+
+def _run(llm, prompt, n_new):
+    seen = []
+    buffered = []
+
+    def choose_from_prefix_length(input_ids, scores):
+        # Token id is a pure function of the prefix length, in 1..5.
+        seen.append([int(token) for token in input_ids])
+        buffered.append(int(llm.n_tokens))
+        scores[:] = -np.inf
+        scores[(len(input_ids) % 5) + 1] = 10.0
+        return scores
+
+    processor = llama_cpp.LogitsProcessorList([choose_from_prefix_length])
+    out = []
+    for token in llm.generate(prompt, temp=0.0, logits_processor=processor):
+        out.append(int(token))
+        if len(out) == n_new:
+            break
+    return out, seen, buffered
+
+
+def test_generate_logits_processor_sees_accepted_prefix_only(monkeypatch):
+    _install_sampler(monkeypatch)
+    # The first accepted token repeats an ngram, so prompt lookup drafts ahead.
+    prompt = [4, 1, 2, 3, 4]
+    n_new = 8
+
+    plain = _llama(draft=False)
+    plain_out, plain_seen, plain_buffered = _run(plain, prompt, n_new)
+    drafted = _llama(draft=True)
+    draft_out, draft_seen, draft_buffered = _run(drafted, prompt, n_new)
+
+    assert plain_out == draft_out
+    assert plain_seen == draft_seen
+    assert len(plain_out) == n_new
+
+    for generated, prefix in enumerate(plain_seen):
+        assert prefix == prompt + plain_out[:generated]
+        assert len(prefix) == plain_buffered[generated]
+
+    for generated, prefix in enumerate(draft_seen):
+        assert prefix == prompt + draft_out[:generated]
+        # The evaluated buffer still contains the draft suffix here.
+        assert len(prefix) <= draft_buffered[generated]
+
+    assert any(
+        len(prefix) < n_tokens for prefix, n_tokens in zip(draft_seen, draft_buffered)
+    )
+    # The length-derived token is what the callback wrote, not the raw argmax.
+    assert plain_out[0] == (len(prompt) % 5) + 1
+    assert plain_out[0] != 7
+
+
+def test_min_tokens_processor_ignores_unverified_drafts(monkeypatch):
+    """``MinTokensLogitsProcessor`` must not treat draft tokens as generated."""
+    _install_sampler(monkeypatch)
+    prompt = [4, 1, 2, 3, 4]
+    eos = 7
+    llm = _llama(draft=True)
+    accepted = []
+    masked = []
+    min_tokens = llama_cpp.MinTokensLogitsProcessor(min_tokens=3, token_eos=eos)
+
+    def observe(input_ids, scores):
+        accepted.append(len(input_ids) - len(prompt))
+        masked.append(bool(np.isneginf(scores[eos])))
+        scores[:] = -np.inf
+        scores[(len(input_ids) % 5) + 1] = 10.0
+        return scores
+
+    processor = llama_cpp.LogitsProcessorList([min_tokens, observe])
+    out = []
+    for token in llm.generate(prompt, temp=0.0, logits_processor=processor):
+        out.append(int(token))
+        if len(out) == 4:
+            break
+
+    # One call per accepted token. Drafts must not advance the count to 3 early.
+    assert accepted == [0, 1, 2, 3]
+    assert masked == [True, True, True, False]
+    assert out[0] == (len(prompt) % 5) + 1
+
+
+def test_generate_without_processor_is_unchanged_by_draft_tokens(monkeypatch):
+    """No callback is installed, so a draft model must not change scripted output."""
+
+    script = [1, 0, 1, 0, 2, 2]
+
+    class Scripted(_ScriptedSampler):
+        def __init__(self):
+            super().__init__()
+            self._script = list(script)
+
+        def sample(self, ctx, idx):
+            del ctx, idx
+            if self.apply_func is not None:
+                raise AssertionError(
+                    "logits processor callback should not be installed"
+                )
+            return self._script.pop(0)
+
+    monkeypatch.setattr(llama_module.internals, "LlamaSampler", Scripted)
+    prompt = [1, 0, 1, 0, 1, 0]
+    n_new = len(script)
+
+    def collect(draft):
+        llm = _llama(draft)
+        out = []
+        for token in llm.generate(prompt, temp=0.0, logits_processor=None):
+            out.append(int(token))
+            if len(out) == n_new:
+                break
+        return out
+
+    assert collect(False) == script
+    assert collect(True) == script