Faster loglikelihood scoring in mlx_lm.evaluate

mlx_lm/evaluate.py

@@ -23,7 +23,7 @@
 from tqdm import tqdm
 
 from .generate import batch_generate
-from .models.cache import make_prompt_cache
+from .models.cache import can_trim_prompt_cache, make_prompt_cache, trim_prompt_cache
 from .sample_utils import make_sampler
 from .utils import load
 
@@ -96,13 +96,25 @@ def __init__(
             self.use_chat_template = self.tokenizer.chat_template is not None
         self._sampler = sampler
 
-    def _process_prompt(self, prompt, step_size: int = 2048):
+    def _process_prompt(self, prompt, step_size: int = 2048, min_split: int = 512):
         prompt = mx.array(prompt)[None]
         cache = make_prompt_cache(self._model)
-        for i in range(0, prompt.shape[1], step_size):
-            logits = self._model(prompt[:, i : i + step_size], cache=cache)
+        length = prompt.shape[1]
+        # When the prompt's final chunk is long, process the final token by
+        # itself. The preceding tokens then only evaluate the cache state so
+        # their logits are never computed, and the vocabulary projection only
+        # runs for the one position whose log probabilities are needed. When
+        # the final chunk is shorter than ``min_split``, the projection costs
+        # less than the extra model call, so the chunk is processed whole.
+        last_len = (
+            1 if (length - 1) % step_size >= min_split else (length - 1) % step_size + 1
+        )
+        rest = prompt[:, : length - last_len]
+        for i in range(0, rest.shape[1], step_size):
+            self._model(rest[:, i : i + step_size], cache=cache)
             mx.eval([c.state for c in cache])
             mx.clear_cache()
+        logits = self._model(prompt[:, length - last_len :], cache=cache)
         logprobs = nn.log_softmax(logits[:, -1, :].astype(mx.float32))
         return logprobs, cache
 
@@ -208,14 +220,18 @@ def loglikelihood(self, requests) -> list[tuple[float, bool]]:
             # If the entire prompt got truncated ignore the question
             if prefix_l == 0:
                 long_completions += 1
-                all_scores.extend([-float("inf")] * len(rs))
-                all_is_greedy.extend([False] * len(rs))
+                scores.extend([-float("inf")] * len(rs))
+                is_greedy.extend([False] * len(rs))
                 continue
 
             # model scoring, returns num_requests x (logp, is_greedy, length).
             logprobs, cache = self._process_prompt(prefix)
             max_idx = mx.argmax(logprobs).item()
 
+            # When the cache is trimmable, score every continuation with the
+            # same cache and rewind it in-between instead of copying it for
+            # each continuation.
+            reuse_cache = can_trim_prompt_cache(cache)
             for s in full_sequences:
                 inputs = s[len(prefix) :]
                 # The logprobs from the last token of the prompt are
@@ -226,8 +242,11 @@ def loglikelihood(self, requests) -> list[tuple[float, bool]]:
                 if len(inputs) == 1:
                     continue
                 score, _, ig = self._score_fn(
-                    mx.array(inputs)[None, :], cache=copy.deepcopy(cache)
+                    mx.array(inputs)[None, :],
+                    cache=cache if reuse_cache else copy.deepcopy(cache),
                 )
+                if reuse_cache:
+                    trim_prompt_cache(cache, len(inputs) - 1)
                 scores[-1] += mx.sum(score).item()
                 is_greedy[-1] &= mx.all(ig).item()
 

tests/test_evaluate.py

@@ -4,6 +4,7 @@
 from unittest.mock import MagicMock, patch
 
 import mlx.core as mx
+import mlx.nn as nn
 
 from mlx_lm.evaluate import MLXLM
 
@@ -55,5 +56,71 @@ def mock_score_fn(batch):
         self.assertEqual(len(call_args_list[2][0][0]), 1)  # Third batch: 1 item
 
 
+class TestMLXLMLoglikelihood(unittest.TestCase):
+    """End-to-end tests for loglikelihood scoring with a real model."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.lm = MLXLM("mlx-community/Qwen1.5-0.5B-Chat-4bit")
+        base = (
+            "The city council met on Tuesday to discuss the new transit "
+            "plan, which includes additional bus routes and longer service "
+            "hours for the northern districts. "
+        )
+        ids = cls.lm.tokenizer.encode(base * 60, add_special_tokens=False)[:700]
+        cls.long_context = cls.lm.tokenizer.decode(ids)
+
+    def _reference(self, context, continuation):
+        """Score a continuation with a single forward pass over the full
+        sequence."""
+        prefix = self.lm._tokenize([context])[0]
+        full = self.lm._tokenize([context + continuation])[0]
+        logits = self.lm._model(mx.array(full[:-1])[None])
+        logprobs = nn.log_softmax(logits[0].astype(mx.float32), axis=-1)
+        score = 0.0
+        greedy = True
+        for pos in range(len(prefix) - 1, len(full) - 1):
+            target = full[pos + 1]
+            score += logprobs[pos, target].item()
+            greedy &= mx.argmax(logprobs[pos]).item() == target
+        return score, greedy
+
+    def test_loglikelihood_matches_full_forward(self):
+        cases = [
+            ("The capital of France is", " Paris"),
+            ("The capital of France is", " the city of Paris, which is known"),
+            ("The capital of France is", " Berlin"),
+            # A long context exercises the split prefill path
+            (self.long_context, " The council approved the plan."),
+        ]
+        requests = [MagicMock(args=case) for case in cases]
+        results = self.lm.loglikelihood(requests)
+        self.assertEqual(len(results), len(cases))
+        for (score, greedy), case in zip(results, cases):
+            ref_score, ref_greedy = self._reference(*case)
+            self.assertLess(abs(score - ref_score), 2e-1)
+            self.assertEqual(greedy, ref_greedy)
+
+    def test_loglikelihood_continuation_order_invariance(self):
+        # Continuations of different lengths after a common long prefix.
+        # Scoring them in either order must give the same results, which
+        # checks that scoring one continuation does not contaminate the
+        # cached prefix state used by the others.
+        context = self.long_context
+        continuations = [
+            " the lazy dog. " * 20,
+            " The fence.",
+            " A log lies on the other side of the river.",
+        ]
+        requests = [MagicMock(args=(context, c)) for c in continuations]
+        forward = self.lm.loglikelihood(requests)
+        backward = self.lm.loglikelihood(list(reversed(requests)))
+        for (score_f, greedy_f), (score_b, greedy_b) in zip(
+            forward, reversed(backward)
+        ):
+            self.assertAlmostEqual(score_f, score_b, places=3)
+            self.assertEqual(greedy_f, greedy_b)
+
+
 if __name__ == "__main__":
     unittest.main()