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

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": {},

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
 

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
 
 

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)
 

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)
 

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)

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
             ]

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._run_endpoint_test("/generate")
-            cls.fetch_result("baseline", cls.baseline_args)
+
-        )
+    @classmethod
-        captured_reference_experts = asyncio.run(
+    def test_return_routed_experts_chat_completions(cls):
-            cls.fetch_result("reference", cls.reference_args)
+        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(
+            return asyncio.run(cls._collect_results_async())
-                DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST,
+        finally:
-                DEFAULT_URL_FOR_TEST,
+            kill_process_tree(process.pid)
-                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+
-                other_args=other_args,
+    @classmethod
-            )
+    async def _collect_results_async(cls):
-            async with aiohttp.ClientSession() as session:
+        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",
+                            f"{DEFAULT_URL_FOR_TEST}{endpoint}",
-                            {
+                            payload_builder(text),
-                                "text": text,
-                                "sampling_params": cls.sampling_args,
-                                "return_routed_experts": True,
-                                "max_new_tokens": 100,
-                            },
                         )
                     )
                     for text in cls.texts
                 ]
                 # return value shape: List[[seq_len, num_layers, topk]...]
                 http_result = await asyncio.gather(*tasks)
-        except Exception as e:
+                results[endpoint] = [
-            raise e
+                    response_extractor(res).reshape(-1, 48, 8) for res in http_result
-        finally:
+                ]
-            kill_process_tree(process.pid)
+        return results
 
-        result = [
+    @classmethod
-            extract_routed_experts_from_meta_info(res).reshape(-1, 48, 8)
+    def _build_generate_payload(cls, text):
-            for res in http_result
+        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)