Fix llama.cpp chat token ID parity

evalution/engines/llama_cpp_engine.py

@@ -9,13 +9,15 @@
 import importlib
 import sys
 import threading
-from collections.abc import Iterable, Iterator
+from collections.abc import Iterable, Iterator, Mapping, Sequence
 from contextlib import closing, suppress
 from dataclasses import asdict, dataclass, field
 from itertools import chain, islice
 from pathlib import Path
 from typing import Any
 
+import numpy as np
+
 from evalution.config import Model
 from evalution.engines.base import (
     BaseEngineDeviceConfig,
@@ -63,6 +65,7 @@ class LlamaCpp(BaseEngineDeviceConfig, SharedEngineConfig):
     chat_format: str | None = None
     verbose: bool = False
     logits_all: bool = True
+    native_loglikelihood_batching: bool = True
     llama_kwargs: dict[str, Any] = field(default_factory=dict)
 
     def build(self, model: Model) -> BaseInferenceSession:
@@ -182,6 +185,7 @@ def describe_execution(self) -> dict[str, Any]:
             "gpu_offload_supported": self.gpu_offload_supported,
             "n_gpu_layers": self.effective_n_gpu_layers,
             "flash_attn": self.config.flash_attn,
+            "native_loglikelihood_batching": self.config.native_loglikelihood_batching,
             "max_model_len": self._max_scoring_input_length(),
         }
 
@@ -352,6 +356,8 @@ def loglikelihood(
 
         del batch_size
         prepared_requests = [self._prepare_loglikelihood_request(request) for request in requests]
+        if all(len(target_ids) == 1 for _prefix_ids, target_ids, _metadata in prepared_requests):
+            return self._score_single_token_continuations(prepared_requests)
         chunk_counts: list[int] = []
         chunk_outputs: list[list[LoglikelihoodOutput]] = []
         with self._generation_lock:
@@ -392,6 +398,143 @@ def loglikelihood(
             )
         return outputs
 
+    def _score_single_token_continuations(
+        self,
+        prepared_requests: list[tuple[list[int], list[int], dict[str, Any]]],
+    ) -> list[LoglikelihoodOutput]:
+        """Score shared-prefix one-token choices once per prefix without changing logits."""
+
+        if (
+            self.config.native_loglikelihood_batching
+            and getattr(self.llm, "_ctx", None) is not None
+            and getattr(self.llm, "_batch", None) is not None
+        ):
+            return self._score_single_token_continuations_native(prepared_requests)
+
+        grouped: dict[tuple[int, ...], list[tuple[int, int, dict[str, Any]]]] = {}
+        for index, (prefix_ids, target_ids, metadata) in enumerate(prepared_requests):
+            effective_prefix = prefix_ids or [self._prefix_token_id()]
+            grouped.setdefault(tuple(effective_prefix), []).append((index, int(target_ids[0]), metadata))
+
+        outputs: list[LoglikelihoodOutput | None] = [None] * len(prepared_requests)
+        with self._generation_lock:
+            for prefix, choices in grouped.items():
+                self.llm.reset()
+                self.llm.eval(list(prefix))
+                token_logprobs = self.llama_module.Llama.logits_to_logprobs(self.llm._scores)[len(prefix) - 1]
+                greedy_token = int(token_logprobs.argmax())
+                for index, token_id, metadata in choices:
+                    outputs[index] = LoglikelihoodOutput(
+                        logprob=float(token_logprobs[token_id]),
+                        is_greedy=greedy_token == token_id,
+                        token_count=1,
+                        metadata=dict(metadata),
+                    )
+
+        if any(output is None for output in outputs):
+            raise RuntimeError("llama.cpp shared-prefix scorer did not produce every requested output")
+        return [output for output in outputs if output is not None]
+
+    def _score_single_token_continuations_native(
+        self,
+        prepared_requests: list[tuple[list[int], list[int], dict[str, Any]]],
+    ) -> list[LoglikelihoodOutput]:
+        """Score multiple distinct prefixes in official llama_batch sequence lanes."""
+
+        grouped: dict[tuple[int, ...], list[tuple[int, int, dict[str, Any]]]] = {}
+        for index, (prefix_ids, target_ids, metadata) in enumerate(prepared_requests):
+            effective_prefix = prefix_ids or [self._prefix_token_id()]
+            grouped.setdefault(tuple(effective_prefix), []).append((index, int(target_ids[0]), metadata))
+
+        ctx = getattr(self.llm, "_ctx", None)
+        batch = getattr(self.llm, "_batch", None)
+        if ctx is None or batch is None:
+            raise RuntimeError("native llama.cpp loglikelihood batching requires low-level internals")
+
+        outputs: list[LoglikelihoodOutput | None] = [None] * len(prepared_requests)
+        pending = list(grouped.items())
+        max_sequences = self._native_sequence_capacity()
+        max_context_tokens = self._max_input_tokens()
+        max_batch_tokens = max(int(self.llm.n_batch), 1)
+        vocabulary_size = int(self.llm.n_vocab())
+
+        with self._generation_lock:
+            while pending:
+                selected: list[tuple[tuple[int, ...], list[tuple[int, int, dict[str, Any]]]]] = []
+                selected_tokens = 0
+                while pending and len(selected) < max_sequences:
+                    prefix, choices = pending[0]
+                    if len(prefix) > max_context_tokens:
+                        raise ValueError("loglikelihood prefix exceeds llama.cpp context window")
+                    if selected and selected_tokens + len(prefix) > max_context_tokens:
+                        break
+                    pending.pop(0)
+                    selected.append((prefix, choices))
+                    selected_tokens += len(prefix)
+
+                ctx.kv_cache_clear()
+                cursors = [0] * len(selected)
+                prefix_logprobs: list[np.ndarray | None] = [None] * len(selected)
+                while any(cursor < len(selected[index][0]) for index, cursor in enumerate(cursors)):
+                    batch.reset()
+                    remaining_capacity = max_batch_tokens
+                    active = [
+                        index
+                        for index, cursor in enumerate(cursors)
+                        if cursor < len(selected[index][0])
+                    ]
+                    for offset, selected_index in enumerate(active):
+                        if remaining_capacity <= 0:
+                            break
+                        prefix = selected[selected_index][0]
+                        active_count = len(active) - offset
+                        take = min(
+                            len(prefix) - cursors[selected_index],
+                            max(remaining_capacity // active_count, 1),
+                        )
+                        chunk = list(prefix[cursors[selected_index] : cursors[selected_index] + take])
+                        self._append_llama_batch_tokens(
+                            batch=batch,
+                            tokens=chunk,
+                            start_pos=cursors[selected_index],
+                            seq_id=selected_index,
+                            request_logits=cursors[selected_index] + take == len(prefix),
+                        )
+                        cursors[selected_index] += take
+                        remaining_capacity -= take
+                    ctx.decode(batch)
+                    for selected_index in active:
+                        if cursors[selected_index] != len(selected[selected_index][0]):
+                            continue
+                        logits_index = next(
+                            index
+                            for index in range(batch.batch.n_tokens - 1, -1, -1)
+                            if batch.batch.seq_id[index][0] == selected_index
+                            and batch.batch.logits[index]
+                        )
+                        prefix_logprobs[selected_index] = np.ctypeslib.as_array(
+                            ctx.get_logits_ith(logits_index),
+                            shape=(vocabulary_size,),
+                        ).copy()
+
+                for selected_index, (_prefix, choices) in enumerate(selected):
+                    logits = prefix_logprobs[selected_index]
+                    if logits is None:
+                        raise RuntimeError("native llama.cpp batch did not return prefix logits")
+                    token_logprobs = self.llama_module.Llama.logits_to_logprobs(logits[None, :])[0]
+                    greedy_token = int(token_logprobs.argmax())
+                    for index, token_id, metadata in choices:
+                        outputs[index] = LoglikelihoodOutput(
+                            logprob=float(token_logprobs[token_id]),
+                            is_greedy=greedy_token == token_id,
+                            token_count=1,
+                            metadata=dict(metadata),
+                        )
+
+        if any(output is None for output in outputs):
+            raise RuntimeError("native llama.cpp batch did not produce every requested output")
+        return [output for output in outputs if output is not None]
+
     def loglikelihood_continuous(
         self,
         requests: Iterable[tuple[Any, LoglikelihoodRequest]],
@@ -876,7 +1019,10 @@ def _generate_one(self, request: GenerationRequest) -> GenerationOutput:
 
         prompt_text, prompt_tokens = self._prepare_generation_prompt(request)
         response = self.llm.create_completion(
-            prompt=prompt_tokens if request.input_ids is not None else prompt_text,
+            # Always pass the already-prepared token ids.  Re-tokenizing an HF
+            # rendered chat prompt inside llama.cpp can duplicate BOS or parse
+            # special-token text differently from the pinned tokenizer.
+            prompt=prompt_tokens,
             max_tokens=request.max_new_tokens,
             temperature=request.temperature if request.do_sample else 0.0,
             stop=list(request.stop) if request.stop else None,
@@ -895,6 +1041,7 @@ def _generate_one(self, request: GenerationRequest) -> GenerationOutput:
     def _prepare_generation_prompt(self, request: GenerationRequest) -> tuple[str, list[int]]:
         """Render one request into the prompt text and prompt tokens consumed by llama.cpp."""
 
+        rendered_token_ids: list[int] | None = None
         if request.rendered_prompt is not None:
             prompt_text = request.rendered_prompt
         elif request.messages is not None:
@@ -914,6 +1061,23 @@ def _prepare_generation_prompt(self, request: GenerationRequest) -> tuple[str, l
                     request.messages,
                     **template_kwargs,
                 )
+                encoded = apply_chat_template(
+                    request.messages,
+                    **{**template_kwargs, "tokenize": True},
+                )
+                if isinstance(encoded, Mapping):
+                    encoded = encoded.get("input_ids")
+                if hasattr(encoded, "tolist"):
+                    encoded = encoded.tolist()
+                if isinstance(encoded, tuple):
+                    encoded = list(encoded)
+                if isinstance(encoded, list) and encoded and isinstance(encoded[0], list):
+                    if len(encoded) != 1:
+                        raise ValueError("chat template returned more than one token sequence")
+                    encoded = encoded[0]
+                if not isinstance(encoded, Sequence) or isinstance(encoded, (str, bytes)):
+                    raise TypeError("chat template must return a list of token ids")
+                rendered_token_ids = [int(token_id) for token_id in encoded]
             else:
                 prompt_text = self._messages_display_prompt(request.messages)
         elif request.prompt is not None:
@@ -924,6 +1088,8 @@ def _prepare_generation_prompt(self, request: GenerationRequest) -> tuple[str, l
         prompt_tokens = (
             list(request.input_ids)
             if request.input_ids is not None
+            else rendered_token_ids
+            if rendered_token_ids is not None
             else self._tokenize_text(prompt_text, add_bos=True)
         )
         return prompt_text, prompt_tokens

tests/test_engine_tools_threading.py

@@ -30,9 +30,9 @@ class RecordingChatTemplateTokenizer:
     def __init__(self) -> None:
         self.calls: list[dict[str, Any]] = []
 
-    def apply_chat_template(self, messages: Any, **kwargs: Any) -> str:
+    def apply_chat_template(self, messages: Any, **kwargs: Any) -> str | list[int]:
         self.calls.append(kwargs)
-        return "<rendered>"
+        return [1, 2, 3] if kwargs.get("tokenize") else "<rendered>"
 
 
 def _chat_request() -> GenerationRequest:

tests/test_llama_cpp_engine.py

@@ -35,8 +35,9 @@ def __init__(self) -> None:
     def apply_chat_template(self, messages, *, tokenize=False, add_generation_prompt=True):
         """Render chat messages into one deterministic prompt string."""
 
-        del tokenize, add_generation_prompt
-        return "\n".join(f"{message['role']}: {message['content']}" for message in messages)
+        del add_generation_prompt
+        rendered = "\n".join(f"{message['role']}: {message['content']}" for message in messages)
+        return [1, *map(ord, rendered)] if tokenize else rendered
 
 
 class FakeLlamaRuntime:
@@ -365,7 +366,7 @@ def test_llama_cpp_session_generate_uses_completion_and_chat_paths() -> None:
 
     assert session.llm.create_completion_calls == [
         {
-            "prompt": "Hello",
+            "prompt": [1, 72, 101, 108, 108, 111],
             "max_tokens": 256,
             "temperature": 0.0,
             "stop": None,
@@ -380,9 +381,10 @@ def test_llama_cpp_session_generate_uses_completion_and_chat_paths() -> None:
             "temperature": 0.0,
             "stop": None,
             "seed": None,
-            "stream": False,
-            "logprobs": False,
-        }
+                "stream": False,
+                "logprobs": False,
+                "tools": None,
+            }
     ]
     assert outputs == [
         GenerationOutput(
@@ -422,7 +424,7 @@ def test_llama_cpp_session_generate_renders_messages_with_prepare_tokenizer() ->
     assert session.llm.create_chat_completion_calls == []
     assert session.llm.create_completion_calls == [
         {
-            "prompt": "user: Hi",
+            "prompt": [1, 117, 115, 101, 114, 58, 32, 72, 105],
             "max_tokens": 256,
             "temperature": 0.0,
             "stop": None,
@@ -566,3 +568,22 @@ def test_llama_cpp_session_loglikelihood_scores_continuation_tokens() -> None:
             metadata={"suite": "demo"},
         )
     ]
+
+
+def test_llama_cpp_session_scores_shared_prefix_single_token_choices_once() -> None:
+    """Verify multiple-choice continuations reuse one byte-identical prefix evaluation."""
+
+    session = _build_session()
+
+    outputs = session.loglikelihood(
+        [
+            LoglikelihoodRequest(context="ab", continuation="\x02", metadata={"choice": "c"}),
+            LoglikelihoodRequest(context="ab", continuation="\x03", metadata={"choice": "d"}),
+        ]
+    )
+
+    assert session.llm.eval_calls == [[1, ord("a"), ord("b")]]
+    assert outputs == [
+        LoglikelihoodOutput(logprob=-10.0, is_greedy=False, token_count=1, metadata={"choice": "c"}),
+        LoglikelihoodOutput(logprob=-10.0, is_greedy=False, token_count=1, metadata={"choice": "d"}),
+    ]