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,69 @@
|
||||
import unittest
|
||||
|
||||
from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
|
||||
|
||||
|
||||
class _InnerGrammar:
|
||||
def __init__(self):
|
||||
self.accepted_tokens = []
|
||||
|
||||
def accept_token(self, token: int):
|
||||
self.accepted_tokens.append(token)
|
||||
|
||||
def is_terminated(self):
|
||||
return False
|
||||
|
||||
def rollback(self, k: int):
|
||||
self.accepted_tokens = self.accepted_tokens[:-k]
|
||||
|
||||
def allocate_vocab_mask(self, vocab_size: int, batch_size: int, device):
|
||||
return None
|
||||
|
||||
def fill_vocab_mask(self, vocab_mask, idx: int):
|
||||
return None
|
||||
|
||||
def move_vocab_mask(self, vocab_mask, device):
|
||||
return vocab_mask
|
||||
|
||||
@property
|
||||
def apply_vocab_mask(self):
|
||||
return None
|
||||
|
||||
def copy(self):
|
||||
ret = _InnerGrammar()
|
||||
ret.accepted_tokens = list(self.accepted_tokens)
|
||||
return ret
|
||||
|
||||
@property
|
||||
def finished(self):
|
||||
return False
|
||||
|
||||
|
||||
class TestReasonerGrammarCurrentToken(unittest.TestCase):
|
||||
def test_wrapper_tracks_current_token_for_disagg_prebuilt_dedup(self):
|
||||
inner = _InnerGrammar()
|
||||
grammar = ReasonerGrammarObject(inner, think_end_id=7)
|
||||
grammar.maybe_init_reasoning(True)
|
||||
|
||||
grammar.accept_token(10) # still in thinking; inner grammar is untouched
|
||||
self.assertEqual(grammar.current_token, 10)
|
||||
self.assertEqual(inner.accepted_tokens, [])
|
||||
|
||||
grammar.accept_token(7) # exits thinking
|
||||
self.assertEqual(grammar.current_token, 7)
|
||||
self.assertEqual(inner.accepted_tokens, [])
|
||||
|
||||
grammar.accept_token(58) # generation token accepted by inner grammar
|
||||
self.assertEqual(grammar.current_token, 58)
|
||||
self.assertEqual(inner.accepted_tokens, [58])
|
||||
|
||||
# Mirrors disaggregation/decode_schedule_batch_mixin.py:process_prebuilt.
|
||||
# Once current_token is tracked on the wrapper, a re-prebuilt request
|
||||
# skips re-accepting the already accepted token.
|
||||
if grammar.current_token is None:
|
||||
grammar.accept_token(58)
|
||||
self.assertEqual(inner.accepted_tokens, [58])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user