[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
2026-04-10 21:26:04 -07:00
committed by GitHub
co-authored by Claude Opus 4.6
parent c7f93a2ce7
commit 8eac618a8d
2 changed files with 34 additions and 97 deletions
+17 -48
View File
@@ -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: