From b1b57289915a1f1b77d038686d17c1d0698db116 Mon Sep 17 00:00:00 2001 From: Qubitium Date: Fri, 4 Sep 2026 08:14:46 +0000 Subject: [PATCH 1/3] fix(llama-cpp): preserve rendered chat token ids --- evalution/engines/llama_cpp_engine.py | 27 +++++++++++++++++++++++++-- tests/test_engine_tools_threading.py | 4 ++-- tests/test_llama_cpp_engine.py | 16 +++++++++------- 3 files changed, 36 insertions(+), 11 deletions(-) diff --git a/evalution/engines/llama_cpp_engine.py b/evalution/engines/llama_cpp_engine.py index f921501..0988478 100644 --- a/evalution/engines/llama_cpp_engine.py +++ b/evalution/engines/llama_cpp_engine.py @@ -9,7 +9,7 @@ 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 @@ -876,7 +876,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 +898,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 +918,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 +945,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 diff --git a/tests/test_engine_tools_threading.py b/tests/test_engine_tools_threading.py index f7eb7de..6fe067f 100644 --- a/tests/test_engine_tools_threading.py +++ b/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 "" + return [1, 2, 3] if kwargs.get("tokenize") else "" def _chat_request() -> GenerationRequest: diff --git a/tests/test_llama_cpp_engine.py b/tests/test_llama_cpp_engine.py index ea74c38..cd7c4f8 100644 --- a/tests/test_llama_cpp_engine.py +++ b/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, From 570d6bc18ccfea16a7a753bc5d77105cc3aa9803 Mon Sep 17 00:00:00 2001 From: Qubitium Date: Fri, 4 Sep 2026 08:24:27 +0000 Subject: [PATCH 2/3] perf(llama-cpp): reuse shared multiple-choice prefixes --- evalution/engines/llama_cpp_engine.py | 32 +++++++++++++++++++++++++++ tests/test_llama_cpp_engine.py | 19 ++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/evalution/engines/llama_cpp_engine.py b/evalution/engines/llama_cpp_engine.py index 0988478..8339617 100644 --- a/evalution/engines/llama_cpp_engine.py +++ b/evalution/engines/llama_cpp_engine.py @@ -352,6 +352,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 +394,36 @@ 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.""" + + 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 loglikelihood_continuous( self, requests: Iterable[tuple[Any, LoglikelihoodRequest]], diff --git a/tests/test_llama_cpp_engine.py b/tests/test_llama_cpp_engine.py index cd7c4f8..2ecdd5e 100644 --- a/tests/test_llama_cpp_engine.py +++ b/tests/test_llama_cpp_engine.py @@ -568,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"}), + ] From 1a7b4da354281c5b002b3b706ebbee9dadff6e1b Mon Sep 17 00:00:00 2001 From: Qubitium Date: Fri, 4 Sep 2026 09:29:19 +0000 Subject: [PATCH 3/3] perf(llama-cpp): batch distinct scoring prefixes --- evalution/engines/llama_cpp_engine.py | 111 ++++++++++++++++++++++++++ 1 file changed, 111 insertions(+) diff --git a/evalution/engines/llama_cpp_engine.py b/evalution/engines/llama_cpp_engine.py index 8339617..6512517 100644 --- a/evalution/engines/llama_cpp_engine.py +++ b/evalution/engines/llama_cpp_engine.py @@ -16,6 +16,8 @@ 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(), } @@ -400,6 +404,13 @@ def _score_single_token_continuations( ) -> 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()] @@ -424,6 +435,106 @@ def _score_single_token_continuations( 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]],