minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[RL] Fix crash when the reqs in a batch have a mix of return_routed_experts = True and False. (#26423)

Browse Source 
Co-authored-by: root <root@slurm-h200-209-231.slurm-compute.tenant-slurm.svc.cluster.local>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
This commit is contained in:
Byron Hsu
2026-05-29 20:46:10 -07:00
committed by GitHub
co-authored by root Cursor Cheng Wan
parent 804f01a823
commit 6f1c9fc77b
3 changed files with 141 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler_components/output_streamer.py
+41 -12
View File
@@ -121,8 +121,21 @@ class SchedulerOutputStreamer:
skip_req: Optional[Req] = None,
is_idle_batch: bool = False,
):
return_hidden_states = any(
req.return_hidden_states for req in reqs if req is not skip_req
)
return_routed_experts = any(
req.return_routed_experts for req in reqs if req is not skip_req
)
return_indexer_topk = any(
req.return_indexer_topk for req in reqs if req is not skip_req
)
acc = _GenerationStreamAccumulator(
return_logprob=return_logprob,
return_hidden_states=return_hidden_states,
return_routed_experts=return_routed_experts,
return_indexer_topk=return_indexer_topk,
spec_algorithm=self.spec_algorithm,
disaggregation_mode=self.disaggregation_mode,
default_stream_interval=self.server_args.stream_interval,
@@ -229,6 +242,9 @@ class SchedulerOutputStreamer:
@dataclass(slots=True, kw_only=True)
class _GenerationStreamAccumulator:
return_logprob: bool
return_hidden_states: bool
return_routed_experts: bool
return_indexer_topk: bool
spec_algorithm: Any
disaggregation_mode: DisaggregationMode
default_stream_interval: int
@@ -256,9 +272,9 @@ class _GenerationStreamAccumulator:
spec_num_correct_drafts: list = field(default_factory=list)
spec_correct_drafts_histogram: list = field(default_factory=list)
retraction_counts: list = field(default_factory=list)
output_hidden_states: list = field(default_factory=list)
routed_experts: list = field(default_factory=list)
indexer_topk: list = field(default_factory=list)
output_hidden_states: Optional[list] = None
routed_experts: Optional[list] = None
indexer_topk: Optional[list] = None
customized_info: dict = field(default_factory=dict)
time_stats: list = field(default_factory=list)
input_token_logprobs_val: Optional[list] = None
@@ -275,6 +291,13 @@ class _GenerationStreamAccumulator:
output_token_ids_logprobs_idx: Optional[list] = None
def __post_init__(self) -> None:
if self.return_hidden_states:
self.output_hidden_states = []
if self.return_routed_experts:
self.routed_experts = []
if self.return_indexer_topk:
self.indexer_topk = []
if self.return_logprob:
self.input_token_logprobs_val = []
self.input_token_logprobs_idx = []
@@ -434,12 +457,18 @@ class _GenerationStreamAccumulator:
self.output_token_ids_logprobs_val.append([])
self.output_token_ids_logprobs_idx.append([])
if req.return_hidden_states:
self.output_hidden_states.append(req.hidden_states)
if req.return_routed_experts:
self.routed_experts.append(req.routed_experts)
if req.return_indexer_topk:
self.indexer_topk.append(req.indexer_topk)
if self.return_hidden_states:
self.output_hidden_states.append(
req.hidden_states if req.return_hidden_states else None
)
if self.return_routed_experts:
self.routed_experts.append(
req.routed_experts if req.return_routed_experts else None
)
if self.return_indexer_topk:
self.indexer_topk.append(
req.indexer_topk if req.return_indexer_topk else None
)
if req.customized_info is not None:
for k, v in req.customized_info.items():
@@ -486,9 +515,9 @@ class _GenerationStreamAccumulator:
output_token_ids_logprobs_val=self.output_token_ids_logprobs_val,
output_token_ids_logprobs_idx=self.output_token_ids_logprobs_idx,
output_token_entropy_val=None,
output_hidden_states=self.output_hidden_states or None,
routed_experts=self.routed_experts or None,
indexer_topk=self.indexer_topk or None,
output_hidden_states=self.output_hidden_states,
routed_experts=self.routed_experts,
indexer_topk=self.indexer_topk,
customized_info=self.customized_info,
placeholder_tokens_idx=None,
placeholder_tokens_val=None,
python/sglang/srt/managers/tokenizer_manager.py
+3 -1
View File
@@ -1784,7 +1784,9 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
]
if getattr(recv_obj, "output_hidden_states", None):
meta_info["hidden_states"] = recv_obj.output_hidden_states[i]
hidden_states = recv_obj.output_hidden_states[i]
if hidden_states is not None:
meta_info["hidden_states"] = hidden_states
if getattr(recv_obj, "routed_experts", None):
val = recv_obj.routed_experts[i]
if val is not None: