chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Add return routed experts to the completions and chat/completions endpoints (#17434)

Browse Source
This commit is contained in:
Mansoor
2026-01-23 15:12:36 -05:00
committed by GitHub
parent 5438cd20ce
commit bdaa3de075
8 changed files with 225 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
python/sglang/srt/entrypoints/openai/protocol.py
View File

@@ -232,6 +232,7 @@ class CompletionRequest(BaseModel):
top_p: float = 1.0
user: Optional[str] = None
return_hidden_states: bool = False
return_routed_experts: bool = False
# Extra parameters for SRT backend only and will be ignored by OpenAI models.
top_k: int = -1
@@ -280,6 +281,22 @@ class CompletionRequest(BaseModel):
return v
class SglExt(BaseModel):
"""SGLang extension fields for OpenAI-compatible responses.
Future SGLang-specific extensions to OpenAI-compatible response objects
should be added as fields here rather than directly on the choice object.
"""
routed_experts: Optional[str] = None
@model_serializer(mode="wrap")
def _serialize(self, handler):
data = handler(self)
# Remove None fields to keep response clean
return {k: v for k, v in data.items() if v is not None}
class CompletionResponseChoice(BaseModel):
index: int
text: str
@@ -287,12 +304,15 @@ class CompletionResponseChoice(BaseModel):
finish_reason: Optional[Literal["stop", "length", "content_filter", "abort"]] = None
matched_stop: Union[None, int, str] = None
hidden_states: Optional[object] = None
sgl_ext: Optional[SglExt] = None
@model_serializer(mode="wrap")
def _serialize(self, handler):
data = handler(self)
if self.hidden_states is None:
data.pop("hidden_states", None)
if self.sgl_ext is None:
data.pop("sgl_ext", None)
return data
@@ -313,12 +333,15 @@ class CompletionResponseStreamChoice(BaseModel):
finish_reason: Optional[Literal["stop", "length", "content_filter", "abort"]] = None
matched_stop: Union[None, int, str] = None
hidden_states: Optional[object] = None
sgl_ext: Optional[SglExt] = None
@model_serializer(mode="wrap")
def _serialize(self, handler):
data = handler(self)
if self.hidden_states is None:
data.pop("hidden_states", None)
if self.sgl_ext is None:
data.pop("sgl_ext", None)
return data
@@ -502,6 +525,7 @@ class ChatCompletionRequest(BaseModel):
default="auto", examples=["none"]
) # noqa
return_hidden_states: bool = False
return_routed_experts: bool = False
reasoning_effort: Optional[Literal["low", "medium", "high"]] = Field(
default="medium",
description="Constrains effort on reasoning for reasoning models. "
@@ -731,12 +755,15 @@ class ChatCompletionResponseChoice(BaseModel):
] = None
matched_stop: Union[None, int, str] = None
hidden_states: Optional[object] = None
sgl_ext: Optional[SglExt] = None
@model_serializer(mode="wrap")
def _serialize(self, handler):
data = handler(self)
if self.hidden_states is None:
data.pop("hidden_states", None)
if self.sgl_ext is None:
data.pop("sgl_ext", None)
return data
@@ -756,12 +783,15 @@ class DeltaMessage(BaseModel):
reasoning_content: Optional[str] = None
tool_calls: Optional[List[ToolCall]] = Field(default=None, examples=[None])
hidden_states: Optional[object] = None
sgl_ext: Optional[SglExt] = None
@model_serializer(mode="wrap")
def _serialize(self, handler):
data = handler(self)
if self.hidden_states is None:
data.pop("hidden_states", None)
if self.sgl_ext is None:
data.pop("sgl_ext", None)
return data

30
python/sglang/srt/entrypoints/openai/serving_chat.py
View File

@@ -28,6 +28,7 @@ from sglang.srt.entrypoints.openai.protocol import (
FunctionResponse,
LogProbs,
MessageProcessingResult,
SglExt,
ToolCall,
ToolCallProcessingResult,
ToolChoice,
@@ -37,6 +38,7 @@ from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
from sglang.srt.entrypoints.openai.usage_processor import UsageProcessor
from sglang.srt.entrypoints.openai.utils import (
process_hidden_states_from_ret,
process_routed_experts_from_ret,
to_openai_style_logprobs,
)
from sglang.srt.function_call.core_types import ToolCallItem
@@ -298,6 +300,7 @@ class OpenAIServingChat(OpenAIServingBase):
bootstrap_room=request.bootstrap_room,
data_parallel_rank=request.data_parallel_rank,
return_hidden_states=request.return_hidden_states,
return_routed_experts=request.return_routed_experts,
rid=request.rid,
extra_key=self._compute_extra_key(request),
require_reasoning=self._get_reasoning_from_request(request),
@@ -609,6 +612,7 @@ class OpenAIServingChat(OpenAIServingBase):
completion_tokens = {}
cached_tokens = {}
hidden_states = {}
routed_experts = {}
try:
async for content in self.tokenizer_manager.generate_request(
@@ -620,6 +624,7 @@ class OpenAIServingChat(OpenAIServingBase):
completion_tokens[index] = content["meta_info"]["completion_tokens"]
cached_tokens[index] = content["meta_info"].get("cached_tokens", 0)
hidden_states[index] = content["meta_info"].get("hidden_states", None)
routed_experts[index] = content["meta_info"].get("routed_experts", None)
# Handle logprobs
choice_logprobs = None
@@ -801,6 +806,27 @@ class OpenAIServingChat(OpenAIServingBase):
)
yield f"data: {hidden_states_chunk.model_dump_json()}\n\n"
if request.return_routed_experts and routed_experts:
for index, choice_routed_experts in routed_experts.items():
if choice_routed_experts is not None:
routed_experts_chunk = ChatCompletionStreamResponse(
id=content["meta_info"]["id"],
created=int(time.time()),
choices=[
ChatCompletionResponseStreamChoice(
index=index,
delta=DeltaMessage(
sgl_ext=SglExt(
routed_experts=choice_routed_experts
)
),
finish_reason=None,
)
],
model=request.model,
)
yield (f"data: {routed_experts_chunk.model_dump_json()}\n\n")
# Additional usage chunk
if request.stream_options and request.stream_options.include_usage:
usage = UsageProcessor.calculate_streaming_usage(
@@ -867,6 +893,7 @@ class OpenAIServingChat(OpenAIServingBase):
# Handle hidden states
hidden_states = process_hidden_states_from_ret(ret_item, request)
routed_experts = process_routed_experts_from_ret(ret_item, request)
finish_reason = ret_item["meta_info"]["finish_reason"]
text = ret_item["text"]
@@ -926,6 +953,9 @@ class OpenAIServingChat(OpenAIServingBase):
else None
),
hidden_states=hidden_states,
sgl_ext=(
SglExt(routed_experts=routed_experts) if routed_experts else None
),
)
choices.append(choice_data)

30
python/sglang/srt/entrypoints/openai/serving_completions.py
View File

@@ -14,11 +14,13 @@ from sglang.srt.entrypoints.openai.protocol import (
CompletionResponseStreamChoice,
CompletionStreamResponse,
ErrorResponse,
SglExt,
)
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
from sglang.srt.entrypoints.openai.usage_processor import UsageProcessor
from sglang.srt.entrypoints.openai.utils import (
process_hidden_states_from_ret,
process_routed_experts_from_ret,
to_openai_style_logprobs,
)
from sglang.srt.managers.io_struct import GenerateReqInput
@@ -118,6 +120,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
bootstrap_room=request.bootstrap_room,
data_parallel_rank=request.data_parallel_rank,
return_hidden_states=request.return_hidden_states,
return_routed_experts=request.return_routed_experts,
rid=request.rid,
extra_key=self._compute_extra_key(request),
priority=request.priority,
@@ -203,6 +206,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
completion_tokens = {}
cached_tokens = {}
hidden_states = {}
routed_experts = {}
try:
async for content in self.tokenizer_manager.generate_request(
@@ -215,6 +219,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
completion_tokens[index] = content["meta_info"]["completion_tokens"]
cached_tokens[index] = content["meta_info"].get("cached_tokens", 0)
hidden_states[index] = content["meta_info"].get("hidden_states", None)
routed_experts[index] = content["meta_info"].get("routed_experts", None)
stream_buffer = stream_buffers.get(index, "")
# Handle echo for first chunk
@@ -311,6 +316,27 @@ class OpenAIServingCompletion(OpenAIServingBase):
)
yield f"data: {hidden_states_chunk.model_dump_json()}\n\n"
if request.return_routed_experts and routed_experts:
for index, choice_routed_experts in routed_experts.items():
if choice_routed_experts is not None:
routed_experts_chunk = CompletionStreamResponse(
id=content["meta_info"]["id"],
created=created,
object="text_completion",
choices=[
CompletionResponseStreamChoice(
index=index,
text="",
sgl_ext=SglExt(
routed_experts=choice_routed_experts
),
finish_reason=None,
)
],
model=request.model,
)
yield (f"data: {routed_experts_chunk.model_dump_json()}\n\n")
# Handle final usage chunk
if request.stream_options and request.stream_options.include_usage:
usage = UsageProcessor.calculate_streaming_usage(
@@ -409,6 +435,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
# Handle hidden states
hidden_states = process_hidden_states_from_ret(ret_item, request)
routed_experts = process_routed_experts_from_ret(ret_item, request)
finish_reason = ret_item["meta_info"]["finish_reason"]
@@ -423,6 +450,9 @@ class OpenAIServingCompletion(OpenAIServingBase):
else None
),
hidden_states=hidden_states,
sgl_ext=(
SglExt(routed_experts=routed_experts) if routed_experts else None
),
)
choices.append(choice_data)

13
python/sglang/srt/entrypoints/openai/utils.py
View File

@@ -70,3 +70,16 @@ def process_hidden_states_from_ret(
if hidden_states is not None:
hidden_states = hidden_states[-1] if len(hidden_states) > 1 else []
return hidden_states
def process_routed_experts_from_ret(
ret_item: Dict[str, Any],
request: Union[
ChatCompletionRequest,
CompletionRequest,
],
) -> Optional[str]:
"""Process routed experts from a ret item in non-streaming response."""
if not getattr(request, "return_routed_experts", False):
return None
return ret_item["meta_info"].get("routed_experts", None)

4
python/sglang/srt/managers/detokenizer_manager.py
View File

@@ -328,14 +328,14 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
def _extract_routed_experts(
self, recv_obj: BatchTokenIDOutput
) -> List[List[int]] | None:
) -> list[str | None] | None:
routed_experts = None
if recv_obj.routed_experts is not None:
routed_experts = [
(
pybase64.b64encode(routed_experts.numpy().tobytes()).decode("utf-8")
if routed_experts is not None
else []
else None
)
for routed_experts in recv_obj.routed_experts
]