[Hotfix] Fix performance regression introduced in #1261 (#1276)

#1261 added a logit cache, which was being concatenated to in a loop.

This PR never caches any logits except for the final one, but if
`include_all_uncached_tokens` is `True`, then `get_logits` will return
an array of shape `(len(tokens)-num_cached, vocab_size)`. Otherwise, it
will return an array of shape `(1, vocab_size)`

This "should" include all tokens being fast-forwarded (unless the
fast-forwarded tokens are somehow already in the kv, e.g. if a notebook
cell is re-run). If any are missing, I just call the probabilities `nan`
here, which could in principle cause some artifacts in the widget.

Note that an alternative implementation would have been to keep the full
logit cache, pre-allocating the buffer in advance with shape `(n_ctx,
vocab_size)`, but this feels really unnecessary. Happy to revisit after
we get this merged and the big performance issues are gone.

@nopdive would you mind re-running the repro? Note that you shouldn't
see a speedup for the prompt part (as this PR does not concern
pre-filling), but you should see a speedup for the rest of the
completion. Side note, but we can get rid of the (relatively small, mind
you) prompt slowdown by simply deciding to omit probabilities from
"input" tokens. Lmk if you want that, and I can open a follow-up PR.

Edit:
Intended to fix #1275

Author: hudson-ai

guidance/models/_engine/_engine.py

@@ -212,10 +212,9 @@ def __call__(
                 logits_lat_ms = 0.0
             else:
                 t1 = time.time()
-                logits = self.get_logits(token_ids=tokens, full_sequence=True)
+                logits = self.get_logits(token_ids=tokens, include_all_uncached_tokens=True)
                 logits_lat_ms = (time.time() - t1) * 1000
 
-
             # Important: don't wait on this future until after getting the logits;
             # this allows the mask to be built concurrently with model inference
             mask, ll_response = mask_fut.result()
@@ -229,19 +228,27 @@ def __call__(
 
             if logits is not None:
                 # Not the last one -- that's for the *next* token.
-                ff_logits = logits[max(len(logits)-len(ff_tokens)-1, 0):-1]
+                ff_logits = logits[-len(ff_tokens)-1:-1]
                 ff_probs = (
-                    softmax(np.array(ff_logits))
+                    softmax(ff_logits)
                     if last_temperature < _TEMPERATURE_EPSILON
-                    else softmax(np.array(ff_logits) / last_temperature)
+                    else softmax(ff_logits / last_temperature)
                 )
-                if ff_probs.shape[0] != len(ff_tokens):
-                    assert ff_probs.shape[0] == len(ff_tokens) - 1
+
+                if len(ff_tokens) == len(tokens) and ff_probs.shape[0] == len(ff_tokens) - 1:
                     # We didn't have a BOS token, so we need to fake the first token prob (say... 1?)
-                    ff_probs = np.concatenate(
-                        (np.ones((1, ff_probs.shape[1])), ff_probs), axis=0
+                    ff_probs = np.pad(
+                        ff_probs, [(1, 0), (0, 0)], mode='constant', constant_values=1.0
+                    )
+                elif ff_probs.shape[0] < len(ff_tokens):
+                    # Not enough logits were returned despite include_all_uncached_tokens=True, probably due to
+                    # using a mock model that doesn't bother to return logits for uncached tokens (all are uncached...)
+                    ff_probs = np.pad(
+                        ff_probs, [(len(ff_tokens) - ff_probs.shape[0], 0), (0, 0)],
+                        mode='constant', constant_values=np.nan
                     )
             else:
+                # really just for mypy -- we shouldn't need this
                 ff_probs = np.empty(shape=())
 
             gen_tokens = []
@@ -448,5 +455,14 @@ def get_top_k(_probs: NDArray, _k: int = 5) -> list[GenToken]:
         return output
 
     @abstractmethod
-    def get_logits(self, token_ids: list[int], full_sequence: bool = False) -> NDArray:
+    def get_logits(self, token_ids: list[int], include_all_uncached_tokens: bool = False) -> NDArray:
+        """
+        Get the logits for the given token ids.
+        If include_all_uncached_tokens is True:
+            logits for all uncached tokens will be returned, i.e.
+            the return value's shape will be `(len(tokens)-num_cached, vocab_size)`.
+        If include_all_uncached_tokens is False:
+            logits for the last token will be returned, i.e.
+            the return value's shape will be `(1, vocab_size)`.
+        """
         pass

guidance/models/_llama_cpp.py

@@ -148,7 +148,7 @@ def __init__(
                          compute_log_probs=compute_log_probs, enable_backtrack=enable_backtrack,
                          enable_ff_tokens=enable_ff_tokens, enable_monitoring=enable_monitoring, **kwargs)
 
-    def get_logits(self, token_ids, full_sequence=False):
+    def get_logits(self, token_ids: list[int], include_all_uncached_tokens: bool = False) -> np.ndarray:
         """Computes the logits for the given token state.
 
         This overrides a method from the LocalEngine class that is used to get
@@ -167,20 +167,17 @@ def get_logits(self, token_ids, full_sequence=False):
         # check what we have already cached
         num_cached = sum(takewhile(operator.truth, map(operator.eq, token_ids, self._cached_token_ids)))
         if num_cached == len(token_ids):
-            if full_sequence:
-                return self._cached_logits[:num_cached, :]
-            else:
-                return self._cached_logits[[num_cached - 1], :]
+            if num_cached == len(self._cached_token_ids):
+                # last token input is the same as the last cached token, so return the last cached logits
+                return self._cached_logits
+            # we need to pass at least one new token
+            num_cached = num_cached - 1
 
         # clear obsolete parts of kv cache
         llama_cpp.llama_kv_cache_seq_rm(self.model_obj.ctx, -1, num_cached, -1)
 
-        # clear obsolete parts of logit cache
-        if self._cached_logits is not None and len(self._cached_logits) > num_cached:
-            self._cached_logits = self._cached_logits[:num_cached, :]
-
         # eval the model
-        logits_for_each_batch = []
+        logits_for_each_batch: list[np.ndarray] = []
         n_batch = self.model_obj.n_batch
         batch = self._context.batch
         for i in range(num_cached, len(token_ids), n_batch):
@@ -199,31 +196,30 @@ def get_logits(self, token_ids, full_sequence=False):
 
             # get the logits for *this* batch
             llama_logits = llama_cpp.llama_get_logits(self.model_obj.ctx)
-            logits = np.ctypeslib.as_array(
+            logits_for_this_batch = np.ctypeslib.as_array(
                 llama_logits,
                 shape=(n_tokens, self._n_vocab),
             ).copy()
-            logits_for_each_batch.append(logits)
+            logits_for_each_batch.append(logits_for_this_batch)
 
         # update the metrics
         self.metrics.engine_output_tokens += 1
         self.metrics.engine_input_tokens += len(token_ids) - num_cached
 
         # save the results
         self._cached_token_ids = token_ids.copy()
+        self._cached_logits = logits_for_each_batch[-1][[-1], :]
 
-        if self._cached_logits is not None:
-            logits_for_each_batch = [self._cached_logits] + logits_for_each_batch
-
-        self._cached_logits = np.concatenate(
-            logits_for_each_batch,
-            axis=0
-        )
-
-        if full_sequence:
-            return self._cached_logits
+        if include_all_uncached_tokens:
+            logits = np.concatenate(
+                logits_for_each_batch,
+                axis=0
+            )
+            assert logits.shape[0] == len(token_ids) - num_cached
+            return logits
         else:
-            return self._cached_logits[[-1], :]
+            return self._cached_logits
+
 
 class LlamaCpp(Model):
     def __init__(

guidance/models/_mock.py

@@ -112,21 +112,7 @@ def get_next_token_with_top_k(
             logits, logits_lat_ms, token_ids, mask, temperature, k, force_return_unmasked_probs
         )
 
-    def get_logits(self, token_ids: list[int], full_sequence: bool = False) -> np.ndarray:
-        """Get the logits for the given token state."""
-        if not full_sequence:
-            return self._get_logits(token_ids).reshape(1, -1)
-        else:
-            # TODO: is it worth it to add a prefix cache here?
-            l0 = self._get_logits([token_ids[0]])
-            # if we are getting the full sequence then we need to compute the logits for all tokens
-            logits = np.zeros((len(token_ids), len(l0)))
-            logits[0] = l0
-            for i in range(1, len(token_ids)):
-                logits[i] = self._get_logits(token_ids[: i + 1])
-            return logits
-
-    def _get_logits(self, token_ids: list[int]) -> np.ndarray:
+    def get_logits(self, token_ids: list[int], include_all_uncached_tokens: bool = False) -> np.ndarray:
         """Pretends to compute the logits for the given token state."""
         if len(token_ids) == 0:
             raise ValueError("token_ids must not be empty")
@@ -154,7 +140,7 @@ def _get_logits(self, token_ids: list[int]) -> np.ndarray:
                         logits[i] += bias
                         bias /= 2  # if we have multiple matches then they apply with decreasing bias
 
-        return logits
+        return logits.reshape(1, -1)
 
     def _get_next_tokens(self, byte_string):
         special_tokens = [

guidance/models/_transformers.py

@@ -459,7 +459,7 @@ def _model(self, model, **kwargs) -> "PreTrainedModel":
             model = transformers_package.AutoModelForCausalLM.from_pretrained(model, **kwargs)
         return model
 
-    def get_logits(self, token_ids: list[int], full_sequence: bool = False):
+    def get_logits(self, token_ids: list[int], include_all_uncached_tokens: bool = False) -> np.ndarray:
         """Computes the logits for the given token state.
 
         This overrides a method from the LocalEngine class that is used to get
@@ -478,10 +478,11 @@ def get_logits(self, token_ids: list[int], full_sequence: bool = False):
         # check what we have already cached
         num_cached = sum(takewhile(operator.truth, map(operator.eq, token_ids, self._cached_token_ids)))
         if num_cached == len(token_ids):
-            if full_sequence:
-                return self._cached_logits[:num_cached, :]
-            else:
-                return self._cached_logits[[num_cached - 1], :]
+            if num_cached == len(self._cached_token_ids):
+                # last token input is the same as the last cached token, so return the last cached logits
+                return self._cached_logits
+            # we need to pass at least one new token
+            num_cached = num_cached - 1
 
         # check how many tokens are in the kv cache and what the max size of the cache is
         past_key_values = self._past_key_values
@@ -558,17 +559,16 @@ def get_logits(self, token_ids: list[int], full_sequence: bool = False):
                     else:
                         self._past_key_values = None
                     past_length = 0
-
-        # clear obsolete parts of logit cache
-        if self._cached_logits is not None and len(self._cached_logits) > num_cached:
-            self._cached_logits = self._cached_logits[:num_cached, :]
+        else:
+            # Should never happen, but just in case -- we're not using all of the cached tokens
+            # (because they're not really in the kv cache somehow)
+            num_cached = past_length
 
         # eval the model
-        assert past_length <= num_cached
-        assert num_cached <= len(token_ids)
+        assert past_length <= len(token_ids)
         new_token_ids = token_ids[past_length:]
         assert len(new_token_ids) > 0
-        logits_for_each_batch = []
+        logits_for_each_batch: list[np.ndarray] = []
         with torch.no_grad():
             # Not all models support batched tokens for some reason
             try:
@@ -628,19 +628,17 @@ def get_logits(self, token_ids: list[int], full_sequence: bool = False):
         # save the results
         self._past_key_values = model_out.past_key_values
         self._cached_token_ids = token_ids.copy()
+        self._cached_logits = logits_for_each_batch[-1][[-1], :]
 
-        if self._cached_logits is not None:
-            logits_for_each_batch = [self._cached_logits] + logits_for_each_batch
-
-        self._cached_logits = np.concatenate(
-            logits_for_each_batch,
-            axis=0
-        )
-
-        if full_sequence:
-            return self._cached_logits
+        if include_all_uncached_tokens:
+            logits = np.concatenate(
+                logits_for_each_batch,
+                axis=0
+            )
+            assert logits.shape[0] == len(token_ids) - num_cached
+            return logits
         else:
-            return self._cached_logits[[-1], :]
+            return self._cached_logits
 
 class Transformers(Model):
     def __init__(