minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

PD streaming: batch notify + SSE fast path (#22658)

Browse Source
This commit is contained in:
inkcherry
2026-04-22 02:21:02 -07:00
committed by GitHub
parent fa85bdf4ed
commit 1e34cd0ba5
3 changed files with 134 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/entrypoints/openai/serving_chat.py
+105 -64
View File
@@ -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:
python/sglang/srt/managers/tokenizer_manager.py
+21 -12
View File
@@ -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 (
python/sglang/srt/server_args.py
+8
View File
@@ -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",