[tokenizer] lazy text accumulation + use deltas directly for streaming (#22548)
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alex Nails
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
c7f93a2ce7
commit
8eac618a8d
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user