From 8eac618a8d5200f590344bb11b21d2b7f99658eb Mon Sep 17 00:00:00 2001 From: Alex Nails Date: Fri, 10 Apr 2026 21:26:04 -0700 Subject: [PATCH] [tokenizer] lazy text accumulation + use deltas directly for streaming (#22548) Co-authored-by: Claude Opus 4.6 (1M context) --- .../sglang/srt/managers/tokenizer_manager.py | 65 +++++------------- test/manual/test_tokenizer_manager.py | 66 +++++-------------- 2 files changed, 34 insertions(+), 97 deletions(-) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index a3f124e67..e3eb3d4ab 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -147,22 +147,20 @@ class ReqState: # For streaming output last_output_offset: int = 0 - last_text_offset: int = 0 - # Buffer non-streaming text until the final response. - buffer_text: bool = False + # Accumulate text lazily so incremental streaming can emit the incoming + # delta directly without rebuilding the full output prefix. text: str = "" text_chunks: List[str] = dataclasses.field(default_factory=list) def append_text(self, chunk: str): - if self.buffer_text: + if chunk: self.text_chunks.append(chunk) - else: - self.text += chunk def get_text(self) -> str: - if self.buffer_text: - return "".join(self.text_chunks) + if self.text_chunks: + self.text += "".join(self.text_chunks) + self.text_chunks.clear() return self.text def get_crash_dump_output(self) -> Dict[Any, Any]: @@ -198,24 +196,6 @@ class ReqState: output_token_ids_logprobs: List[Any] = dataclasses.field(default_factory=list) -def make_req_state( - out_list: List[Dict[Any, Any]], - finished: bool, - event: asyncio.Event, - obj: Union[GenerateReqInput, EmbeddingReqInput], - time_stats: APIServerReqTimeStats, -) -> ReqState: - is_streaming_request = getattr(obj, "stream", False) - return ReqState( - out_list, - finished, - event, - obj, - time_stats, - buffer_text=not is_streaming_request, - ) - - def _slice_streaming_output_meta_info( meta_info: Dict[Any, Any], last_output_offset: int, @@ -1717,18 +1697,18 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): incremental = ( self.server_args.incremental_streaming_output and is_stream ) + delta_text = recv_obj.output_strs[i] + delta_output_ids = recv_obj.output_ids[i] output_offset = state.last_output_offset - state.append_text(recv_obj.output_strs[i]) - state.output_ids.extend(recv_obj.output_ids[i]) + state.append_text(delta_text) + state.output_ids.extend(delta_output_ids) if is_stream: if incremental: - output_token_ids = state.output_ids[output_offset:] + output_token_ids = delta_output_ids _slice_streaming_output_meta_info(meta_info, output_offset) state.last_output_offset = len(state.output_ids) - text = state.get_text() - output_text = text[state.last_text_offset :] - state.last_text_offset = len(text) + output_text = delta_text else: output_token_ids = state.output_ids.copy() output_text = state.get_text() @@ -1750,12 +1730,13 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): incremental = ( self.server_args.incremental_streaming_output and is_stream ) + delta_output_ids = recv_obj.output_ids[i] output_offset = state.last_output_offset - state.output_ids.extend(recv_obj.output_ids[i]) + state.output_ids.extend(delta_output_ids) if is_stream: if incremental: - output_token_ids = state.output_ids[output_offset:] + output_token_ids = delta_output_ids _slice_streaming_output_meta_info(meta_info, output_offset) state.last_output_offset = len(state.output_ids) else: @@ -2463,13 +2444,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): if not hasattr(obj, "is_single") or obj.is_single: time_stats = APIServerReqTimeStats(disagg_mode=self.disaggregation_mode) - state = make_req_state( - [], - False, - asyncio.Event(), - obj, - time_stats, - ) + state = ReqState([], False, asyncio.Event(), obj, time_stats) self.rid_to_state[obj.rid] = state if self.server_args.enable_trace: @@ -2485,13 +2460,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): else: for i in range(len(obj.rid)): time_stats = APIServerReqTimeStats(disagg_mode=self.disaggregation_mode) - state = make_req_state( - [], - False, - asyncio.Event(), - obj[i], - time_stats, - ) + state = ReqState([], False, asyncio.Event(), obj[i], time_stats) self.rid_to_state[obj.rid[i]] = state if self.server_args.enable_trace: diff --git a/test/manual/test_tokenizer_manager.py b/test/manual/test_tokenizer_manager.py index b3b646b97..d0febc75e 100644 --- a/test/manual/test_tokenizer_manager.py +++ b/test/manual/test_tokenizer_manager.py @@ -11,19 +11,17 @@ python3 -m unittest test_tokenizer_manager.TestTokenizerResultExtraction python3 -m unittest test_tokenizer_manager.TestTokenizerManagerIntegration python3 -m unittest test_tokenizer_manager.TestReqStateTextBuffering python3 -m unittest test_tokenizer_manager.TestReqStateCrashDump -python3 -m unittest test_tokenizer_manager.TestMakeReqState """ import asyncio import unittest from unittest.mock import Mock, patch -from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput +from sglang.srt.managers.io_struct import GenerateReqInput from sglang.srt.managers.tokenizer_manager import ( InputFormat, ReqState, TokenizerManager, - make_req_state, ) from sglang.srt.observability.req_time_stats import APIServerReqTimeStats from sglang.srt.server_args import PortArgs, ServerArgs @@ -415,7 +413,7 @@ class TestTokenizerManagerIntegration(unittest.TestCase): self.assertIsNone(result_token_type_ids) -def _make_state(buffer_text: bool = False) -> ReqState: +def _make_state() -> ReqState: """Create a minimal ReqState for testing.""" obj = Mock(spec=GenerateReqInput) return ReqState( @@ -424,25 +422,26 @@ def _make_state(buffer_text: bool = False) -> ReqState: event=asyncio.Event(), obj=obj, time_stats=APIServerReqTimeStats(), - buffer_text=buffer_text, ) class TestReqStateTextBuffering(unittest.TestCase): """Test ReqState.append_text / get_text in both buffering modes.""" - def test_streaming_mode_concatenates_directly(self): - state = _make_state(buffer_text=False) - state.append_text("hello ") - state.append_text("world") - self.assertEqual(state.get_text(), "hello world") - self.assertEqual(state.text_chunks, []) - - def test_buffer_mode_collects_chunks(self): - state = _make_state(buffer_text=True) + def test_collects_chunks_lazily(self): + state = _make_state() state.append_text("hello ") state.append_text("world") self.assertEqual(state.text, "") + self.assertEqual(state.text_chunks, ["hello ", "world"]) + self.assertEqual(state.get_text(), "hello world") + self.assertEqual(state.text_chunks, []) + + def test_get_text_preserves_materialized_prefix(self): + state = _make_state() + state.append_text("hello ") + self.assertEqual(state.get_text(), "hello ") + state.append_text("world") self.assertEqual(state.get_text(), "hello world") @@ -450,21 +449,21 @@ class TestReqStateCrashDump(unittest.TestCase): """Test ReqState.get_crash_dump_output.""" def test_empty_state(self): - state = _make_state(buffer_text=False) + state = _make_state() self.assertEqual(state.get_crash_dump_output(), {}) def test_with_text_only(self): - state = _make_state(buffer_text=False) + state = _make_state() state.append_text("partial output") self.assertEqual(state.get_crash_dump_output(), {"text": "partial output"}) def test_with_output_ids_only(self): - state = _make_state(buffer_text=False) + state = _make_state() state.output_ids = [1, 2, 3] self.assertEqual(state.get_crash_dump_output(), {"output_ids": [1, 2, 3]}) def test_with_text_and_output_ids(self): - state = _make_state(buffer_text=False) + state = _make_state() state.append_text("hello") state.output_ids = [10, 20] self.assertEqual( @@ -473,36 +472,5 @@ class TestReqStateCrashDump(unittest.TestCase): ) -class TestMakeReqState(unittest.TestCase): - """Test make_req_state factory function.""" - - def _call(self, *, obj_stream=None): - if obj_stream is not None: - obj = Mock(spec=GenerateReqInput) - obj.stream = obj_stream - else: - obj = Mock(spec=EmbeddingReqInput) - del obj.stream - return make_req_state( - out_list=[], - finished=False, - event=asyncio.Event(), - obj=obj, - time_stats=APIServerReqTimeStats(), - ) - - def test_streaming_request_does_not_buffer(self): - state = self._call(obj_stream=True) - self.assertFalse(state.buffer_text) - - def test_non_streaming_request_buffers(self): - state = self._call(obj_stream=False) - self.assertTrue(state.buffer_text) - - def test_embedding_request_always_buffers(self): - state = self._call() - self.assertTrue(state.buffer_text) - - if __name__ == "__main__": unittest.main(verbosity=2)