fix: correctly size Llama.scores when draft_model implies logits_all

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- fix: size `Llama.scores` from `self._logits_all` instead of the raw `logits_all` argument, so passing `draft_model` without also explicitly passing `logits_all=True` no longer crashes generation past `n_batch` tokens with `ValueError: could not broadcast input array ... into shape (0,)` by @yadukpb
+
 ## [0.3.36]
 
 - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098

llama_cpp/llama.py

@@ -482,7 +482,7 @@ def free_lora_adapter():
         self._requires_eval = True
         self.input_ids: npt.NDArray[np.intc] = np.ndarray((n_ctx,), dtype=np.intc)
         self.scores: npt.NDArray[np.single] = np.ndarray(
-            (n_ctx if logits_all == True else n_batch, self._n_vocab), dtype=np.single
+            (n_ctx if self._logits_all else n_batch, self._n_vocab), dtype=np.single
         )
 
         self._mirostat_mu = ctypes.c_float(

tests/test_llama.py

@@ -267,6 +267,35 @@ def logit_processor_func(input_ids, logits):
     assert number_1 == number_3
 
 
+def test_draft_model_scores_buffer_sized_for_n_ctx(llama_cpp_transformer_model_path):
+    # Llama.__init__ sets self._logits_all = True whenever draft_model is
+    # given (see the `logits_all if draft_model is None else True` line),
+    # regardless of whether the caller also passed logits_all=True -- but
+    # self.scores used to be allocated from the raw logits_all argument
+    # instead of self._logits_all, so it came out n_batch-sized instead of
+    # n_ctx-sized. eval() writes into self.scores using self._logits_all,
+    # so generation crashed with "could not broadcast input array ...
+    # into shape (0,)" as soon as n_past exceeded n_batch.
+    from llama_cpp.llama_speculative import LlamaPromptLookupDecoding
+
+    llm = llama_cpp.Llama(
+        model_path=llama_cpp_transformer_model_path,
+        n_ctx=128,
+        n_batch=16,
+        draft_model=LlamaPromptLookupDecoding(max_ngram_size=2, num_pred_tokens=4),
+        verbose=False,
+        # Deliberately not passing logits_all=True: draft_model alone
+        # should be enough, per self._logits_all's own logic above.
+    )
+
+    assert llm.scores.shape[0] >= llm.n_batch
+
+    # More tokens than n_batch so eval() must cross that boundary at least
+    # once; this is what used to raise ValueError before the fix.
+    output = llm.create_completion("Once upon a time", max_tokens=30)
+    assert output["choices"][0]["text"]
+
+
 def test_real_llama_repeated_prompt_cache(llama_cpp_model_path):
     model = llama_cpp.Llama(
         llama_cpp_model_path,