Repository navigation
Prompt-lookup decoding advances stateful logits processors (CFG, SynthID watermark) on draft tokens, so output differs from generation without prompt lookup #50051
Description
System Info
- transformers 5.21.0 (also reproduced on main @ 494470d)
- torch 2.11.0+cpu, Python 3.12, Linux (WSL2), CPU only
pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cpu && pip install transformers==5.21.0
Who can help?
Reproduction
With do_sample=False, prompt lookup should only change speed. It does for plain greedy decoding and for a stateless processor (repetition_penalty), but not when guidance_scale or a SynthID watermarking_config is set:
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, SynthIDTextWatermarkingConfig
model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
print("transformers", transformers.__version__)
# Repetitive prompt so that prompt lookup finds drafts.
inputs = tok("a b c d a b c d a b c d a b c d a b c", return_tensors="pt")
n = inputs.input_ids.shape[1]
watermark = SynthIDTextWatermarkingConfig(keys=[654, 400, 836, 123, 340, 443, 597, 160, 57, 29], ngram_len=5)
for name, kwargs in {
"plain greedy (control)": {},
"repetition_penalty=1.3 (stateless processor, control)": {"repetition_penalty": 1.3},
"guidance_scale=1.5": {"guidance_scale": 1.5},
"SynthID watermark": {"watermarking_config": watermark},
}.items():
a, b = (model.generate(**inputs, do_sample=False, max_new_tokens=40, prompt_lookup_num_tokens=k, **kwargs)[0, n:].tolist()
for k in (None, 3))
first = next((i for i, (x, y) in enumerate(zip(a + [None], b + [None])) if x != y), None)
print(f"{name}: identical={a == b} first_mismatch={first}")Relevant output (5.21.0; main gives the same four lines):
plain greedy (control): identical=True first_mismatch=None
repetition_penalty=1.3 (stateless processor, control): identical=True first_mismatch=None
guidance_scale=1.5: identical=False first_mismatch=1
SynthID watermark: identical=False first_mismatch=0
Expected behavior
Generation with prompt_lookup_num_tokens returns the same tokens as generation without it, as in the two controls.
Cause
Prompt lookup calls the same processor instances at two points where plain generation would not:
- Candidate filtering, before verification:
PromptLookupCandidateGeneratorcalls the logits-processor list on fake logits once per draft token (candidate_generator.py:1116). - Verification:
_assisted_decodingcalls it once per candidate position, including positions after the first rejected draft token (utils.py:4090).
UnbatchedClassifierFreeGuidanceLogitsProcessor keeps the unconditional context and its KV cache across calls (logits_process.py:2325-2356). SynthIDTextWatermarkLogitsProcessor keeps an n-gram context, a call counter and a context history (logits_process.py:2748-2781). Each extra call advances that state, and nothing restores it. Stage 1 runs before verification, so even the first verified position can see the wrong state. This fits the SynthID output differing at the first new token.
generate() already refuses assisted generation for stateful models (utils.py:1871); there is no equivalent check for stateful logits processors. do_sample=True goes through the same calls, but the script above only tests greedy decoding.
A fix would need to cover both stages, for example by refusing (or falling back to non-assisted decoding) when a stateful processor is active, or by keeping the processor state tied to the accepted prefix (no state updates from candidate filtering, and verification-time state rebuilt for the accepted tokens only). Happy to test a patch.
Activity
Duo-ee commentedon Oct 28, 2026
Independent reproduction on a stock CPU install, plus a few measurements that may help pin the fix down.
Environment: transformers 5.15.0, torch 2.13.0+cpu, Python 3.10.11. Same tiny model and prompt as in the report, and all four lines match:
So this is not a 5.21 regression — 5.15.0 shows it too.
Call counts, taken by wrapping the processor's
__call__for a 40-token greedy run withprompt_lookup_num_tokens=3:SynthIDTextWatermarkLogitsProcessor: 40 calls without prompt lookup, 46 with itUnbatchedClassifierFreeGuidanceLogitsProcessor: 40 without, 46 withThe
input_idslengths at which those calls enter are[20, 21, 22, 20, 21, 22, 23, 21, 22, 23, ...](prompt is 20 tokens including BOS). The length goes backwards: the candidate-filtering pass re-enters the processor with the un-extended sequence, which is the extra state advance. For the watermark processor the very first verified position already sees a shiftedstate.num_calls, consistent withfirst_mismatch=0.One detail worth stating explicitly for whoever writes the patch: the stage-1 call is unconditional today.
PromptLookupCandidateGeneratortakeslogits_processoras an optional argument (candidate_generator.py:1053on main 95bafdc), bututils.py:1318-1325always passes the fullLogitsProcessorList, and the filtering loop atcandidate_generator.py:1109-1124runs whenever it is non-empty. Dropping the argument is not reachable fromgenerate().On the two options in the report: the "refuse or fall back to greedy" one needs a way to mark a processor as stateful, since no such flag exists today (
LogitsProcessorsubclasses expose nothing like_is_stateful;generate()only checks the model, atutils.py:1871). The "restore state for the accepted prefix" one needs save/restore per stateful processor; for CFG that includes the unconditional KV cache (logits_process.py:2356), which is not cheap. Either way there is more plumbing than it looks, but I do not have a strong preference between the two.Pothan0 commentedon Oct 29, 2026
Confirmed on release 5.20.2 (torch 2.14.1+cpu, Python 3.10, CPU only) with your script: the two controls are identical, while guidance_scale=1.5 diverges at generated-token index 1 and the SynthID watermark at index 0 (first_mismatch over the continuation). I could check main only by source (it needs Python >=3.11): candidate_generator.py:1115 and utils.py:4090 are unchanged, and the stateful-model refusal at utils.py:1875 still has no processor-level equivalent.
Proposal for a first PR, scoped to what I reproduced: when prompt_lookup_num_tokens is active, refuse or fall back to plain decoding if the effective logits-processor list contains a known stateful built-in (UnbatchedClassifierFreeGuidanceLogitsProcessor, WatermarkLogitsProcessor, SynthIDTextWatermarkLogitsProcessor), mirroring the stateful-model check. The same two call sites also serve assistant-model decoding, which I have not tested. This is not a complete fix: a custom stateful processor passed via logits_processor= would still advance state.
Would you prefer this narrow guard, or the full fix tying processor state to the accepted prefix, which seems to need a statefulness contract on LogitsProcessor?
prasadovhal commentedon Oct 29, 2026
Thanks for the detailed investigation and reproductions.
I agree that the narrower guard proposed by @Pothan0 could be a reasonable first fix.
One small refinement: I'd prefer an explicit error for unsupported stateful processors rather than silently falling back to standard decoding. This would make it clear when prompt-lookup acceleration cannot be used.
We could keep the initial PR limited to known stateful built-in processors, with regression tests for CFG, SynthID, and unaffected stateless processors. The custom processor case and accepted-prefix state management could remain separate work.
Would the maintainers be comfortable with this scope?
If nobody is already working on the fix, I'd be happy to take it up and open a focused PR.