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

18
docs/basic_usage/openai_api_completions.ipynb
View File

@@ -368,6 +368,15 @@
" print(chunk.choices[0].delta.content, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Returning Routed Experts (MoE Models)\n",
"\n",
"For MoE models, set `return_routed_experts: true` in `extra_body` to return expert routing data. Requires `--enable-return-routed-experts` server flag. The `routed_experts` field will be returned in the `sgl_ext` object on each choice, containing base64-encoded int32 expert IDs as a flattened array with logical shape `[num_tokens, num_layers, top_k]`."
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -453,6 +462,15 @@
"print_highlight(f\"Response: {response}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Returning Routed Experts (MoE Models)\n",
"\n",
"For MoE models, set `return_routed_experts: true` in `extra_body` to return expert routing data. Requires `--enable-return-routed-experts` server flag. The `routed_experts` field will be returned in the `sgl_ext` object on each choice, containing base64-encoded int32 expert IDs as a flattened array with logical shape `[num_tokens, num_layers, top_k]`."
]
},
{
"cell_type": "markdown",
"metadata": {},

1
docs/basic_usage/sampling_params.md
View File

@@ -25,6 +25,7 @@ The `/generate` endpoint accepts the following parameters in JSON format. For de
| lora_path | `Optional[Union[List[Optional[str]], Optional[str]]] = None` | The path to the LoRA. |
| custom_logit_processor | `Optional[Union[List[Optional[str]], str]] = None` | Custom logit processor for advanced sampling control. Must be a serialized instance of `CustomLogitProcessor` using its `to_str()` method. For usage see below. |
| return_hidden_states | `Union[List[bool], bool] = False` | Whether to return hidden states. |
| return_routed_experts | `bool = False` | Whether to return routed experts for MoE models. Requires `--enable-return-routed-experts` server flag. Returns base64-encoded int32 expert IDs as a flattened array with logical shape `[num_tokens, num_layers, top_k]`. |
## Sampling parameters

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
]

132
test/registered/rl/test_return_routed_experts.py
View File

@@ -21,7 +21,7 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=180, suite="stage-c-test-large-4-gpu")
register_cuda_ci(est_time=360, suite="stage-c-test-large-4-gpu")
SHAREGPT_URL = (
"https://huggingface.co/datasets/anon8231489123/"
@@ -81,15 +81,42 @@ class TestReturnRoutedExperts(CustomTestCase):
if not cls.texts:
raise ValueError("No valid texts found in the dataset")
cls.texts = cls.texts[:100]
cls._endpoints = [
(
"/generate",
cls._build_generate_payload,
extract_routed_experts_from_meta_info,
),
(
"/v1/chat/completions",
cls._build_chat_payload,
extract_routed_experts_from_openai_response,
),
(
"/v1/completions",
cls._build_completion_payload,
extract_routed_experts_from_openai_response,
),
]
cls.baseline_results = cls._collect_results(cls.baseline_args)
cls.reference_results = cls._collect_results(cls.reference_args)
@classmethod
def test_return_routed_experts(cls):
captured_baseline_experts = asyncio.run(
cls.fetch_result("baseline", cls.baseline_args)
)
captured_reference_experts = asyncio.run(
cls.fetch_result("reference", cls.reference_args)
)
cls._run_endpoint_test("/generate")
@classmethod
def test_return_routed_experts_chat_completions(cls):
cls._run_endpoint_test("/v1/chat/completions")
@classmethod
def test_return_routed_experts_completions(cls):
cls._run_endpoint_test("/v1/completions")
@classmethod
def _run_endpoint_test(cls, endpoint):
captured_baseline_experts = cls.baseline_results[endpoint]
captured_reference_experts = cls.reference_results[endpoint]
check_all_experts_id_valid(captured_baseline_experts)
check_all_experts_id_valid(captured_reference_experts)
@@ -114,43 +141,69 @@ class TestReturnRoutedExperts(CustomTestCase):
), f"Too many mismatches: {num_mismatches} out of {num_baseline_topks} ({num_mismatches/num_baseline_topks:.4%})"
@classmethod
async def fetch_result(cls, title, other_args):
def _collect_results(
cls,
other_args,
):
process = popen_launch_server(
DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST,
DEFAULT_URL_FOR_TEST,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
try:
process = popen_launch_server(
DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST,
DEFAULT_URL_FOR_TEST,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
async with aiohttp.ClientSession() as session:
return asyncio.run(cls._collect_results_async())
finally:
kill_process_tree(process.pid)
@classmethod
async def _collect_results_async(cls):
results = {}
async with aiohttp.ClientSession() as session:
for endpoint, payload_builder, response_extractor in cls._endpoints:
tasks = [
asyncio.create_task(
make_request(
session,
f"{DEFAULT_URL_FOR_TEST}/generate",
{
"text": text,
"sampling_params": cls.sampling_args,
"return_routed_experts": True,
"max_new_tokens": 100,
},
f"{DEFAULT_URL_FOR_TEST}{endpoint}",
payload_builder(text),
)
)
for text in cls.texts
]
# return value shape: List[[seq_len, num_layers, topk]...]
http_result = await asyncio.gather(*tasks)
except Exception as e:
raise e
finally:
kill_process_tree(process.pid)
results[endpoint] = [
response_extractor(res).reshape(-1, 48, 8) for res in http_result
]
return results
result = [
extract_routed_experts_from_meta_info(res).reshape(-1, 48, 8)
for res in http_result
]
@classmethod
def _build_generate_payload(cls, text):
return {
"text": text,
"sampling_params": cls.sampling_args,
"return_routed_experts": True,
"max_new_tokens": 100,
}
return result
@classmethod
def _build_chat_payload(cls, text):
return {
"messages": [{"role": "user", "content": text}],
"temperature": 0,
"max_tokens": 100,
"return_routed_experts": True,
}
@classmethod
def _build_completion_payload(cls, text):
return {
"prompt": text,
"temperature": 0,
"max_tokens": 100,
"return_routed_experts": True,
}
async def make_request(session, url, payload):
@@ -159,6 +212,23 @@ async def make_request(session, url, payload):
return await response.json()
def extract_routed_experts_from_openai_response(response):
if "error" in response:
raise ValueError(f"OpenAI response error: {response['error']}")
choices = response.get("choices", [])
if not choices:
raise ValueError("OpenAI response has no choices.")
sgl_ext = choices[0].get("sgl_ext", None)
if sgl_ext is None:
raise ValueError("OpenAI response missing sgl_ext.")
routed_experts = sgl_ext.get("routed_experts", None)
if routed_experts is None:
raise ValueError("OpenAI response sgl_ext missing routed_experts.")
return extract_routed_experts_from_meta_info(
{"meta_info": {"routed_experts": routed_experts}}
)
def check_all_experts_id_valid(experts: List[List[List[int]]]):
tensor_list = [torch.tensor(lst) for lst in experts]
padded_tensor = pad_sequence(tensor_list, batch_first=True, padding_value=0)