[dexhunter]

TL;DR

• Reduce work in LlamaPromptLookupDecoding.find_candidate_pred_tokens during prompt lookup decoding.
• Compare shifted token slices, skip impossible matches, and select the earliest continuation.
• Lookup benchmark: 89.02 ms → 25.83 ms per 4,320 calls, 71.0% less time.
• pytest tests/test_llama_speculative.py tests/test_llama_chat_format.py tests/test_llama_grammar.py plus native version/tokenization tests: eight passed.

What & Why

Each suffix length currently builds sliding windows and compares every token in every window.
It also collects all matching positions before returning one result.

The new path compares shifted slices for integer token arrays and skips suffix lengths without an earlier matching first token.
Integer suffixes match themselves, so selecting the first full match keeps the same result and slice.

Validation

Measured on Linux x86_64, AMD EPYC 9R14, Python 3.12.13, and NumPy 2.5.4, with CPU affinity and one BLAS thread.
The unchanged native library used a Release CPU build with OpenMP and BLAS disabled.
These measurements cover lookup time, not end-to-end model generation.

Results

Situation	Before	After	Change
32 tokens, repeated seven-token pattern	11.83 µs	4.10 µs	65.4% less time
64 unique tokens, no continuation	24.43 µs	4.26 µs	82.6% less time
128 repeated tokens, NumPy integer continuation count	14.44 µs	14.34 µs	No measurable change; 0.7% difference

All three rows use max_ngram_size=2 and ten requested tokens.
The last row passes np.int64(10) and exercises the existing matching path.
Times are per lookup, calculated from 5,120-call batches with nine timing comparisons.
The third row's difference is within measurement noise.

To compare locally, run this command on both revisions:

python -m timeit -r 9 -n 5120 \
  -s 'import numpy as np; from llama_cpp.llama_speculative import LlamaPromptLookupDecoding' \
  -s 'ids = np.resize(np.arange(7, dtype=np.intc), 32); count = 10' \
  'LlamaPromptLookupDecoding(2, count)(ids)'

For the second row, use ids = np.arange(64, dtype=np.intc).
For the last row, resize the seven-token pattern to 128 elements and set count = np.int64(10).

Native checks on the patched source:

python -m pytest tests/test_llama_speculative.py tests/test_llama_chat_format.py tests/test_llama_grammar.py tests/test_llama.py::test_llama_cpp_version tests/test_llama.py::test_llama_cpp_tokenization -q
make lint

Ruff lint and formatting passed across all 29 Python and test files.
The four-file suite including all of tests/test_llama.py reached its ten-minute limit during test_real_llama, after three passes.
The baseline validation had the same model-inference limitation; a full-suite pass is not claimed.

Why this is safe and easy to maintain

The change stays within one method and keeps one shared continuation loop.
The fallback preserves existing behavior for other inputs and unusual argument types.
The lookup tests and seeded checks cover values, slice offsets, dtypes, strides, read-only inputs, and errors.

AI assistance

Codex gpt-6.1-sol generated candidates tracked with Weco Observe.
Codex performed the source selection, validation, and PR preparation; no human review is claimed.

Weco share link: https://dashboard.weco.ai/share/YGQjNo2S_3mnwHD2B150wuOXpqR4MQhL