Skip to content

fix accepted-prefix consistency for streaming, history and stopping - #69

Open
Gxj230958 wants to merge 1 commit into
hao-ai-lab:mainfrom
Gxj230958:codex/consistent-block-commit
Open

Gxj230958 wants to merge 1 commit into
hao-ai-lab:mainfrom
Gxj230958:codex/consistent-block-commit

Conversation

@Gxj230958

Copy link
Copy Markdown

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-argument super(); 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.

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

1 participant