fix: preserve completion usage and trim stop tokens - #312
Draft
Andrewxu313 wants to merge 1 commit into
Draft
Conversation
Use the original prompt length for response usage after sequence re-entry, and route completion detokenization through the shared EOS-trimming helper.
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.
Description
Preserve OpenAI prompt-token usage across sequence re-entry and trim model stop tokens consistently from completion text.
Motivation
Sequence re-entry mutates the internal prompt length while recomputing KV state, but response usage must retain the caller's original prompt. Two completion detokenization paths also bypassed the existing EOS-aware helper and could leak stop tokens.
Milestone (optional)
GLM-5.2 release closeout.
Type of Change
model— add/extend model support (models/**+ registration seam + model kernels only)kernel— add/optimize a compute kernel (batchgen_kernels/**+ in-tree kernel dirs)core— change scheduling/serving/runtime scaffolding (the only type that may)fix— narrow bug fix (+ a regression test)infra— build / CI / packaging / scripts / Dockerdocs— documentation onlyFile changes
batchgen/batchgen_worker.pytests/test_completion_usage_accounting.pyValidation
PR_TYPE=fix.Checklist
PR Merge Policy Contract — pre-merge checklist
tests/.