diff --git a/evalution/benchmarks/mmlu.py b/evalution/benchmarks/mmlu.py index 7bb3e0c..93cd8e9 100644 --- a/evalution/benchmarks/mmlu.py +++ b/evalution/benchmarks/mmlu.py @@ -356,9 +356,30 @@ def _evaluate_continuous( else {} ) - def iter_request_stream() -> Any: - """Iterate over request stream. Keep the nested traversal explicit so ordering and metadata stay aligned.""" - for index, doc in enumerate(docs): + def iter_subject_groups() -> Any: + """Yield contiguous subject groups without reading past an explicit row limit.""" + group: list[dict[str, Any]] = [] + group_subject: str | None = None + seen = 0 + for doc in docs: + subject_key = normalize_subset_token(str(doc["subject"])) + if group and subject_key != group_subject: + yield group + group = [] + group.append(doc) + group_subject = subject_key + seen += 1 + # ``limit_docs`` uses islice for streaming datasets. Do not ask that + # iterator for one more row merely to detect the final group boundary. + if self.max_rows is not None and seen >= self.max_rows: + break + if group: + yield group + + def iter_request_stream(group_docs: list[dict[str, Any]], start_index: int) -> Any: + """Iterate one subject group while preserving global sample indices.""" + for local_index, doc in enumerate(group_docs): + index = start_index + local_index subject = str(doc["subject"]) subject_key = normalize_subset_token(subject) prompt = _fewshot_prompt( @@ -378,20 +399,21 @@ def iter_request_stream() -> Any: metadata=dict(request_progress_metadata), ) - for (sample_index, choice_index), output in loglikelihood_continuous( - iter_request_stream(), - batch_size=self.batch_size, - ): - sample_choice_scores[int(sample_index)].append( - ( - int(choice_index), - output.logprob, - output.token_count, + for group_docs in iter_subject_groups(): + group_start = len(sample_docs) + for (sample_index, choice_index), output in loglikelihood_continuous( + iter_request_stream(group_docs, group_start), + batch_size=self.batch_size, + ): + sample_choice_scores[int(sample_index)].append( + ( + int(choice_index), + output.logprob, + output.token_count, + ) ) - ) - if score_bar is not None: - score_bar.next().draw() - + if score_bar is not None: + score_bar.next().draw() executed = len(sample_docs) if total is None: logger.info("%s: executed %d sample(s)", task_name, executed) diff --git a/evalution/engines/base.py b/evalution/engines/base.py index 146705e..c78ae5c 100644 --- a/evalution/engines/base.py +++ b/evalution/engines/base.py @@ -268,6 +268,20 @@ class BaseEngineTransformersRuntimeConfig(BaseEngineDeviceConfig): """Define the base engine transformers runtime config helper class.""" attn_implementation: str | None = None device_map: str | dict[str, Any] | None = None + # Enable scorer-side reuse for repeated long prefixes in loglikelihood workloads. + # Generation continuous batching owns a separate KV cache and cannot serve direct scoring. + # The optimization is exact and narrowly applies to duplicate one-token-choice prefixes; + # callers can still disable it for compatibility with ``loglikelihood_prefix_cache=False``. + loglikelihood_prefix_cache: bool = True + # Prefill eligible shared prefixes before their scoring batch so scoring lookups are warm. + # The prefill itself is reported separately from steady-state cache-hit metrics. + loglikelihood_prefix_cache_prewarm: bool = True + # Limit one batched prefill to avoid an unbounded activation/cache peak. + loglikelihood_prefix_cache_prewarm_batch_size: int = 32 + # MMLU-style subject groups can release their prefix entries as soon as the group completes. + loglikelihood_prefix_cache_release_after_group: bool = True + loglikelihood_prefix_cache_min_tokens: int = 256 + loglikelihood_prefix_cache_max_entries: int = 32 @dataclass diff --git a/evalution/engines/transformers_common.py b/evalution/engines/transformers_common.py index 3400e87..aa45e6d 100644 --- a/evalution/engines/transformers_common.py +++ b/evalution/engines/transformers_common.py @@ -10,6 +10,7 @@ import random import sys import threading +from collections import OrderedDict from collections.abc import Mapping from contextlib import contextmanager, suppress from contextvars import ContextVar @@ -85,6 +86,38 @@ ) +def _is_recoverable_cache_runtime_error(exc: RuntimeError) -> bool: + """Identify narrow cache-API runtime mismatches that can safely use the old scorer.""" + try: + import torch + + if isinstance(exc, torch.OutOfMemoryError): + return False + except (ImportError, AttributeError): + pass + message = str(exc).lower() + fatal_markers = ( + "out of memory", + "illegal memory access", + "device-side assert", + "launch failure", + "misaligned address", + "unspecified launch failure", + "cuda error", + ) + if any(marker in message for marker in fatal_markers): + return False + compatibility_markers = ( + "past_key_values", + "batch_select_indices", + "batch_repeat_interleave", + "cache object", + "cache type", + "unsupported cache", + ) + return any(marker in message for marker in compatibility_markers) + + @dataclass(slots=True) class _ScoringChunk: # Track one model forward slice plus which token span contributes to the score. @@ -144,6 +177,23 @@ class BaseTransformerSession(BaseInferenceSession): init=False, repr=False, ) + # Direct loglikelihood forwards do not pass through the generation manager. Keep a small, + # explicit LRU of immutable prefix caches for repeated-prefix scoring workloads. + _loglikelihood_prefix_cache: OrderedDict[tuple[int, ...], Any] = field( + default_factory=OrderedDict, + init=False, + repr=False, + ) + _loglikelihood_prefix_cache_hits: int = field(default=0, init=False, repr=False) + _loglikelihood_prefix_cache_misses: int = field(default=0, init=False, repr=False) + _loglikelihood_prefix_cache_warmup_hits: int = field(default=0, init=False, repr=False) + _loglikelihood_prefix_cache_warmup_misses: int = field(default=0, init=False, repr=False) + _loglikelihood_prefix_cache_warmup_batches: int = field(default=0, init=False, repr=False) + _loglikelihood_prefix_cache_warmup_prefixes: int = field(default=0, init=False, repr=False) + _loglikelihood_prefix_cache_released_prefixes: int = field(default=0, init=False, repr=False) + _loglikelihood_prefix_cache_release_calls: int = field(default=0, init=False, repr=False) + _loglikelihood_prefix_cache_disabled: bool = field(default=False, init=False, repr=False) + _loglikelihood_prefix_cache_reported: bool = field(default=False, init=False, repr=False) # Run fixed-batch generation for engines that do not own a paged continuous batching manager. def generate( @@ -179,10 +229,13 @@ def loglikelihood( with self._generation_lock: prepared_requests = [self._prepare_loglikelihood_request(request) for request in requests] effective_batch_size = batch_size or self._resolve_scoring_batch_size(prepared_requests) - return self._score_prepared_loglikelihood_requests( - prepared_requests, - batch_size=effective_batch_size, - ) + try: + return self._score_prepared_loglikelihood_requests( + prepared_requests, + batch_size=effective_batch_size, + ) + finally: + self._release_loglikelihood_prefix_cache_after_call() # Keep log-likelihood request submission decoupled from caller iteration so suites can stream # rows lazily while this session keeps fixed-size scoring batches full on a worker thread. @@ -240,14 +293,17 @@ def consume_requests( if prepared_batch: self._emit_scored_loglikelihood_batch(prepared_batch, put_result) - yield from stream_request_results( - items, - producer_name=f"{type(self).__name__}.loglikelihood_request_producer", - consumer_name=f"{type(self).__name__}.loglikelihood_request_consumer", - process_requests=consume_requests, - require_non_main_thread=self.request_executor_requires_non_main_thread, - request_queue_max_size=max(effective_batch_size * 2, 1), - ) + try: + yield from stream_request_results( + items, + producer_name=f"{type(self).__name__}.loglikelihood_request_producer", + consumer_name=f"{type(self).__name__}.loglikelihood_request_consumer", + process_requests=consume_requests, + require_non_main_thread=self.request_executor_requires_non_main_thread, + request_queue_max_size=max(effective_batch_size * 2, 1), + ) + finally: + self._release_loglikelihood_prefix_cache_after_call() return iterator() @@ -405,6 +461,17 @@ def gc(self) -> None: with self._state_lock: self.stop_criteria_cache.clear() self.auto_batch_size_cache.clear() + self._loglikelihood_prefix_cache.clear() + self._loglikelihood_prefix_cache_hits = 0 + self._loglikelihood_prefix_cache_misses = 0 + self._loglikelihood_prefix_cache_warmup_hits = 0 + self._loglikelihood_prefix_cache_warmup_misses = 0 + self._loglikelihood_prefix_cache_warmup_batches = 0 + self._loglikelihood_prefix_cache_warmup_prefixes = 0 + self._loglikelihood_prefix_cache_released_prefixes = 0 + self._loglikelihood_prefix_cache_release_calls = 0 + self._loglikelihood_prefix_cache_disabled = False + self._loglikelihood_prefix_cache_reported = False self.execution_logged = False gc.collect() with suppress(Exception): @@ -1043,6 +1110,117 @@ def _score_chunks( if score_bar is not None: batch_index = (start // batch_size) + 1 score_bar.subtitle(f"batch={batch_index}/{total_batches} batch_size={batch_size}") + + # Multiple-choice suites commonly submit one-token continuations for + # the same context (for example, `` A``, `` B``, `` C``, `` D`` in + # MMLU). The ordinary path below appends each candidate token and + # re-runs the complete decoder prefill for every choice. When a + # scoring batch contains duplicate prefixes, run each prefix once + # and gather all candidate logits from its final position. This is + # exact for causal models: the logit at the last context token is + # the next-token distribution used to score every one-token + # continuation. Keep the optimization narrowly guarded so + # multi-token and non-identical requests retain the established + # token-by-token semantics. + pad_token_id = getattr(self.tokenizer, "pad_token_id", None) + if pad_token_id is None: + pad_token_id = self._prefix_token_id() + choice_groups: dict[tuple[int, ...], list[tuple[int, _ScoringChunk]]] = {} + if batch and all( + chunk.score_count == 1 and chunk.score_start == len(chunk.input_ids) - 1 + for chunk in batch + ): + for batch_index, chunk in enumerate(batch): + prefix = tuple(chunk.input_ids[: chunk.score_start]) + choice_groups.setdefault(prefix, []).append((batch_index, chunk)) + if choice_groups and len(choice_groups) < len(batch): + self._prewarm_loglikelihood_prefix_cache_for_batch( + choice_groups, + pad_token_id=int(pad_token_id), + ) + cached_choice_outputs = self._score_one_token_choice_prefix_cache_partitioned( + batch, + choice_groups=choice_groups, + pad_token_id=int(pad_token_id), + ) + if cached_choice_outputs is not None: + scored_chunks.extend(cached_choice_outputs) + continue + encoded = None + logits = None + try: + prefixes = list(choice_groups) + padded_length = max(len(prefix) for prefix in prefixes) + padded_rows = [ + list(prefix) + ([int(pad_token_id)] * (padded_length - len(prefix))) + for prefix in prefixes + ] + encoded = { + "input_ids": torch.tensor( + padded_rows, + dtype=torch.long, + device=self.input_device, + ) + } + with self._scoring_attention_context(): + with torch.inference_mode(): + outputs = self.model(**encoded) + logits = outputs.logits + context_lengths = torch.tensor( + [len(prefix) for prefix in prefixes], + dtype=torch.long, + device=logits.device, + ) + row_indices = torch.arange( + len(prefixes), + dtype=torch.long, + device=logits.device, + ) + final_logits = logits[row_indices, context_lengths - 1, :] + final_log_probs = torch.log_softmax(final_logits, dim=-1) + final_greedy = final_logits.argmax(dim=-1) + # One packed D2H transfer for the reduced choice rows, + # matching the no-sync reduction policy of the normal path. + packed_rows: list[torch.Tensor] = [] + for group_index, prefix in enumerate(prefixes): + for _batch_index, chunk in choice_groups[prefix]: + target = int(chunk.input_ids[chunk.score_start]) + packed_rows.append( + torch.stack( + ( + final_log_probs[group_index, target], + (final_greedy[group_index] == target).to(final_log_probs.dtype), + ) + ) + ) + host_rows = torch.stack(packed_rows).detach().cpu().tolist() + # Restore the original batch order; choice_groups is keyed + # by prefix for the compact forward, not request ordering. + output_by_batch_index: dict[int, tuple[float, bool]] = {} + packed_cursor = 0 + for prefix in prefixes: + for batch_index, _chunk in choice_groups[prefix]: + logprob, is_greedy = host_rows[packed_cursor] + output_by_batch_index[batch_index] = (float(logprob), bool(is_greedy)) + packed_cursor += 1 + for batch_index, chunk in enumerate(batch): + logprob, is_greedy = output_by_batch_index[batch_index] + scored_chunks.append( + LoglikelihoodOutput( + logprob=logprob, + is_greedy=is_greedy, + token_count=1, + metadata=dict(chunk.metadata), + ) + ) + if score_bar is not None: + score_bar.next().draw() + finally: + if logits is not None: + del logits + if encoded is not None: + del encoded + continue encoded = None logits = None try: @@ -1086,6 +1264,13 @@ def _score_chunks( shift_log_probs = torch.log_softmax(logits[:, :-1, :], dim=-1) shift_labels = encoded["input_ids"][:, 1:] + # Keep reductions and greedy checks on device until every row in + # this forward has been assembled. ``Tensor.item()`` and + # ``torch.equal()`` both synchronize CUDA streams, so doing + # either inside this loop serializes continuous scoring one + # continuation at a time. + batch_logprobs: list[Any] = [] + batch_greedy: list[Any] = [] for row_index, chunk in enumerate(batch): if logits_to_keep is not None: # `logits_to_keep` returns only the tail window logits, so absolute @@ -1100,10 +1285,27 @@ def _score_chunks( sample_targets = shift_labels[row_index, shift_start:shift_end] gathered = sample_log_probs.gather(-1, sample_targets.unsqueeze(-1)).squeeze(-1) greedy_tokens = sample_log_probs.argmax(dim=-1) + + batch_logprobs.append(gathered.sum()) + batch_greedy.append(torch.all(greedy_tokens == sample_targets)) + + # Preserve the original reduction dtype while packing the bool + # flag into the same transfer. This is one synchronization and + # one contiguous D2H copy per forward batch rather than one per + # scored continuation. + packed = torch.stack( + [ + torch.stack(batch_logprobs), + torch.stack(batch_greedy).to(dtype=batch_logprobs[0].dtype), + ], + dim=1, + ) + host_rows = packed.detach().cpu().tolist() + for chunk, (logprob, is_greedy) in zip(batch, host_rows, strict=True): scored_chunks.append( LoglikelihoodOutput( - logprob=float(gathered.sum().item()), - is_greedy=bool(torch.equal(greedy_tokens, sample_targets)), + logprob=float(logprob), + is_greedy=bool(is_greedy), token_count=chunk.score_count, metadata=dict(chunk.metadata), ) @@ -1117,6 +1319,592 @@ def _score_chunks( del encoded return scored_chunks + # Reuse one long prefix KV prefill for a batch of divergent one-token choices. This is kept + # separate from the exact-prefix choice deduplication above: MMLU questions share a subject + # few-shot prefix, but their full contexts are different. The cache is enabled by default for + # transformer sessions and remains narrowly guarded; callers can disable it when a custom + # model wrapper does not support cache-enabled prefix forwards. + def _score_one_token_choice_prefix_cache( + self, + batch: list[_ScoringChunk], + *, + choice_groups: dict[tuple[int, ...], list[tuple[int, _ScoringChunk]]], + pad_token_id: int, + ) -> list[LoglikelihoodOutput] | None: + """Score one-token choices through a shared prefix KV cache when the batch is cacheable.""" + import torch + + if not self._loglikelihood_prefix_cache_enabled() or not choice_groups: + return None + + min_tokens = self._loglikelihood_prefix_cache_min_tokens() + contexts = list(choice_groups) + if not contexts or any(len(context) <= min_tokens for context in contexts): + return None + + # Use a fixed prefix length so the same cache key remains discoverable when a subject + # crosses a scoring-batch boundary. The prompt construction is responsible for making + # this prefix identical; unrelated batches simply fall back to the established scorer. + prefix_keys = {tuple(context[:min_tokens]) for context in contexts} + if len(prefix_keys) != 1: + return None + prefix_key = next(iter(prefix_keys)) + suffixes = [context[min_tokens:] for context in contexts] + if any(not suffix for suffix in suffixes): + return None + + try: + prefix_cache = self._get_loglikelihood_prefix_cache( + prefix_key, + prefix_token_id=pad_token_id, + ) + if prefix_cache is None: + prefix_ids = torch.tensor( + [list(prefix_key)], + dtype=torch.long, + device=self.input_device, + ) + with self._scoring_attention_context(), torch.inference_mode(): + prefix_outputs = self.model( + input_ids=prefix_ids, + use_cache=True, + ) + prefix_cache = getattr(prefix_outputs, "past_key_values", None) + if prefix_cache is None: + self._disable_loglikelihood_prefix_cache( + "model did not return past_key_values for a cache-enabled prefill" + ) + return None + self._put_loglikelihood_prefix_cache(prefix_key, prefix_cache) + self._loglikelihood_prefix_cache_misses += 1 + else: + self._loglikelihood_prefix_cache_hits += 1 + + max_suffix_length = max(len(suffix) for suffix in suffixes) + suffix_rows = [ + list(suffix) + [pad_token_id] * (max_suffix_length - len(suffix)) + for suffix in suffixes + ] + suffix_ids = torch.tensor( + suffix_rows, + dtype=torch.long, + device=self.input_device, + ) + branch_cache = self._repeat_loglikelihood_cache(deepcopy(prefix_cache), len(suffixes)) + if branch_cache is None: + self._disable_loglikelihood_prefix_cache( + "past_key_values does not support batch cache expansion" + ) + return None + with self._scoring_attention_context(), torch.inference_mode(): + outputs = self.model( + input_ids=suffix_ids, + past_key_values=branch_cache, + use_cache=True, + ) + logits = outputs.logits + row_indices = torch.arange( + len(suffixes), + dtype=torch.long, + device=logits.device, + ) + final_positions = torch.tensor( + [len(suffix) - 1 for suffix in suffixes], + dtype=torch.long, + device=logits.device, + ) + final_logits = logits[row_indices, final_positions, :] + final_log_probs = torch.log_softmax(final_logits, dim=-1) + final_greedy = final_logits.argmax(dim=-1) + + packed_rows: list[torch.Tensor] = [] + output_indices: list[int] = [] + for group_index, context in enumerate(contexts): + for batch_index, chunk in choice_groups[context]: + target = int(chunk.input_ids[chunk.score_start]) + packed_rows.append( + torch.stack( + ( + final_log_probs[group_index, target], + (final_greedy[group_index] == target).to(final_log_probs.dtype), + ) + ) + ) + output_indices.append(batch_index) + host_rows = torch.stack(packed_rows).detach().cpu().tolist() + output_by_batch_index: dict[int, LoglikelihoodOutput] = {} + for batch_index, (logprob, is_greedy) in zip(output_indices, host_rows, strict=True): + chunk = batch[batch_index] + output_by_batch_index[batch_index] = LoglikelihoodOutput( + logprob=float(logprob), + is_greedy=bool(is_greedy), + token_count=1, + metadata=dict(chunk.metadata), + ) + self._report_loglikelihood_prefix_cache_once() + return [output_by_batch_index[index] for index in range(len(batch))] + except (TypeError, AttributeError, NotImplementedError) as exc: + self._disable_loglikelihood_prefix_cache(str(exc)) + return None + except RuntimeError as exc: + if not _is_recoverable_cache_runtime_error(exc): + raise + self._disable_loglikelihood_prefix_cache(str(exc)) + return None + + # Partition mixed-task scoring batches by their fixed shared-prefix key. MMLU rows are + # subject-contiguous in the source dataset, but a batch can straddle two subjects; optimizing + # the eligible buckets independently keeps the remaining contexts on the exact old path. + def _score_one_token_choice_prefix_cache_partitioned( + self, + batch: list[_ScoringChunk], + *, + choice_groups: dict[tuple[int, ...], list[tuple[int, _ScoringChunk]]], + pad_token_id: int, + ) -> list[LoglikelihoodOutput] | None: + """Score cacheable prefix buckets and fall back for ineligible contexts.""" + if not self._loglikelihood_prefix_cache_enabled() or not choice_groups: + return None + min_tokens = self._loglikelihood_prefix_cache_min_tokens() + buckets: dict[tuple[int, ...], list[tuple[int, ...]]] = {} + for context in choice_groups: + if len(context) <= min_tokens: + continue + key = tuple(context[:min_tokens]) + buckets.setdefault(key, []).append(context) + + optimized_by_batch_index: dict[int, LoglikelihoodOutput] = {} + for contexts in buckets.values(): + # A cache hit can make a singleton worthwhile after a prior batch populated it; a + # cache miss is intentionally deferred until at least two divergent contexts share it. + if len(contexts) == 1 and ( + self._get_loglikelihood_prefix_cache( + contexts[0][:min_tokens], + prefix_token_id=pad_token_id, + ) + is None + ): + continue + local_indices = [ + batch_index + for context in contexts + for batch_index, _chunk in choice_groups[context] + ] + local_index_by_global = { + global_index: local_index for local_index, global_index in enumerate(local_indices) + } + local_batch = [batch[index] for index in local_indices] + local_groups = { + context: [ + (local_index_by_global[batch_index], chunk) + for batch_index, chunk in choice_groups[context] + ] + for context in contexts + } + local_outputs = self._score_one_token_choice_prefix_cache( + local_batch, + choice_groups=local_groups, + pad_token_id=pad_token_id, + ) + if local_outputs is None: + continue + for local_index, output in enumerate(local_outputs): + optimized_by_batch_index[local_indices[local_index]] = output + + if not optimized_by_batch_index: + return None + + remaining_indices = [ + index for index in range(len(batch)) if index not in optimized_by_batch_index + ] + fallback_by_batch_index: dict[int, LoglikelihoodOutput] = {} + if remaining_indices: + remaining_batch = [batch[index] for index in remaining_indices] + fallback_outputs = self._score_chunks( + remaining_batch, + batch_size=max(len(remaining_batch), 1), + ) + for local_index, output in enumerate(fallback_outputs): + fallback_by_batch_index[remaining_indices[local_index]] = output + + ordered_outputs: list[LoglikelihoodOutput] = [] + for index in range(len(batch)): + output = optimized_by_batch_index.get(index) + if output is None: + output = fallback_by_batch_index[index] + ordered_outputs.append(output) + return ordered_outputs + + def _loglikelihood_prefix_cache_enabled(self) -> bool: + """Return whether scorer-side prefix KV reuse was explicitly requested.""" + return bool( + getattr(self.config, "loglikelihood_prefix_cache", False) + and not self._loglikelihood_prefix_cache_disabled + and self._loglikelihood_prefix_cache_max_entries() > 0 + ) + + def _loglikelihood_prefix_cache_prewarm_enabled(self) -> bool: + """Return whether eligible shared prefixes should be filled before scoring.""" + return bool( + self._loglikelihood_prefix_cache_enabled() + and getattr(self.config, "loglikelihood_prefix_cache_prewarm", False) + ) + + def _release_loglikelihood_prefix_cache_after_call(self) -> None: + """Drop scorer-side KV entries when one loglikelihood request lifecycle ends.""" + if not self._loglikelihood_prefix_cache_enabled() and not self._loglikelihood_prefix_cache: + return + try: + release_after_call = bool( + getattr(self.config, "loglikelihood_prefix_cache_release_after_group", True) + ) + except (TypeError, ValueError): + release_after_call = True + if release_after_call: + self.release_loglikelihood_prefix_cache() + + def _loglikelihood_prefix_cache_prewarm_batch_size(self) -> int: + """Resolve the bounded number of distinct prefixes in one prefill forward.""" + try: + return max( + int(getattr(self.config, "loglikelihood_prefix_cache_prewarm_batch_size", 32)), + 1, + ) + except (TypeError, ValueError): + return 32 + + def _prewarm_loglikelihood_prefix_cache_for_batch( + self, + choice_groups: dict[tuple[int, ...], list[tuple[int, _ScoringChunk]]], + *, + pad_token_id: int, + ) -> None: + """Prefill one shared prefix per eligible choice batch before suffix scoring. + + Multiple-choice rows are grouped by their fixed token prefix rather than by answer + label. A prefix is worth prewarming only when at least two distinct question contexts + in the current batch share it; a singleton would pay a prefill without avoiding any + subsequent work. The scoring path then observes a cache hit for every eligible lookup, + while the compulsory model prefill is reported separately as a warmup miss. + """ + if not self._loglikelihood_prefix_cache_prewarm_enabled() or not choice_groups: + return + + min_tokens = self._loglikelihood_prefix_cache_min_tokens() + prefix_contexts: dict[tuple[int, ...], set[tuple[int, ...]]] = {} + for context in choice_groups: + if len(context) <= min_tokens: + continue + prefix_key = tuple(context[:min_tokens]) + prefix_contexts.setdefault(prefix_key, set()).add(context) + + eligible_prefixes = [ + prefix_key + for prefix_key, contexts in prefix_contexts.items() + if len(contexts) >= 2 + ] + if not eligible_prefixes: + return + max_entries = self._loglikelihood_prefix_cache_max_entries() + # Never prefill more entries than the LRU can retain. Otherwise early prefixes would be + # evicted before this scoring batch consumes them, turning the prewarm into wasted work. + eligible_prefixes = eligible_prefixes[:max_entries] + self._loglikelihood_prefix_cache_warmup_prefixes += len(eligible_prefixes) + + import torch + + missing_prefixes: list[tuple[int, ...]] = [] + for prefix_key in eligible_prefixes: + cached = self._get_loglikelihood_prefix_cache( + prefix_key, + prefix_token_id=pad_token_id, + ) + if cached is None: + missing_prefixes.append(prefix_key) + else: + self._loglikelihood_prefix_cache_warmup_hits += 1 + + if not missing_prefixes: + return + + batch_size = min( + self._loglikelihood_prefix_cache_prewarm_batch_size(), + max_entries, + ) + for start in range(0, len(missing_prefixes), batch_size): + prefix_batch = missing_prefixes[start : start + batch_size] + try: + prefix_ids = torch.tensor( + [list(prefix_key) for prefix_key in prefix_batch], + dtype=torch.long, + device=self.input_device, + ) + with self._scoring_attention_context(), torch.inference_mode(): + prefix_outputs = self.model( + input_ids=prefix_ids, + use_cache=True, + ) + prefix_cache = getattr(prefix_outputs, "past_key_values", None) + if prefix_cache is None: + self._disable_loglikelihood_prefix_cache( + "model did not return past_key_values for a cache-enabled prefill" + ) + return + for index, prefix_key in enumerate(prefix_batch): + selected_cache = ( + prefix_cache + if len(prefix_batch) == 1 + else self._select_loglikelihood_cache_row(prefix_cache, index) + ) + if selected_cache is None: + # Unknown cache wrappers may not expose per-layer tensors. Preserve the + # optimization by retrying this chunk as independent one-prefix forwards + # rather than disabling caching for the whole session. + for fallback_prefix_key in prefix_batch: + self._prefill_loglikelihood_prefix(fallback_prefix_key) + self._loglikelihood_prefix_cache_warmup_batches += 1 + len(prefix_batch) + self._loglikelihood_prefix_cache_warmup_misses += len(prefix_batch) + break + self._put_loglikelihood_prefix_cache(prefix_key, selected_cache) + else: + self._loglikelihood_prefix_cache_warmup_batches += 1 + self._loglikelihood_prefix_cache_warmup_misses += len(prefix_batch) + except (TypeError, AttributeError, NotImplementedError) as exc: + self._disable_loglikelihood_prefix_cache(str(exc)) + return + except RuntimeError as exc: + if not _is_recoverable_cache_runtime_error(exc): + raise + self._disable_loglikelihood_prefix_cache(str(exc)) + return + + def _prefill_loglikelihood_prefix(self, prefix_key: tuple[int, ...]) -> None: + """Prefill one prefix for cache wrappers that cannot split a batched result.""" + import torch + + prefix_ids = torch.tensor( + [list(prefix_key)], + dtype=torch.long, + device=self.input_device, + ) + with self._scoring_attention_context(), torch.inference_mode(): + prefix_outputs = self.model( + input_ids=prefix_ids, + use_cache=True, + ) + prefix_cache = getattr(prefix_outputs, "past_key_values", None) + if prefix_cache is None: + raise TypeError("model did not return past_key_values for a cache-enabled prefill") + self._put_loglikelihood_prefix_cache(prefix_key, prefix_cache) + + def _select_loglikelihood_cache_row(self, cache: Any, index: int) -> Any | None: + """Extract one independent batch row from a prefetched KV cache.""" + import torch + + selected = deepcopy(cache) + # Native HF caches may shard/offload layers independently. Build one index on each + # key/value tensor's device instead of passing a single input-device index to all layers. + layers = getattr(selected, "layers", None) + if layers is not None: + if not layers: + return selected + for layer in layers: + keys = getattr(layer, "keys", None) + values = getattr(layer, "values", None) + if keys is None or values is None: + return None + key_indices = torch.tensor([index], dtype=torch.long, device=keys.device) + value_indices = torch.tensor([index], dtype=torch.long, device=values.device) + layer.keys = keys.index_select(0, key_indices) + layer.values = values.index_select(0, value_indices) + return selected + + # Older wrappers expose separate key/value lists rather than ``layers``. + key_cache = getattr(selected, "key_cache", None) + value_cache = getattr(selected, "value_cache", None) + if isinstance(key_cache, (tuple, list)) and isinstance(value_cache, (tuple, list)): + if len(key_cache) != len(value_cache): + return None + selected_keys = [] + selected_values = [] + for keys, values in zip(key_cache, value_cache, strict=True): + if keys is None or values is None: + return None + key_indices = torch.tensor([index], dtype=torch.long, device=keys.device) + value_indices = torch.tensor([index], dtype=torch.long, device=values.device) + selected_keys.append(keys.index_select(0, key_indices)) + selected_values.append(values.index_select(0, value_indices)) + selected.key_cache = ( + tuple(selected_keys) if isinstance(key_cache, tuple) else selected_keys + ) + selected.value_cache = ( + tuple(selected_values) if isinstance(value_cache, tuple) else selected_values + ) + return selected + + # Unknown cache wrappers are handled by the caller's unbatched-prefill fallback. Avoid + # invoking a selector that may internally broadcast one index tensor across devices. + if isinstance(selected, (tuple, list)): + selected_layers = [] + for layer in selected: + if not isinstance(layer, (tuple, list)) or len(layer) < 2: + return None + tensors = [] + for tensor in layer[:2]: + if not hasattr(tensor, "index_select"): + return None + indices = torch.tensor([index], dtype=torch.long, device=tensor.device) + tensors.append(tensor.index_select(0, indices)) + selected_layers.append( + tuple(tensors) + tuple(layer[2:]) + ) + return tuple(selected_layers) if isinstance(selected, tuple) else selected_layers + return None + + def _loglikelihood_prefix_cache_min_tokens(self) -> int: + """Resolve the minimum prefix length that justifies a second prefill.""" + try: + return max(int(getattr(self.config, "loglikelihood_prefix_cache_min_tokens", 256)), 1) + except (TypeError, ValueError): + return 256 + + def _loglikelihood_prefix_cache_max_entries(self) -> int: + """Resolve the bounded GPU cache capacity.""" + try: + return max(int(getattr(self.config, "loglikelihood_prefix_cache_max_entries", 32)), 0) + except (TypeError, ValueError): + return 32 + + def _get_loglikelihood_prefix_cache( + self, + prefix_key: tuple[int, ...], + *, + prefix_token_id: int, + ) -> Any | None: + """Look up one prefix cache and retain it as the most recently used entry.""" + del prefix_token_id # Reserved for future cache-key validation across tokenizer changes. + with self._state_lock: + cached = self._loglikelihood_prefix_cache.get(prefix_key) + if cached is not None: + self._loglikelihood_prefix_cache.move_to_end(prefix_key) + return cached + + def _put_loglikelihood_prefix_cache(self, prefix_key: tuple[int, ...], prefix_cache: Any) -> None: + """Insert one prefix cache and evict the oldest entry before GPU memory grows unbounded.""" + with self._state_lock: + self._loglikelihood_prefix_cache[prefix_key] = prefix_cache + self._loglikelihood_prefix_cache.move_to_end(prefix_key) + max_entries = self._loglikelihood_prefix_cache_max_entries() + while len(self._loglikelihood_prefix_cache) > max_entries: + self._loglikelihood_prefix_cache.popitem(last=False) + + def release_loglikelihood_prefix_cache( + self, + prefix_keys: list[tuple[int, ...]] | tuple[tuple[int, ...], ...] | None = None, + ) -> int: + """Release scorer prefix KV entries whose request group has completed. + + Dropping the Python references releases the underlying device tensors when no temporary + suffix branch still owns them. The CUDA allocator may retain freed blocks for reuse, so + this deliberately avoids an implicit ``empty_cache`` synchronization in the hot path. + """ + with self._state_lock: + if prefix_keys is None: + released = len(self._loglikelihood_prefix_cache) + self._loglikelihood_prefix_cache.clear() + else: + released = 0 + for prefix_key in prefix_keys: + if self._loglikelihood_prefix_cache.pop(tuple(prefix_key), None) is not None: + released += 1 + self._loglikelihood_prefix_cache_release_calls += 1 + self._loglikelihood_prefix_cache_released_prefixes += released + return released + + def _repeat_loglikelihood_cache(self, cache: Any, repeats: int) -> Any | None: + """Expand one prefix cache for a divergent suffix batch without mutating the stored entry.""" + repeat_interleave = getattr(cache, "batch_repeat_interleave", None) + if callable(repeat_interleave): + repeat_interleave(repeats) + return cache + if isinstance(cache, (tuple, list)): + # Support legacy Transformers tuple caches used by older model wrappers. + repeated_layers = [] + for layer in cache: + if not isinstance(layer, (tuple, list)) or len(layer) < 2: + return None + for tensor in layer[:2]: + if not hasattr(tensor, "repeat_interleave"): + return None + repeated_layers.append( + tuple(tensor.repeat_interleave(repeats, dim=0) for tensor in layer[:2]) + + tuple(layer[2:]) + ) + return tuple(repeated_layers) if isinstance(cache, tuple) else repeated_layers + return None + + def _disable_loglikelihood_prefix_cache(self, reason: str) -> None: + """Disable the optional optimization after a model-specific cache incompatibility.""" + if self._loglikelihood_prefix_cache_disabled: + return + self._loglikelihood_prefix_cache_disabled = True + get_logger().warning( + "disabling loglikelihood prefix KV cache after model incompatibility: %s", + reason, + ) + + def _report_loglikelihood_prefix_cache_once(self) -> None: + """Emit one auditable hit/miss report after the first successful cached scoring batch.""" + if self._loglikelihood_prefix_cache_reported: + return + self._loglikelihood_prefix_cache_reported = True + total = self._loglikelihood_prefix_cache_hits + self._loglikelihood_prefix_cache_misses + hit_rate = self._loglikelihood_prefix_cache_hits / total if total else 0.0 + get_logger().info( + "loglikelihood prefix KV cache active: hits=%d misses=%d hit_rate=%.1f%% " + "warmup_hits=%d warmup_misses=%d steady_state_hit_rate=%.1f%% entries=%d", + self._loglikelihood_prefix_cache_hits, + self._loglikelihood_prefix_cache_misses, + hit_rate * 100.0, + self._loglikelihood_prefix_cache_warmup_hits, + self._loglikelihood_prefix_cache_warmup_misses, + hit_rate * 100.0, + len(self._loglikelihood_prefix_cache), + ) + + def loglikelihood_prefix_cache_stats(self) -> dict[str, int | float]: + """Return scorer-side prefix KV cache counters for benchmark telemetry.""" + lookups = self._loglikelihood_prefix_cache_hits + self._loglikelihood_prefix_cache_misses + warmup_lookups = ( + self._loglikelihood_prefix_cache_warmup_hits + + self._loglikelihood_prefix_cache_warmup_misses + ) + return { + "hits": self._loglikelihood_prefix_cache_hits, + "misses": self._loglikelihood_prefix_cache_misses, + "lookups": lookups, + "hit_rate": self._loglikelihood_prefix_cache_hits / lookups if lookups else 0.0, + "warmup_hits": self._loglikelihood_prefix_cache_warmup_hits, + "warmup_misses": self._loglikelihood_prefix_cache_warmup_misses, + "warmup_lookups": warmup_lookups, + "warmup_hit_rate": ( + self._loglikelihood_prefix_cache_warmup_hits / warmup_lookups + if warmup_lookups + else 0.0 + ), + "steady_state_hits": self._loglikelihood_prefix_cache_hits, + "steady_state_misses": self._loglikelihood_prefix_cache_misses, + "steady_state_lookups": lookups, + "steady_state_hit_rate": ( + self._loglikelihood_prefix_cache_hits / lookups if lookups else 0.0 + ), + "warmup_batches": self._loglikelihood_prefix_cache_warmup_batches, + "warmup_prefixes": self._loglikelihood_prefix_cache_warmup_prefixes, + "released_prefixes": self._loglikelihood_prefix_cache_released_prefixes, + "release_calls": self._loglikelihood_prefix_cache_release_calls, + "entries": len(self._loglikelihood_prefix_cache), + } + # Reuse the shared chunk scorer for both eager and continuous log-likelihood submission paths. def _score_prepared_loglikelihood_requests( self, diff --git a/pyproject.toml b/pyproject.toml index 956dee5..b192e17 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "Evalution" -version = "0.0.16" +version = "0.0.17" description = "Modern LLM model evaluation for Transformers, SGLang, vLLM, TensorRT-LLM, llama.cpp, GPTQModel, OpenAI-compatible HTTP backends, and OpenVINO." readme = "README.md" requires-python = ">=3.10" diff --git a/tests/test_mmlu.py b/tests/test_mmlu.py index ca400d6..8d3611a 100644 --- a/tests/test_mmlu.py +++ b/tests/test_mmlu.py @@ -339,6 +339,75 @@ def loglikelihood_continuous(self, requests, *, batch_size=None): ] +def test_mmlu_processes_subject_groups_for_session_lifecycle(monkeypatch) -> None: + """Submit contiguous subjects separately so the session can release each lifecycle.""" + test = Dataset.from_list( + [ + { + "question": "Algebra question", + "subject": "abstract_algebra", + "choices": ["a", "b", "c", "d"], + "answer": 0, + }, + { + "question": "Learning question", + "subject": "machine_learning", + "choices": ["a", "b", "c", "d"], + "answer": 0, + }, + ] + ) + + def fake_load_suite_dataset( + loader, + *, + task_name, + dataset_path, + dataset_name, + split, + cache_dir, + stream, + purpose=None, + ): + del loader, task_name, dataset_path, dataset_name, cache_dir, stream, purpose + assert split == "test" + return test, 0.0 + + class GroupReleaseSession: + """Provide continuous scoring and an observable prefix-cache release hook.""" + + config = SimpleNamespace( + loglikelihood_prefix_cache=True, + loglikelihood_prefix_cache_release_after_group=True, + ) + + def __init__(self) -> None: + self.group_calls = 0 + + def loglikelihood_continuous(self, requests, *, batch_size=None): + assert batch_size == 8 + self.group_calls += 1 + for (sample_index, choice_index), _request in requests: + yield (sample_index, choice_index), LoglikelihoodOutput( + logprob=0.0 if choice_index == 0 else -1.0, + is_greedy=choice_index == 0, + token_count=1, + ) + + session = GroupReleaseSession() + monkeypatch.setattr(mmlu_module, "load_suite_dataset", fake_load_suite_dataset) + result = evalution.benchmarks.mmlu( + subsets="all", + num_fewshot=0, + max_rows=2, + batch_size=8, + stream=True, + ).evaluate(session) + + assert len(result.samples) == 2 + assert session.group_calls == 2 + + def test_mmlu_subset_node_filter_uses_distinct_result_name(monkeypatch) -> None: """Verify MMLU subset node filter uses distinct result name.""" test = Dataset.from_list( diff --git a/tests/test_transformer.py b/tests/test_transformer.py index 8c47935..69b09b4 100644 --- a/tests/test_transformer.py +++ b/tests/test_transformer.py @@ -42,6 +42,10 @@ def test_transformer_defaults_batch_size_to_auto() -> None: assert engine.q_padding_interval_size == 0 assert engine.kv_padding_interval_size == 0 assert engine.max_cached_graphs == 0 + assert engine.loglikelihood_prefix_cache is True + assert engine.loglikelihood_prefix_cache_prewarm is True + assert engine.loglikelihood_prefix_cache_prewarm_batch_size == 32 + assert engine.loglikelihood_prefix_cache_release_after_group is True assert engine.to_dict()["batch_size"] == "auto" assert engine.to_dict()["seed"] is None assert engine.to_dict()["manual_eviction"] is False @@ -53,6 +57,8 @@ def test_transformer_defaults_batch_size_to_auto() -> None: assert engine.to_dict()["q_padding_interval_size"] == 0 assert engine.to_dict()["kv_padding_interval_size"] == 0 assert engine.to_dict()["max_cached_graphs"] == 0 + assert engine.to_dict()["loglikelihood_prefix_cache"] is True + assert engine.to_dict()["loglikelihood_prefix_cache_prewarm"] is True def test_resolve_input_device_keeps_cpu_only_hf_device_maps_on_cpu() -> None: @@ -2313,6 +2319,376 @@ def __call__(self, *, input_ids, attention_mask=None): assert outputs[1].logprob == pytest.approx(expected_token_logprob, abs=1e-6) +def test_transformer_session_loglikelihood_deduplicates_one_token_choice_prefixes() -> None: + """Score repeated one-token choices from one prefix forward.""" + + class FakeTokenizer: + pad_token_id = 0 + eos_token_id = 1 + padding_side = "right" + + class FakeModel: + def __init__(self) -> None: + self.config = SimpleNamespace(max_position_embeddings=8) + self.calls: list[tuple[int, int]] = [] + + def __call__(self, *, input_ids, **kwargs): + assert kwargs == {} + self.calls.append(tuple(input_ids.shape)) + batch_size, sequence_length = input_ids.shape + vocab_size = 16 + logits = torch.full((batch_size, sequence_length, vocab_size), -2.0) + # Make the final prefix position predict the next token with a + # distinct score for each candidate, matching a causal LM. + for row in range(batch_size): + logits[row, -1, 4] = 3.0 + logits[row, -1, 7] = 2.0 + logits[row, -1, 8] = 1.0 + logits[row, -1, 9] = 0.0 + return SimpleNamespace(logits=logits) + + model = FakeModel() + session = TransformersSession( + config=Transformers(batch_size=4), + model_config=Model(path="/tmp/model"), + model=model, + tokenizer=FakeTokenizer(), + input_device=torch.device("cpu"), + ) + + outputs = session.loglikelihood( + [ + LoglikelihoodRequest(context_input_ids=[5, 6], continuation_input_ids=[4]), + LoglikelihoodRequest(context_input_ids=[5, 6], continuation_input_ids=[7]), + LoglikelihoodRequest(context_input_ids=[5, 6], continuation_input_ids=[8]), + LoglikelihoodRequest(context_input_ids=[5, 6], continuation_input_ids=[9]), + ], + batch_size=4, + ) + + assert model.calls == [(1, 2)] + assert [output.token_count for output in outputs] == [1, 1, 1, 1] + assert [output.is_greedy for output in outputs] == [True, False, False, False] + reference_logits = torch.full((16,), -2.0) + reference_logits[[4, 7, 8, 9]] = torch.tensor([3.0, 2.0, 1.0, 0.0]) + reference_logprobs = torch.log_softmax(reference_logits, dim=0) + expected = [float(reference_logprobs[index].item()) for index in [4, 7, 8, 9]] + assert [output.logprob for output in outputs] == pytest.approx(expected, abs=1e-6) + + +def test_transformer_session_loglikelihood_reuses_shared_prefix_kv_cache() -> None: + """Reuse one prefix prefill for divergent one-token choice contexts.""" + + class FakeTokenizer: + pad_token_id = 0 + eos_token_id = 1 + padding_side = "right" + + class FakeCache: + def __init__(self) -> None: + self.repeats = 1 + + def batch_repeat_interleave(self, repeats: int) -> None: + self.repeats = repeats + + class FakeModel: + def __init__(self) -> None: + self.config = SimpleNamespace(max_position_embeddings=32) + self.calls: list[tuple[tuple[int, int], bool, bool]] = [] + + def __call__(self, *, input_ids, use_cache=False, past_key_values=None, **kwargs): + assert kwargs == {} + self.calls.append((tuple(input_ids.shape), bool(use_cache), past_key_values is not None)) + batch_size, sequence_length = input_ids.shape + logits = torch.full((batch_size, sequence_length, 16), -2.0) + logits[:, -1, 4] = 3.0 + if past_key_values is None: + assert use_cache is True + return SimpleNamespace(logits=logits, past_key_values=FakeCache()) + assert use_cache is True + assert past_key_values.repeats == batch_size + return SimpleNamespace(logits=logits) + + model = FakeModel() + session = TransformersSession( + config=Transformers( + batch_size=8, + loglikelihood_prefix_cache=True, + loglikelihood_prefix_cache_prewarm=False, + loglikelihood_prefix_cache_release_after_group=False, + loglikelihood_prefix_cache_min_tokens=2, + loglikelihood_prefix_cache_max_entries=4, + ), + model_config=Model(path="/tmp/model"), + model=model, + tokenizer=FakeTokenizer(), + input_device=torch.device("cpu"), + ) + requests = [ + LoglikelihoodRequest(context_input_ids=[5, 6, 7], continuation_input_ids=[4]), + LoglikelihoodRequest(context_input_ids=[5, 6, 7], continuation_input_ids=[8]), + LoglikelihoodRequest(context_input_ids=[5, 6, 8], continuation_input_ids=[4]), + LoglikelihoodRequest(context_input_ids=[5, 6, 8], continuation_input_ids=[8]), + ] + + outputs = session.loglikelihood(requests, batch_size=4) + assert len(outputs) == 4 + assert model.calls == [((1, 2), True, False), ((2, 1), True, True)] + assert session.loglikelihood_prefix_cache_stats() == { + "hits": 0, + "misses": 1, + "lookups": 1, + "hit_rate": 0.0, + "warmup_hits": 0, + "warmup_misses": 0, + "warmup_lookups": 0, + "warmup_hit_rate": 0.0, + "steady_state_hits": 0, + "steady_state_misses": 1, + "steady_state_lookups": 1, + "steady_state_hit_rate": 0.0, + "warmup_batches": 0, + "warmup_prefixes": 0, + "released_prefixes": 0, + "release_calls": 0, + "entries": 1, + } + session.loglikelihood(requests, batch_size=4) + assert model.calls == [ + ((1, 2), True, False), + ((2, 1), True, True), + ((2, 1), True, True), + ] + assert session.loglikelihood_prefix_cache_stats() == { + "hits": 1, + "misses": 1, + "lookups": 2, + "hit_rate": 0.5, + "warmup_hits": 0, + "warmup_misses": 0, + "warmup_lookups": 0, + "warmup_hit_rate": 0.0, + "steady_state_hits": 1, + "steady_state_misses": 1, + "steady_state_lookups": 2, + "steady_state_hit_rate": 0.5, + "warmup_batches": 0, + "warmup_prefixes": 0, + "released_prefixes": 0, + "release_calls": 0, + "entries": 1, + } + + +def test_transformer_session_loglikelihood_prewarm_reports_steady_state_hits() -> None: + """Prefill a shared prefix once and keep the scoring lookup phase warm.""" + + class FakeTokenizer: + pad_token_id = 0 + eos_token_id = 1 + padding_side = "right" + + class FakeCache: + def __init__(self) -> None: + self.repeats = 1 + + def batch_repeat_interleave(self, repeats: int) -> None: + self.repeats = repeats + + class FakeModel: + def __init__(self) -> None: + self.config = SimpleNamespace(max_position_embeddings=32) + self.calls: list[tuple[tuple[int, int], bool, bool]] = [] + + def __call__(self, *, input_ids, use_cache=False, past_key_values=None, **kwargs): + assert kwargs == {} + self.calls.append((tuple(input_ids.shape), bool(use_cache), past_key_values is not None)) + batch_size, sequence_length = input_ids.shape + logits = torch.full((batch_size, sequence_length, 16), -2.0) + logits[:, -1, 4] = 3.0 + if past_key_values is None: + assert use_cache is True + return SimpleNamespace(logits=logits, past_key_values=FakeCache()) + assert use_cache is True + assert past_key_values.repeats == batch_size + return SimpleNamespace(logits=logits) + + model = FakeModel() + session = TransformersSession( + config=Transformers( + batch_size=8, + loglikelihood_prefix_cache=True, + loglikelihood_prefix_cache_prewarm=True, + loglikelihood_prefix_cache_release_after_group=False, + loglikelihood_prefix_cache_min_tokens=2, + loglikelihood_prefix_cache_max_entries=4, + ), + model_config=Model(path="/tmp/model"), + model=model, + tokenizer=FakeTokenizer(), + input_device=torch.device("cpu"), + ) + requests = [ + LoglikelihoodRequest(context_input_ids=[5, 6, 7], continuation_input_ids=[4]), + LoglikelihoodRequest(context_input_ids=[5, 6, 7], continuation_input_ids=[8]), + LoglikelihoodRequest(context_input_ids=[5, 6, 8], continuation_input_ids=[4]), + LoglikelihoodRequest(context_input_ids=[5, 6, 8], continuation_input_ids=[8]), + ] + + outputs = session.loglikelihood(requests, batch_size=4) + assert len(outputs) == 4 + assert model.calls == [((1, 2), True, False), ((2, 1), True, True)] + assert session.loglikelihood_prefix_cache_stats() == { + "hits": 1, + "misses": 0, + "lookups": 1, + "hit_rate": 1.0, + "warmup_hits": 0, + "warmup_misses": 1, + "warmup_lookups": 1, + "warmup_hit_rate": 0.0, + "steady_state_hits": 1, + "steady_state_misses": 0, + "steady_state_lookups": 1, + "steady_state_hit_rate": 1.0, + "warmup_batches": 1, + "warmup_prefixes": 1, + "released_prefixes": 0, + "release_calls": 0, + "entries": 1, + } + assert session.release_loglikelihood_prefix_cache() == 1 + released_stats = session.loglikelihood_prefix_cache_stats() + assert released_stats["entries"] == 0 + assert released_stats["released_prefixes"] == 1 + assert released_stats["release_calls"] == 1 + + +def test_transformer_session_loglikelihood_prewarm_batches_prefixes() -> None: + """Batch distinct shared prefixes and split their cache rows for suffix scoring.""" + + class FakeTokenizer: + pad_token_id = 0 + eos_token_id = 1 + padding_side = "right" + + class FakeCache: + def __init__(self) -> None: + self.repeats = 1 + self.layers = [ + SimpleNamespace( + keys=torch.empty((2, 1, 1, 1)), + values=torch.empty((2, 1, 1, 1)), + ), + SimpleNamespace( + keys=torch.empty((2, 1, 1, 1), device="meta"), + values=torch.empty((2, 1, 1, 1), device="meta"), + ), + ] + + def batch_select_indices(self, indices) -> None: + raise AssertionError("per-layer cache selection must not use one global device") + + def batch_repeat_interleave(self, repeats: int) -> None: + self.repeats = repeats + + class FakeModel: + def __init__(self) -> None: + self.config = SimpleNamespace(max_position_embeddings=32) + self.calls: list[tuple[tuple[int, int], bool, bool]] = [] + + def __call__(self, *, input_ids, use_cache=False, past_key_values=None, **kwargs): + assert kwargs == {} + self.calls.append((tuple(input_ids.shape), bool(use_cache), past_key_values is not None)) + batch_size, sequence_length = input_ids.shape + logits = torch.full((batch_size, sequence_length, 16), -2.0) + logits[:, -1, 4] = 3.0 + if past_key_values is None: + return SimpleNamespace(logits=logits, past_key_values=FakeCache()) + assert past_key_values.repeats == batch_size + return SimpleNamespace(logits=logits) + + model = FakeModel() + session = TransformersSession( + config=Transformers( + batch_size=8, + loglikelihood_prefix_cache=True, + loglikelihood_prefix_cache_prewarm=True, + loglikelihood_prefix_cache_prewarm_batch_size=2, + loglikelihood_prefix_cache_min_tokens=2, + loglikelihood_prefix_cache_max_entries=4, + ), + model_config=Model(path="/tmp/model"), + model=model, + tokenizer=FakeTokenizer(), + input_device=torch.device("cpu"), + ) + requests = [ + LoglikelihoodRequest(context_input_ids=[5, 6, 7], continuation_input_ids=[4]), + LoglikelihoodRequest(context_input_ids=[5, 6, 7], continuation_input_ids=[8]), + LoglikelihoodRequest(context_input_ids=[5, 6, 8], continuation_input_ids=[4]), + LoglikelihoodRequest(context_input_ids=[5, 6, 8], continuation_input_ids=[8]), + LoglikelihoodRequest(context_input_ids=[9, 10, 11], continuation_input_ids=[4]), + LoglikelihoodRequest(context_input_ids=[9, 10, 11], continuation_input_ids=[8]), + LoglikelihoodRequest(context_input_ids=[9, 10, 12], continuation_input_ids=[4]), + LoglikelihoodRequest(context_input_ids=[9, 10, 12], continuation_input_ids=[8]), + ] + + outputs = session.loglikelihood(requests, batch_size=8) + assert len(outputs) == 8 + assert model.calls == [ + ((2, 2), True, False), + ((2, 1), True, True), + ((2, 1), True, True), + ] + stats = session.loglikelihood_prefix_cache_stats() + assert stats["hits"] == 2 + assert stats["misses"] == 0 + assert stats["steady_state_hit_rate"] == 1.0 + assert stats["warmup_batches"] == 1 + assert stats["warmup_prefixes"] == 2 + assert stats["warmup_misses"] == 2 + assert stats["released_prefixes"] == 2 + assert stats["release_calls"] == 1 + assert stats["entries"] == 0 + + +def test_transformer_session_loglikelihood_prefix_cache_propagates_oom() -> None: + """Do not hide allocator failures behind an uncached retry.""" + + class FakeTokenizer: + pad_token_id = 0 + eos_token_id = 1 + padding_side = "right" + + class OOMModel: + config = SimpleNamespace(max_position_embeddings=32) + + def __call__(self, **kwargs): + del kwargs + raise torch.OutOfMemoryError("CUDA out of memory") + + session = TransformersSession( + config=Transformers( + batch_size=4, + loglikelihood_prefix_cache=True, + loglikelihood_prefix_cache_min_tokens=2, + ), + model_config=Model(path="/tmp/model"), + model=OOMModel(), + tokenizer=FakeTokenizer(), + input_device=torch.device("cpu"), + ) + requests = [ + LoglikelihoodRequest(context_input_ids=[5, 6, 7], continuation_input_ids=[4]), + LoglikelihoodRequest(context_input_ids=[5, 6, 7], continuation_input_ids=[8]), + ] + + with pytest.raises(torch.OutOfMemoryError, match="out of memory"): + session.loglikelihood(requests, batch_size=2) + assert session._loglikelihood_prefix_cache_disabled is False + + def test_transformer_session_loglikelihood_sorts_requests_by_total_length_before_scoring( monkeypatch, ) -> None: diff --git a/tests/test_transformer_loglikelihood_buffer.py b/tests/test_transformer_loglikelihood_buffer.py new file mode 100644 index 0000000..c0c884e --- /dev/null +++ b/tests/test_transformer_loglikelihood_buffer.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from contextlib import nullcontext +from types import SimpleNamespace + +import torch + +from evalution.engines.transformers_common import BaseTransformerSession, _ScoringChunk + + +def test_score_chunks_batches_cuda_scalar_reductions() -> None: + class Model: + def __call__(self, input_ids, **_kwargs): + batch, length = input_ids.shape + logits = torch.full((batch, length, 16), -4.0) + for row in range(batch): + for position in range(length - 1): + logits[row, position, int(input_ids[row, position + 1])] = 4.0 + return SimpleNamespace(logits=logits) + + session = SimpleNamespace( + tokenizer=SimpleNamespace(pad_token_id=0), + input_device=torch.device("cpu"), + model=Model(), + _scoring_attention_context=lambda: nullcontext(), + ) + chunks = [ + _ScoringChunk( + request_index=index, + input_ids=[index + 1, index + 2], + score_start=1, + score_count=1, + metadata={"_evalution_disable_loglikelihood_chunk_progress": True}, + ) + for index in range(3) + ] + + outputs = BaseTransformerSession._score_chunks(session, chunks, batch_size=2) + + assert len(outputs) == len(chunks) + assert all(output.is_greedy for output in outputs) + assert all(output.token_count == 1 for output in outputs) + assert outputs[0].logprob == outputs[1].logprob == outputs[2].logprob