[Arham-Begani]

PR description

Fixes #306.

The first predicted token was appended before the generation loop, but the EOS check ran only after subsequent predictions. An immediate EOS could therefore trigger extra model calls and append ordinary text after the stop token. Removing special tokens during decoding does not remove that extra text.

Move the existing EOS check to the beginning of each additional decoding iteration. This handles the first and later predictions uniformly and preserves the existing EOS configuration, position IDs and token-limit behavior.

Local regression evidence (tests omitted from PR at maintainer request)

• Scalar and list-valued EOS settings, including an immediate EOS.
• EOS after a normal token, an exhausted token budget, no configured EOS and a one-token budget.
• Assert the generated token IDs, number of model calls, position IDs and cache reuse.
• Run the complete pipeline with a tiny randomly initialized Llama and an in-memory tokenizer; equal logits force immediate EOS. No model weights or tokenizer assets are downloaded.

The eight regression checks were preserved locally after SimJeg requested removing the added tests from the PR. Only kvpress/pipeline.py remains in the submitted diff, with the production change identical to the contributor-reviewed patch.

On unchanged main (7136c224105ba7b046a6e538742c3589698680b8), the new tests produce 3 failures and 5 passes. After the fix, all 8 pass. The scripted immediate-EOS case drops from four model calls to one; the real pipeline drops from two to one. These are deterministic call-count checks, not GPU throughput claims.

Validation

• Windows, Python 3.12.10, PyTorch 2.6.0+cpu, Transformers 5.2.0, pytest 7.4.4.
• Before the requested test-file removal: python -m pytest tests/test_generation_stopping.py -q: 8 passed. These results validate the identical production code retained in the current PR.
• Black and isort checks passed for the production file and the original local regression file.
• Repository-wide flake8 passes.
• python -m mypy . --check-untyped-defs: no issues in 85 source files.
• SPDX header scan: all 117 Python files have the required header; git diff --check passes.
• The Makefile requires Bash/make, unavailable in this Windows shell. Its flake8, mypy and SPDX checks were run directly and pass.
• Full make test, FlashAttention and GPU benchmarks were not run. The repository's agent instructions request testing new/modified tests; this change does not modify a press or a compression algorithm.
• Repository-wide Black/isort checks find existing formatting problems in untouched files; no unrelated formatting is included.

Related #301 changes EOS fallback configuration. This fix addresses the first-token stopping boundary and preserves the current configuration behavior.

Checklist

• Tests are working (make test) — targeted CPU regression suite passes; full GPU suite not run
• Code is formatted correctly (make style, on errors try make format) — equivalent style checks pass as detailed above; changed files pass Black/isort
• Copyright header is included
• All commits are signed-off using git commit -s — contributor personally authorized the DCO sign-off for the reviewed patch

This change adds no new press; the new-press checklist is not applicable.

AI assistance

An automated AI coding agent investigated the bug, wrote the implementation and tests, and ran the validation on behalf of Arham-Begani. The contributor supplied personal DCO authorization for the reviewed patch. The revised commit retains that sign-off on the identical production change and omits only the added test file as requested by the maintainer; it adds no new code beyond the certified patch.

🤖🤖🤖

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Arham-Begani]

This draft is ready for the repository's runner-vetting step at commit a173f14. The automated DCO check passed. The focused CPU suite passes all 8 cases after reproducing 3 failures on unchanged main; flake8, mypy and SPDX checks also pass, as detailed in the PR description.

An authorized NVIDIA vetter will need to approve CI before the upstream test workflow can run. Please vet this commit when appropriate. I am keeping the PR in draft pending that validation; the full GPU suite has not been run locally.

🤖🤖🤖

[SimJeg]

Hi thanks for the PR, please remove the added tests

[Arham-Begani]

Removed the added test file as requested in fa2161b. The PR now changes only kvpress/pipeline.py; its production code is identical to the previously submitted fix.

The eight regression checks are retained locally and were rerun after removing the file from the branch: all 8 pass, including the tiny CPU Llama pipeline check. The PR description now clearly identifies these as local evidence rather than tests included in the diff. Full NVIDIA GPU CI still requires authorized runner vetting for this head, so the PR remains draft.

This revision and reply were prepared with automated AI assistance.

🤖🤖🤖

[SimJeg]

/ok to test fa2161b

[SimJeg]

Looks good to me !