Skip to content

Fix first token handling in speculative decoding - #2308

Open
Excelius-Wang wants to merge 1 commit into
Lightning-AI:mainfrom
Excelius-Wang:fix/speculative-first-token
Open

Fix first token handling in speculative decoding#2308
Excelius-Wang wants to merge 1 commit into
Lightning-AI:mainfrom
Excelius-Wang:fix/speculative-first-token

Conversation

@Excelius-Wang

Copy link
Copy Markdown

What changed

  • return the target model's prefill sample as the first completion token
  • stop immediately when that token matches a stop token
  • keep the generated sequence within max_returned_tokens when a speculative round returns k + 1 tokens
  • add CPU regression coverage for the prefill, stop-token, and output-limit paths

Why

The target model samples the first completion token during prefill, but generate() previously
initialized the output accumulator as empty and only appended tokens from later speculative
decoding rounds. This dropped the first token from the generated text and could return one fewer
token than requested. It also skipped stop-token handling for that first sample.

Fixes #2307.

Validation

  • pytest -q tests/test_generate_speculatively.py: 19 passed, 1 skipped (GPU-only)
  • ruff check litgpt/generate/speculative_decoding.py tests/test_generate_speculatively.py
  • ruff format --check litgpt/generate/speculative_decoding.py tests/test_generate_speculatively.py
  • all applicable pre-commit hooks passed for the two changed Python files; the Prettier hook was skipped because it only targets JSON, YAML, and TOML files

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.

Speculative decoding drops the first generated token

1 participant