diff --git a/docs/basic_usage/openai_api_completions.ipynb b/docs/basic_usage/openai_api_completions.ipynb index d498f13ed..e89dfd57f 100644 --- a/docs/basic_usage/openai_api_completions.ipynb +++ b/docs/basic_usage/openai_api_completions.ipynb @@ -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": {}, diff --git a/docs/basic_usage/sampling_params.md b/docs/basic_usage/sampling_params.md index e27d844e6..a1848d41d 100644 --- a/docs/basic_usage/sampling_params.md +++ b/docs/basic_usage/sampling_params.md @@ -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 diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 060e9659a..40ad9f3fb 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -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 diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index f7d79b7f8..2acf07cd5 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -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) diff --git a/python/sglang/srt/entrypoints/openai/serving_completions.py b/python/sglang/srt/entrypoints/openai/serving_completions.py index 8229de122..10c04552c 100644 --- a/python/sglang/srt/entrypoints/openai/serving_completions.py +++ b/python/sglang/srt/entrypoints/openai/serving_completions.py @@ -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) diff --git a/python/sglang/srt/entrypoints/openai/utils.py b/python/sglang/srt/entrypoints/openai/utils.py index 94ac5458d..c4c9baede 100644 --- a/python/sglang/srt/entrypoints/openai/utils.py +++ b/python/sglang/srt/entrypoints/openai/utils.py @@ -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) diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index 50b861e14..a032de8ed 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -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 ] diff --git a/test/registered/rl/test_return_routed_experts.py b/test/registered/rl/test_return_routed_experts.py index 480f2e32a..248acf10d 100644 --- a/test/registered/rl/test_return_routed_experts.py +++ b/test/registered/rl/test_return_routed_experts.py @@ -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)