EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Improve: Use TypeBasedDispatcher in DetokenizerManager (#3117)

Browse Source
This commit is contained in:
fzyzcjy
2025-02-22 11:50:46 +08:00
committed by GitHub
parent a3339d8cac
commit 9087694006
1 changed files with 112 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
python/sglang/srt/managers/detokenizer_manager.py
View File

@@ -32,7 +32,11 @@ from sglang.srt.managers.io_struct import (
) )
from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.utils import configure_logger, get_zmq_socket from sglang.srt.utils import configure_logger, get_zmq_socket
from sglang.utils import find_printable_text, get_exception_traceback from sglang.utils import (
TypeBasedDispatcher,
find_printable_text,
get_exception_traceback,
)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -83,6 +87,13 @@ class DetokenizerManager:
self.decode_status = LimitedCapacityDict(capacity=DETOKENIZER_MAX_STATES) self.decode_status = LimitedCapacityDict(capacity=DETOKENIZER_MAX_STATES)
self._request_dispatcher = TypeBasedDispatcher(
[
(BatchEmbeddingOut, self.handle_batch_embedding_out),
(BatchTokenIDOut, self.handle_batch_token_id_out),
]
)
def trim_matched_stop( def trim_matched_stop(
self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool
): ):
@@ -111,14 +122,14 @@ class DetokenizerManager:
while True: while True:
recv_obj = self.recv_from_scheduler.recv_pyobj() recv_obj = self.recv_from_scheduler.recv_pyobj()
output = self._request_dispatcher(recv_obj)
self.send_to_tokenizer.send_pyobj(output)
if isinstance(recv_obj, BatchEmbeddingOut): def handle_batch_embedding_out(self, recv_obj: BatchEmbeddingOut):
# If it is embedding model, no detokenization is needed. # If it is embedding model, no detokenization is needed.
self.send_to_tokenizer.send_pyobj(recv_obj) return recv_obj
continue
else:
assert isinstance(recv_obj, BatchTokenIDOut)
def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOut):
bs = len(recv_obj.rids) bs = len(recv_obj.rids)
# Initialize decode status # Initialize decode status
@@ -196,8 +207,7 @@ class DetokenizerManager:
) )
) )
self.send_to_tokenizer.send_pyobj( out = BatchStrOut(
BatchStrOut(
rids=recv_obj.rids, rids=recv_obj.rids,
finished_reasons=recv_obj.finished_reasons, finished_reasons=recv_obj.finished_reasons,
output_strs=output_strs, output_strs=output_strs,
@@ -215,12 +225,13 @@ class DetokenizerManager:
output_top_logprobs_idx=recv_obj.output_top_logprobs_idx, output_top_logprobs_idx=recv_obj.output_top_logprobs_idx,
output_hidden_states=recv_obj.output_hidden_states, output_hidden_states=recv_obj.output_hidden_states,
) )
)
# remove decodestatus for completed requests # remove decodestatus for completed requests
for rid in finished_reqs: for rid in finished_reqs:
self.decode_status.pop(rid) self.decode_status.pop(rid)
return out
class LimitedCapacityDict(OrderedDict): class LimitedCapacityDict(OrderedDict):
def __init__(self, capacity: int, *args, **kwargs): def __init__(self, capacity: int, *args, **kwargs):