Encode routed_experts in the detokenizer, off the tokenizer hot path (#24263)

Co-authored-by: fzyzcjy <ch271828n@outlook.com>
Co-authored-by: Yueming Yuan <yym022502@gmail.com>
This commit is contained in:
Liangsheng Yin
2026-05-02 02:44:32 -07:00
committed by GitHub
co-authored by fzyzcjy Yueming Yuan
parent 589f90b368
commit 3259a2c789
3 changed files with 38 additions and 11 deletions
@@ -21,7 +21,9 @@ from collections import OrderedDict, defaultdict
from typing import Dict, List, Optional, Tuple, Union
import psutil
import pybase64
import setproctitle
import torch
import zmq
from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
@@ -320,6 +322,25 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
return output_strs
@staticmethod
def _b64_encode_per_request(
data_list: Optional[List[Optional[torch.Tensor]]],
) -> Optional[List[Optional[str]]]:
"""Encode a per-request list of tensors as base64 strings, off the
tokenizer hot path. Returns None when the input is None; per-item None
stays None.
"""
if data_list is None:
return None
return [
(
pybase64.b64encode(item.numpy().tobytes()).decode("utf-8")
if item is not None
else None
)
for item in data_list
]
def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput):
# If handling idle batch, set output_strs to [].
output_strs = (
@@ -327,6 +348,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
if len(recv_obj.rids) > 0
else []
)
routed_experts = self._b64_encode_per_request(recv_obj.routed_experts)
return BatchStrOutput(
rids=recv_obj.rids,
http_worker_ipcs=recv_obj.http_worker_ipcs,
@@ -355,7 +377,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,
routed_experts=recv_obj.routed_experts,
routed_experts=routed_experts,
customized_info=recv_obj.customized_info,
placeholder_tokens_idx=None,
placeholder_tokens_val=None,
+8 -5
View File
@@ -1101,8 +1101,10 @@ class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
# Hidden states
output_hidden_states: List[List[float]]
# The routed experts for each token, including both input and output tokens
# routed_experts[i] is a tensor of shape (token, layer, top_k) for request i
# Per-request routed experts (input + output tokens), shape
# (token, layer, top_k). DetokenizerManager encodes to base64 into
# BatchStrOutput; on the skip_tokenizer_init path the scheduler sends this
# straight to TokenizerManager, which encodes on demand.
routed_experts: List[Optional[torch.Tensor]]
# The information of placeholder tokens (e.g., image token)
@@ -1163,9 +1165,10 @@ class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
# Hidden states
output_hidden_states: List[List[float]]
# The routed experts for each token, including both input and output tokens
# routed_experts[i] is a tensor of shape (token, layer, top_k) for request i
routed_experts: List[Optional[torch.Tensor]]
# Per-request routed experts, base64-encoded by DetokenizerManager off the
# tokenizer hot path. Underlying tensor shape is (token, layer, top_k);
# see BatchTokenIDOutput.routed_experts.
routed_experts: List[Optional[str]]
# The information of placeholder tokens (e.g., image token)
# idx is the index of the token in the prompt after expansion.
@@ -1703,11 +1703,13 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
if getattr(recv_obj, "output_hidden_states", None):
meta_info["hidden_states"] = recv_obj.output_hidden_states[i]
if getattr(recv_obj, "routed_experts", None):
routed_experts_tensor = recv_obj.routed_experts[i]
if routed_experts_tensor is not None:
meta_info["routed_experts"] = pybase64.b64encode(
routed_experts_tensor.numpy().tobytes()
).decode("utf-8")
val = recv_obj.routed_experts[i]
if val is not None:
# BatchStrOutput is pre-encoded by the detokenizer;
# BatchTokenIDOutput (skip_tokenizer_init) bypasses it.
if isinstance(val, torch.Tensor):
val = pybase64.b64encode(val.numpy().tobytes()).decode("utf-8")
meta_info["routed_experts"] = val
if getattr(recv_obj, "customized_info", None):
for k, v in recv_obj.customized_info.items():
meta_info[k] = v[i]