From 3af4316d540753748d1cf9bd880563313c3093af Mon Sep 17 00:00:00 2001
From: yadukpb <yadukrishnapb1@gmail.com>
Date: Tue, 13 Oct 2026 14:16:38 +0530
Subject: [PATCH] fix: correctly size Llama.scores when draft_model implies
 logits_all

Llama.__init__ already sets self._logits_all = True whenever draft_model
is given, regardless of the raw logits_all constructor argument -- but
self.scores was allocated from that raw argument directly instead of
self._logits_all. That mismatch leaves self.scores sized n_batch instead
of n_ctx whenever a caller passes draft_model=... without also explicitly
passing logits_all=True.

eval() writes into self.scores based on self._logits_all (not the raw
argument), so the undersized buffer crashes generation as soon as n_past
exceeds n_batch:

    ValueError: could not broadcast input array from shape (32000,) into shape (0,)

Fix: size self.scores from self._logits_all, consistent with how eval()
already decides whether to write into it at all.

Found while building a small CPU-only voice agent
(https://github.com/yadukpb/potato-agent) that uses
LlamaPromptLookupDecoding for speculative decoding -- hit this on every
turn past ~512 tokens.

#657 reports a superficially similar "could not broadcast ... into shape
(0,)" message, but that one is input_ids overflowing past n_ctx itself
(genuine context-window exhaustion), a different root cause, so this
isn't folded into that thread. #2110 is also draft_model-related but a
separate (performance) complaint.

Tested:
- git submodule update --init --recursive, make deps (builds the C++
  extension from this checkout with the fix applied), make lint -- clean.
- Added test_draft_model_scores_buffer_sized_for_n_ctx in tests/test_llama.py;
  verified it fails with the exact ValueError above against the pre-fix
  code and passes against the fix.
- Ran tests/test_llama.py, tests/test_llama_speculative.py,
  tests/test_llama_chat_format.py, and tests/test_llama_grammar.py (27
  tests total) against the from-source build: all pass, no regressions.
- CPU only, Linux x86_64, no special build flags.
---
 CHANGELOG.md        |  2 ++
 llama_cpp/llama.py  |  2 +-
 tests/test_llama.py | 29 +++++++++++++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5c34e2729..40adc4148 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
index 14e2f8500..ff86ed9a6 100644
--- a/llama_cpp/llama.py
+++ b/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(
diff --git a/tests/test_llama.py b/tests/test_llama.py
index c1c16e30c..c8adff920 100644
--- a/tests/test_llama.py
+++ b/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,