Preserve speculative finish correctness across reasoning and stop strings
ReasonerGrammarObject wraps an inner grammar, but disaggregated prebuilt replay checks the wrapper current_token to avoid accepting an already accepted token twice. Track the token on the wrapper itself so reasoning grammar follows the same contract as XGrammar. Speculative decode can accept a stop string and EOS in one step. Check stop strings before token-based EOS after sanitizing invalid token ids, and set finished_len at the matched stop position so trailing accepted tokens do not leak. Constraint: Current branch predates upstream helper methods for locating string stop positions, so the stop-string fix is manually ported instead of cherry-picked. Rejected: Direct cherry-pick of bbc853df46 | current schedule_batch.py lacks the upstream helper context. Confidence: high Scope-risk: moderate Directive: Keep vocab-boundary sanitization before string decoding; do not move token-based EOS ahead of stop-string checks without a same-step speculative regression test. Tested: RED/GREEN remote pytest in cjy-glm5-new for constrained current_token and stop-string speculative tests Tested: git diff --check; py_compile for reasoner_grammar_backend.py and schedule_batch.py Not-tested: Full scheduler/disaggregation integration suite
This commit is contained in:
@@ -0,0 +1,61 @@
|
||||
import unittest
|
||||
from array import array
|
||||
|
||||
from sglang.srt.managers.schedule_batch import Req
|
||||
from sglang.srt.sampling.sampling_params import SamplingParams
|
||||
|
||||
|
||||
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)},
|
||||
}
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
eos_token_id = EOS_ID
|
||||
additional_stop_token_ids = None
|
||||
|
||||
def decode(self, ids):
|
||||
return "".join(ID_TO_TEXT[int(i)] for i in ids)
|
||||
|
||||
|
||||
class _NormalizeTokenizer:
|
||||
def encode(self, text, add_special_tokens=False):
|
||||
return list(range(len(text)))
|
||||
|
||||
|
||||
def _make_req(output_ids, *, stop=None, eos_token_ids=frozenset()):
|
||||
sampling_params = SamplingParams(max_new_tokens=1000, stop=stop)
|
||||
sampling_params.normalize(tokenizer=_NormalizeTokenizer())
|
||||
req = Req(
|
||||
rid="r",
|
||||
origin_input_text="",
|
||||
origin_input_ids=array("q", [0]),
|
||||
sampling_params=sampling_params,
|
||||
eos_token_ids=set(eos_token_ids),
|
||||
vocab_size=10_000,
|
||||
)
|
||||
req.tokenizer = _FakeTokenizer()
|
||||
req.output_ids = array("q", output_ids)
|
||||
return req
|
||||
|
||||
|
||||
class TestStopStrSpeculative(unittest.TestCase):
|
||||
def test_stop_str_wins_over_eos_in_same_spec_step(self):
|
||||
# A speculative step may accept both the stop string and EOS. The stop
|
||||
# string must finish first so finished_len trims at STOP instead of EOS.
|
||||
req = _make_req([10, 11, STOP_ID, EOS_ID], stop=["STOP"], eos_token_ids={EOS_ID})
|
||||
|
||||
req.check_finished(new_accepted_len=4)
|
||||
|
||||
self.assertTrue(req.finished())
|
||||
self.assertEqual(req.finished_reason.matched, "STOP")
|
||||
self.assertEqual(req.finished_len, 3)
|
||||
self.assertEqual(list(req.output_ids_through_stop), [10, 11, STOP_ID])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user