Skip to content

fix: correctly size Llama.scores when draft_model implies logits_all - #2384

Open
yadukpb wants to merge 1 commit into
abetlen:mainfrom
yadukpb:fix/draft-model-logits-all-scores-buffer
Open

yadukpb wants to merge 1 commit into
abetlen:mainfrom
yadukpb:fix/draft-model-logits-all-scores-buffer

Conversation

@yadukpb

@yadukpb yadukpb commented Oct 13, 2026

Copy link
Copy Markdown

What changed

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. One-line fix: size self.scores from self._logits_all, consistent with how eval() already uses self._logits_all to decide whether to write into it at all.

Why it matters

Passing draft_model=... without also explicitly passing logits_all=True leaves self.scores sized n_batch instead of n_ctx. eval() writes into it based on self._logits_all, so generation crashes as soon as n_past exceeds n_batch:

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

Found while building a small CPU-only voice agent (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 on closer read that one is input_ids overflowing past n_ctx itself (genuine context-window exhaustion) — a different root cause, so I didn't fold this into that thread. #2110 is also draft_model-related but a separate (performance) complaint.

How 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: constructs a Llama with draft_model=LlamaPromptLookupDecoding(...) and no explicit logits_all, generates past n_batch tokens. 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.

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.

abetlen#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. abetlen#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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants