diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 1173e0b3b..406314925 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -9,6 +9,7 @@ from http import HTTPStatus from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Union import jinja2 +import msgspec import orjson from fastapi import Request from fastapi.responses import ORJSONResponse, StreamingResponse @@ -53,6 +54,75 @@ from sglang.srt.parser.conversation import generate_chat_conv from sglang.srt.parser.jinja_template_utils import process_content_for_template_format from sglang.srt.parser.reasoning_parser import ReasoningParser +_SSE_DATA_B = b"data: " +_SSE_NL_B = b"\n\n" + + +class _StreamDelta(msgspec.Struct, omit_defaults=True): + # OpenAI Python SDK's ChoiceDelta does not declare reasoning_content; it is + # surfaced via pydantic `extra`. With omit_defaults=True, defaulting to + # None would drop the key entirely from the SSE payload, making + # `data.reasoning_content` raise AttributeError on the client. Keep it + # required (no default) so it is always serialized as null or a string. + reasoning_content: Optional[str] + role: Optional[str] = None + content: Optional[str] = None + + +class _StreamChoice(msgspec.Struct): + index: int + delta: _StreamDelta + logprobs: Optional[dict] = None + finish_reason: Optional[str] = None + matched_stop: Union[None, int, str] = None + + +class _StreamChunk(msgspec.Struct, omit_defaults=True): + id: str + object: str + created: int + model: str + choices: List[_StreamChoice] + usage: Optional[dict] = None + + +_stream_encoder = msgspec.json.Encoder() + + +def _fast_sse_content( + chunk_id: str, + created: int, + model: str, + index: int, + role: Optional[str] = None, + content: Optional[str] = None, + reasoning_content: Optional[str] = None, + finish_reason: Optional[str] = None, + logprobs: Optional[dict] = None, + matched_stop: Union[None, int, str] = None, + usage: Optional[dict] = None, +) -> str: + delta = _StreamDelta( + role=role, content=content, reasoning_content=reasoning_content + ) + choice = _StreamChoice( + index=index, + delta=delta, + logprobs=logprobs, + finish_reason=finish_reason, + matched_stop=matched_stop, + ) + chunk = _StreamChunk( + id=chunk_id, + object="chat.completion.chunk", + created=created, + model=model, + choices=[choice], + usage=usage, + ) + return (_SSE_DATA_B + _stream_encoder.encode(chunk) + _SSE_NL_B).decode() + + if TYPE_CHECKING: from sglang.srt.managers.template_manager import TemplateManager from sglang.srt.managers.tokenizer_manager import TokenizerManager @@ -721,7 +791,7 @@ class OpenAIServingChat(OpenAIServingBase): if n_prev_token < total_output_logprobs: choice_logprobs = self._process_streaming_logprobs( content, n_prev_token, total_output_logprobs - ) + ).model_dump() n_prev_tokens[index] = total_output_logprobs finish_reason = content["meta_info"].get("finish_reason", None) @@ -751,20 +821,14 @@ class OpenAIServingChat(OpenAIServingBase): # First chunk with role if is_firsts.get(index, True): is_firsts[index] = False - delta = DeltaMessage(role="assistant", content="") - choice_data = ChatCompletionResponseStreamChoice( - index=index, - delta=delta, - finish_reason=None, - logprobs=None, - ) - chunk = ChatCompletionStreamResponse( - id=content["meta_info"]["id"], + yield _fast_sse_content( + chunk_id=content["meta_info"]["id"], created=int(time.time()), - choices=[choice_data], model=request.model, + index=index, + role="assistant", + content="", ) - yield f"data: {chunk.model_dump_json()}\n\n" stream_started = True offset = stream_offsets.get(index, 0) @@ -781,27 +845,22 @@ class OpenAIServingChat(OpenAIServingBase): index, delta, reasoning_parser_dict, content, request ) if reasoning_text: - choice_data = ChatCompletionResponseStreamChoice( - index=index, - delta=DeltaMessage(reasoning_content=reasoning_text), - finish_reason=None, - ) - chunk = ChatCompletionStreamResponse( - id=content["meta_info"]["id"], - created=int(time.time()), - choices=[choice_data], - model=request.model, - ) - - # Add usage stats if continuous_usage_stats is enabled + usage = None if continuous_usage_stats: - chunk.usage = UsageProcessor.calculate_token_usage( + usage = UsageProcessor.calculate_token_usage( prompt_tokens=prompt_tokens.get(index, 0), reasoning_tokens=reasoning_tokens.get(index, 0), completion_tokens=completion_tokens.get(index, 0), - ) + ).model_dump() - yield f"data: {chunk.model_dump_json()}\n\n" + yield _fast_sse_content( + chunk_id=content["meta_info"]["id"], + created=int(time.time()), + model=request.model, + index=index, + reasoning_content=reasoning_text, + usage=usage, + ) # Handle tool calls if ( @@ -833,29 +892,23 @@ class OpenAIServingChat(OpenAIServingBase): else: # Regular content if delta: - choice_data = ChatCompletionResponseStreamChoice( - index=index, - delta=DeltaMessage(content=delta), - finish_reason=None, - matched_stop=None, - logprobs=choice_logprobs, - ) - chunk = ChatCompletionStreamResponse( - id=content["meta_info"]["id"], - created=int(time.time()), - choices=[choice_data], - model=request.model, - ) - - # Add usage stats if continuous_usage_stats is enabled + usage = None if continuous_usage_stats: - chunk.usage = UsageProcessor.calculate_token_usage( + usage = UsageProcessor.calculate_token_usage( prompt_tokens=prompt_tokens.get(index, 0), reasoning_tokens=reasoning_tokens.get(index, 0), completion_tokens=completion_tokens.get(index, 0), - ) + ).model_dump() - yield f"data: {chunk.model_dump_json()}\n\n" + yield _fast_sse_content( + chunk_id=content["meta_info"]["id"], + created=int(time.time()), + model=request.model, + index=index, + content=delta, + logprobs=choice_logprobs, + usage=usage, + ) # Send finish_reason chunks for each index that completed for idx, finish_reason_data in finish_reasons.items(): @@ -866,27 +919,15 @@ class OpenAIServingChat(OpenAIServingBase): if has_tool_calls.get(idx, False) and finish_reason_type == "stop": final_finish_reason = "tool_calls" - finish_reason_chunk = ChatCompletionStreamResponse( - id=content["meta_info"][ - "id" - ], # NOTE: openai uses the same chatcmpl-id for all indices + matched_stop = finish_reason_data.get("matched") + yield _fast_sse_content( + chunk_id=content["meta_info"]["id"], created=int(time.time()), - choices=[ - ChatCompletionResponseStreamChoice( - index=idx, - delta=DeltaMessage(), - finish_reason=final_finish_reason, - matched_stop=( - finish_reason_data["matched"] - if "matched" in finish_reason_data - else None - ), - ) - ], model=request.model, - usage=None, + index=idx, + finish_reason=final_finish_reason, + matched_stop=matched_stop, ) - yield f"data: {finish_reason_chunk.model_dump_json()}\n\n" # Send hidden states if requested if request.return_hidden_states and hidden_states: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 729ed3cd7..1ac16f718 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -487,14 +487,6 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): def init_request_dispatcher(self): self._result_dispatcher = TypeBasedDispatcher( [ - ( - ( - BatchStrOutput, - BatchEmbeddingOutput, - BatchTokenIDOutput, - ), - self._handle_batch_output, - ), (AbortReq, self._handle_abort_req), (OpenSessionReqOutput, self._handle_open_session_req_output), ( @@ -1623,11 +1615,17 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): while True: with self.soft_watchdog.disable(): recv_obj = await self.recv_from_detokenizer.recv_pyobj() - self._result_dispatcher(recv_obj) + if isinstance( + recv_obj, + (BatchStrOutput, BatchEmbeddingOutput, BatchTokenIDOutput), + ): + await self._handle_batch_output(recv_obj) + else: + self._result_dispatcher(recv_obj) self.last_receive_tstamp = real_time() self.soft_watchdog.feed() - def _handle_batch_output( + async def _handle_batch_output( self, recv_obj: Union[ BatchStrOutput, @@ -1635,6 +1633,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): BatchTokenIDOutput, ], ): + pending_notify: dict[str, ReqState] = {} + batch_notify_size = self.server_args.batch_notify_size for i, rid in enumerate(recv_obj.rids): state = self.rid_to_state.get(rid, None) if state is None: @@ -1833,9 +1833,14 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): if out_dict is not None: state.out_list.append(out_dict) - state.event.set() + pending_notify[rid] = state + + if len(pending_notify) >= batch_notify_size: + for s in pending_notify.values(): + s.event.set() + pending_notify = {} + await asyncio.sleep(0) - # Log metrics and dump if self.enable_metrics and state.obj.log_metrics: self.collect_metrics(state, recv_obj, i) if self.dump_requests_folder and state.finished and state.obj.log_metrics: @@ -1843,6 +1848,10 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): if self.crash_dump_folder and state.finished and state.obj.log_metrics: self.record_request_for_crash_dump(state, out_dict) + # handle_loop awaits next recv immediately + for s in pending_notify.values(): + s.event.set() + # When skip_tokenizer_init is enabled, tokensizer_manager receives # BatchTokenIDOutput. if ( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1e2a1d95e..beb8398f8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -379,6 +379,7 @@ class ServerArgs: pp_max_micro_batch_size: Optional[int] = None pp_async_batch_depth: int = 0 stream_interval: int = 1 + batch_notify_size: int = 16 stream_response_default_include_usage: bool = False incremental_streaming_output: bool = False enable_streaming_session: bool = False @@ -4542,6 +4543,13 @@ class ServerArgs: default=ServerArgs.stream_interval, help="The interval (or buffer size) for streaming in terms of the token length. A smaller value makes streaming smoother, while a larger value makes the throughput higher", ) + parser.add_argument( + "--batch-notify-size", + type=int, + default=ServerArgs.batch_notify_size, + help="Number of streaming notifications to batch before yielding to the event loop. " + "Reduces asyncio wakeup overhead under high concurrency.", + ) parser.add_argument( "--incremental-streaming-output", action="store_true",