fix(schedule_batch): trim stop string when EOS matches in the same step (#28802)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
7c23d2255a
commit
bbc853df46
@@ -1423,15 +1423,19 @@ class Req(ReqDllmMixin):
|
||||
|
||||
new_accepted_tokens = self.output_ids[-new_accepted_len:]
|
||||
|
||||
if self._check_token_based_finish(new_accepted_tokens):
|
||||
return
|
||||
|
||||
# Sanitize out-of-range / NaN token ids before any decode.
|
||||
if self._check_vocab_boundary_finish(new_accepted_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):
|
||||
return
|
||||
|
||||
if self._check_token_based_finish(new_accepted_tokens):
|
||||
return
|
||||
|
||||
def reset_for_retract(self):
|
||||
# Increment retraction count before resetting other state. We should not reset this
|
||||
# since we are tracking the total number of retractions for each request.
|
||||
|
||||
@@ -17,8 +17,10 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
# Token id -> decoded text; decode() concatenates. Distinct symbols so no
|
||||
# accidental cross-matches (10-39 are lowercase letters).
|
||||
STOP_ID = 1
|
||||
EOS_ID = 2
|
||||
ID_TO_TEXT = {
|
||||
STOP_ID: "STOP",
|
||||
EOS_ID: "",
|
||||
**{i: chr(ord("a") + i % 26) for i in range(10, 40)},
|
||||
60: "a",
|
||||
61: ".",
|
||||
@@ -48,7 +50,7 @@ class _MockTokenizerForNormalize:
|
||||
return list(range(len(s))) # One "token" per character
|
||||
|
||||
|
||||
def _make_req(output_ids, stop=None, stop_regex=None):
|
||||
def _make_req(output_ids, stop=None, stop_regex=None, eos_token_ids=frozenset()):
|
||||
sp = SamplingParams(max_new_tokens=1000, stop=stop, stop_regex=stop_regex)
|
||||
sp.normalize(tokenizer=_MockTokenizerForNormalize()) # char-based stop_str_max_len
|
||||
req = Req(
|
||||
@@ -56,7 +58,7 @@ def _make_req(output_ids, stop=None, stop_regex=None):
|
||||
origin_input_text="",
|
||||
origin_input_ids=array("q", [0]),
|
||||
sampling_params=sp,
|
||||
eos_token_ids=set(),
|
||||
eos_token_ids=eos_token_ids,
|
||||
vocab_size=10_000,
|
||||
)
|
||||
req.tokenizer = _FakeTokenizer()
|
||||
@@ -80,6 +82,18 @@ class TestStopStrSpeculative(unittest.TestCase):
|
||||
self.assertEqual(req.finished_reason.matched, "STOP")
|
||||
self.assertEqual(req.finished_len, 4)
|
||||
|
||||
def test_stop_str_wins_over_eos_in_same_step(self):
|
||||
# A step accepting both the stop string and EOS must finish as
|
||||
# FINISH_MATCHED_STR with finished_len past the stop string, not as
|
||||
# FINISH_MATCHED_TOKEN trimming only the EOS token (which leaks "STOP").
|
||||
req = _make_req(
|
||||
[10, 11, STOP_ID, EOS_ID], stop=["STOP"], eos_token_ids={EOS_ID}
|
||||
)
|
||||
req.update_finish_state(new_accepted_len=4)
|
||||
self.assertTrue(req.finished())
|
||||
self.assertEqual(req.finished_reason.matched, "STOP")
|
||||
self.assertEqual(req.finished_len, 3)
|
||||
|
||||
def test_stop_str_at_chunk_end_uses_full_len(self):
|
||||
# Stop is the last token -> loop never matches before the full window ->
|
||||
# fallback returns len(output_ids).
|
||||
|
||||
Reference in New Issue
Block a user