diff --git a/evalution/engines/llama_cpp_engine.py b/evalution/engines/llama_cpp_engine.py index f921501..6512517 100644 --- a/evalution/engines/llama_cpp_engine.py +++ b/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 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..2ecdd5e 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, @@ -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"}), + ]