Conversation
Lookahead could return one token sequence while streaming another: accepted block tokens were omitted from the stream, history repeated the last accepted token, and EOS/custom stopping inside a block was checked too late. Determine the accepted prefix before KV-cache commitment and use it consistently for streaming, history, masks and return values in greedy and sampling paths. Also fix augment_all() compatibility (descriptor-preserving method copy, Transformers constructors, repeated augmentation, disabled-sampling fallback, SDPA config). Relates to issues hao-ai-lab#62, hao-ai-lab#54 and hao-ai-lab#59; does not touch the acceptance-count algorithm of PR hao-ai-lab#68. Developed with AI assistance.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Lookahead can return one token sequence while streaming another: accepted block tokens are omitted, history repeats the last accepted token, and EOS/custom stopping inside a block is checked too late. This change determines the accepted prefix before KV-cache commitment and uses that prefix consistently for streaming, history, masks and return values in greedy and sampling paths.
Public
augment_all()also required compatibility fixes to reach these paths: copy explicitly defined methods while preserving descriptors, retain Transformers constructors, make repeated augmentation safe, and preserve disabled sampling fallback. Default SDPA configuration now avoids an unsupported argument and transplanted zero-argumentsuper(); legacy tuple-cache Lookahead falls back to manual attention, with no SDPA speed claim.Validation on Torch2.7.1+cu128/Transformers4.36.2: 19 CPU tests and two GPU backend cases pass (each backend checks three random seeds, AR and Lookahead against the pre-injection HF reference). Tests exercise real tiny Llama forward passes, multi-token acceptance, multiple EOS IDs, internal stopping, streaming, history and cache lengths. The unmodified public augmentation baseline fails construction; a separately documented baseline with only compatibility fixes fails all11 decoding regressions. The new SDPA test also failed before its fix.
History behavior was reported in issue#62; disabled sampling fallback was reported in issue#54; the SDPA
lookahead-kwarg TypeError was reported in issue#59 and is resolved by the SDPA configuration fix here. This does not modify the acceptance-count algorithm addressed by PR#68. FlashAttention2, multiple GPUs, newer Transformers and return-dict generation are not covered; the upstream return-dict restriction remains. FlashAttention2's ordinary HF Cache path still has a tuple-only assumption and is not claimed compatible.Implementation and tests were prepared with AI coding assistance.