[bugfix] Stop/EOS inside a spec accept run beats the max_new_tokens finish (#33758)

Signed-off-by: Shiyan Deng <dsy842974287@meta.com>
Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>
Co-authored-by: Hanming Lu <hanminglu@meta.com>
This commit is contained in:
Shiyan Deng
2026-08-07 15:18:14 -07:00
committed by GitHub
co-authored by Lu Fang Hanming Lu Hanming Lu
parent 07297049e9
commit b2f9603f93
2 changed files with 231 additions and 8 deletions
+29 -8
View File
@@ -1572,7 +1572,9 @@ class Req(ReqDllmMixin):
def _check_vocab_boundary_finish(self, new_accepted_tokens: List[int] = None):
for i, token_id in enumerate(new_accepted_tokens):
if token_id >= self.vocab_size or token_id < 0:
if token_id < 0 or (
self.vocab_size is not None and token_id >= self.vocab_size
):
offset = len(self.output_ids) - len(new_accepted_tokens) + i
if self.sampling_params.stop_token_ids:
self.output_ids[offset] = next(
@@ -1586,6 +1588,18 @@ class Req(ReqDllmMixin):
return False
def _cap_finished_len_at_max_new_tokens(self) -> None:
"""Demote a stop matched beyond the length budget to a length finish.
Speculative decoding can accept a run that both crosses
``max_new_tokens`` and contains a stop; a stop located past the cap
must not extend the emitted output beyond the cap.
"""
max_new_tokens = self.sampling_params.max_new_tokens
if self.finished_len is not None and self.finished_len > max_new_tokens:
self.finished_reason = FINISH_LENGTH(length=max_new_tokens)
self.finished_len = max_new_tokens
def update_finish_state(self, new_accepted_len: int = 1):
if self.finished():
return
@@ -1595,26 +1609,33 @@ class Req(ReqDllmMixin):
self.to_finish = None
return
if len(self.output_ids) >= self.sampling_params.max_new_tokens:
self.finished_reason = FINISH_LENGTH(
length=self.sampling_params.max_new_tokens
)
self.finished_len = self.sampling_params.max_new_tokens
return
new_accepted_tokens = self.output_ids[-new_accepted_len:]
# Sanitize out-of-range / NaN token ids before any decode.
if self._check_vocab_boundary_finish(new_accepted_tokens):
self._cap_finished_len_at_max_new_tokens()
return
# Stop string beats EOS/stop-token matched in the same step (speculative
# decoding can accept >1 token): token-based would trim only the last
# token and leak the stop string.
if self._check_str_based_finish(new_accepted_len):
self._cap_finished_len_at_max_new_tokens()
return
# Stop token/EOS beats the length cap for the same reason: a spec accept
# run can cross max_new_tokens in the very step the EOS lands, and a
# length-first finish would keep the over-accepted tokens after the EOS
# (up to the cap) in the emitted output.
if self._check_token_based_finish(new_accepted_tokens):
self._cap_finished_len_at_max_new_tokens()
return
if len(self.output_ids) >= self.sampling_params.max_new_tokens:
self.finished_reason = FINISH_LENGTH(
length=self.sampling_params.max_new_tokens
)
self.finished_len = self.sampling_params.max_new_tokens
return
if self.grammar is not None and self.grammar.is_terminated():