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:
co-authored by
fzyzcjy
Yueming Yuan
parent
589f90b368
commit
3259a2c789
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user