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>

python/sglang/srt/managers/detokenizer_manager.py

@@ -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,

python/sglang/srt/managers/io_struct.py

@@ -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
+    # Per-request routed experts (input + output tokens), shape
-    # routed_experts[i] is a tensor of shape (token, layer, top_k) for request i
+    # (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
+    # Per-request routed experts, base64-encoded by DetokenizerManager off the
-    # routed_experts[i] is a tensor of shape (token, layer, top_k) for request i
+    # tokenizer hot path. Underlying tensor shape is (token, layer, top_k);
-    routed_experts: List[Optional[torch.Tensor]]
+    # 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.

python/sglang/srt/managers/tokenizer_manager.py

@@ -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]
+                val = recv_obj.routed_experts[i]
-                if routed_experts_tensor is not None:
+                if val is not None:
-                    meta_info["routed_experts"] = pybase64.b64encode(
+                    # BatchStrOutput is pre-encoded by the detokenizer;
-                        routed_experts_tensor.numpy().tobytes()
+                    # BatchTokenIDOutput (skip_tokenizer_init) bypasses it.
-                    ).decode("utf-8")
+                    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]