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

[Feature] Enable return routed experts (#12162)

Browse Source 
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>
This commit is contained in:
Junrong Lin
2025-12-21 15:16:43 +08:00
committed by GitHub
co-authored by yizhang2077 Liangsheng Yin
parent 8fe3e37468
commit bed301a5ac
27 changed files with 646 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/detokenizer_manager.py
+18
View File
@@ -21,6 +21,7 @@ from collections import OrderedDict
from typing import Dict, List, Union
import psutil
import pybase64
import setproctitle
import zmq
@@ -266,8 +267,24 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
return output_strs
def _extract_routed_experts(self, recv_obj: BatchTokenIDOutput) -> List[List[int]]:
output_routed_experts = None
if recv_obj.output_routed_experts is not None:
output_routed_experts = [
(
pybase64.b64encode(output_routed_experts.numpy().tobytes()).decode(
"utf-8"
)
if output_routed_experts is not None
else []
)
for output_routed_experts in recv_obj.output_routed_experts
]
return output_routed_experts
def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput):
output_strs = self._decode_batch_token_id_output(recv_obj)
output_routed_experts = self._extract_routed_experts(recv_obj)
return BatchStrOutput(
rids=recv_obj.rids,
@@ -294,6 +311,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
output_token_ids_logprobs_idx=recv_obj.output_token_ids_logprobs_idx,
output_token_entropy_val=recv_obj.output_token_entropy_val,
output_hidden_states=recv_obj.output_hidden_states,
output_routed_experts=output_routed_experts,
placeholder_tokens_idx=None,
placeholder_tokens_val=None,
retraction_counts=recv_obj.retraction_counts,
python/sglang/srt/managers/io_struct.py
+14
View File
@@ -23,6 +23,8 @@ from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
import torch
from sglang.srt.lora.lora_registry import LoRARef
from sglang.srt.managers.schedule_batch import BaseFinishReason
from sglang.srt.multimodal.mm_utils import has_valid_data
@@ -196,6 +198,8 @@ class GenerateReqInput(BaseReq, APIServingTimingMixin):
log_metrics: bool = True
# Whether to return hidden states
return_hidden_states: Union[List[bool], bool] = False
# Whether to return captured routed experts
return_routed_experts: bool = False
# The modalities of the image data [image, multi-images, video]
modalities: Optional[List[str]] = None
@@ -618,6 +622,7 @@ class GenerateReqInput(BaseReq, APIServingTimingMixin):
if isinstance(self.return_hidden_states, list)
else self.return_hidden_states
),
return_routed_experts=self.return_routed_experts,
modalities=self.modalities[i] if self.modalities else None,
session_params=self.session_params,
lora_path=self.lora_path[i] if self.lora_path is not None else None,
@@ -687,6 +692,9 @@ class TokenizedGenerateReqInput(BaseReq):
# Whether to return hidden states
return_hidden_states: bool = False
# Whether to return captured routed experts
return_routed_experts: bool = False
# The input embeds
input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None
@@ -953,6 +961,9 @@ class BatchTokenIDOutput(
# Hidden states
output_hidden_states: List[List[float]]
# The routed experts for each output token
output_routed_experts: List[torch.Tensor]
# The information of placeholder tokens (e.g., image token)
# idx is the index of the token in the prompt after expansion.
# val is the length of padded tokens after expansion.
@@ -1032,6 +1043,9 @@ class BatchStrOutput(
# Hidden states
output_hidden_states: List[List[float]]
# The routed experts for each output token
output_routed_experts: List[List[int]]
# The information of placeholder tokens (e.g., image token)
# idx is the index of the token in the prompt after expansion.
# val is the length of padded tokens after expansion.
python/sglang/srt/managers/multi_tokenizer_mixin.py
+3
View File
@@ -270,6 +270,9 @@ def _handle_output_by_index(output, i):
output_hidden_states=_extract_field_by_index(
output, "output_hidden_states", i, check_length=False
),
output_routed_experts=_extract_field_by_index(
output, "output_routed_experts", i, check_length=False
),
placeholder_tokens_idx=None,
placeholder_tokens_val=None,
retraction_counts=_extract_field_by_index(output, "retraction_counts", i),
python/sglang/srt/managers/schedule_batch.py
+12
View File
@@ -489,6 +489,7 @@ class Req:
custom_logit_processor: Optional[str] = None,
require_reasoning: bool = False,
return_hidden_states: bool = False,
return_routed_experts: bool = False,
eos_token_ids: Optional[Set[int]] = None,
bootstrap_host: Optional[str] = None,
bootstrap_port: Optional[int] = None,
@@ -679,6 +680,12 @@ class Req:
self.output_topk_p = None
self.output_topk_index = None
# capture routed experts
self.return_routed_experts = return_routed_experts
self.routed_experts: Optional[torch.Tensor] = (
None # cpu tensor: shape (seqlen, topk)
)
# Embedding (return values)
self.embedding = None
@@ -1043,6 +1050,7 @@ class Req:
self.retraction_count += 1
self.prefix_indices = torch.empty((0,), dtype=torch.int64)
self.routed_experts = None
self.last_node = None
self.swa_uuid_for_lock = None
self.extend_input_len = 0
@@ -1219,6 +1227,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# Whether to return hidden states
return_hidden_states: bool = False
# Whether to return captured experts
return_routed_experts: bool = False
# Whether this batch is prefill-only (no token generation needed)
is_prefill_only: bool = False
@@ -1266,6 +1277,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
device=req_to_token_pool.device,
spec_algorithm=spec_algorithm,
return_hidden_states=any(req.return_hidden_states for req in reqs),
return_routed_experts=any(req.return_routed_experts for req in reqs),
is_prefill_only=all(req.is_prefill_only for req in reqs),
chunked_req=chunked_req,
dllm_config=dllm_config,
python/sglang/srt/managers/scheduler.py
+1
View File
@@ -1413,6 +1413,7 @@ class Scheduler(
custom_logit_processor=recv_req.custom_logit_processor,
require_reasoning=recv_req.require_reasoning,
return_hidden_states=recv_req.return_hidden_states,
return_routed_experts=recv_req.return_routed_experts,
eos_token_ids=self.model_config.hf_eos_token_id,
bootstrap_host=recv_req.bootstrap_host,
bootstrap_port=recv_req.bootstrap_port,
python/sglang/srt/managers/scheduler_output_processor_mixin.py
+18
View File
@@ -9,6 +9,7 @@ import torch
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.environ import envs
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.moe.routed_experts_capturer import get_global_experts_capturer
from sglang.srt.managers.io_struct import (
AbortReq,
BatchEmbeddingOutput,
@@ -62,6 +63,14 @@ class SchedulerOutputProcessorMixin:
trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
self.stream_output(batch.reqs, batch.return_logprob)
def maybe_collect_routed_experts(self: Scheduler, req: Req):
"""Collect routed experts for a finished request."""
req.routed_experts = get_global_experts_capturer().get_routed_experts(
req_pool_idx=req.req_pool_idx,
seqlen=req.seqlen,
req_to_token_pool=self.req_to_token_pool,
)
def process_batch_result_prefill(
self: Scheduler,
batch: ScheduleBatch,
@@ -116,6 +125,7 @@ class SchedulerOutputProcessorMixin:
req.check_finished()
if req.finished():
self.maybe_collect_routed_experts(req)
release_kv_cache(req, self.tree_cache)
req.time_stats.completion_time = time.perf_counter()
elif not batch.decoding_reqs or req not in batch.decoding_reqs:
@@ -372,6 +382,8 @@ class SchedulerOutputProcessorMixin:
req.check_finished(new_accepted_len)
if req.finished():
self.maybe_collect_routed_experts(req)
if self.server_args.disaggregation_decode_enable_offload_kvcache:
# Asynchronously offload KV cache; release_kv_cache will be called after Device->Host transfer completes
if not self.decode_offload_manager.offload_kv_cache(req):
@@ -791,6 +803,7 @@ class SchedulerOutputProcessorMixin:
spec_accepted_tokens = []
retraction_counts = []
output_hidden_states = None
output_routed_experts = None
queue_times = []
forward_entry_times = []
@@ -985,6 +998,10 @@ class SchedulerOutputProcessorMixin:
if output_hidden_states is None:
output_hidden_states = []
output_hidden_states.append(req.hidden_states)
if req.return_routed_experts:
if output_routed_experts is None:
output_routed_experts = []
output_routed_experts.append(req.routed_experts)
if (
req.finished()
@@ -1034,6 +1051,7 @@ class SchedulerOutputProcessorMixin:
output_token_ids_logprobs_idx=output_token_ids_logprobs_idx,
output_token_entropy_val=None,
output_hidden_states=output_hidden_states,
output_routed_experts=output_routed_experts,
placeholder_tokens_idx=None,
placeholder_tokens_val=None,
retraction_counts=retraction_counts,
python/sglang/srt/managers/tokenizer_manager.py
+4
View File
@@ -846,6 +846,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
custom_logit_processor=obj.custom_logit_processor,
require_reasoning=obj.require_reasoning,
return_hidden_states=obj.return_hidden_states,
return_routed_experts=obj.return_routed_experts,
data_parallel_rank=obj.data_parallel_rank,
priority=obj.priority,
extra_key=obj.extra_key,
@@ -1585,6 +1586,9 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
if getattr(recv_obj, "output_hidden_states", None):
meta_info["hidden_states"] = recv_obj.output_hidden_states[i]
if getattr(recv_obj, "output_routed_experts", None):
meta_info["routed_experts"] = recv_obj.output_routed_experts[i]
if isinstance(recv_obj, BatchStrOutput):
state.text += recv_obj.output_strs[i]
if self.server_args.stream_output and state.obj.stream:
python/sglang/srt/managers/tp_worker.py
+1 -9
View File
@@ -317,15 +317,7 @@ class TpModelWorker(BaseTpWorker):
# Profile number of tokens
self.max_total_num_tokens = self.model_runner.max_total_num_tokens
self.max_prefill_tokens = server_args.max_prefill_tokens
self.max_running_requests = min(
(
self.max_total_num_tokens // 2
if server_args.max_running_requests is None
else server_args.max_running_requests
// (server_args.dp_size if server_args.enable_dp_attention else 1)
),
self.model_runner.req_to_token_pool.size,
)
self.max_running_requests = self.model_runner.max_running_requests
assert self.max_running_requests > 0, "max_running_request is zero"
self.max_queued_requests = server_args.max_queued_requests
assert (