Improve: Use TypeBasedDispatcher in DetokenizerManager (#3117)

python/sglang/srt/managers/detokenizer_manager.py

@@ -32,7 +32,11 @@ from sglang.srt.managers.io_struct import (
 )
 from sglang.srt.server_args import PortArgs, ServerArgs
 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__)
 
@@ -83,6 +87,13 @@ class DetokenizerManager:
 
         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(
         self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool
     ):
@@ -111,14 +122,14 @@ class DetokenizerManager:
 
         while True:
             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.
-                self.send_to_tokenizer.send_pyobj(recv_obj)
-                continue
-            else:
-                assert isinstance(recv_obj, BatchTokenIDOut)
+        return recv_obj
 
+    def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOut):
         bs = len(recv_obj.rids)
 
         # Initialize decode status
@@ -196,8 +207,7 @@ class DetokenizerManager:
                 )
             )
 
-            self.send_to_tokenizer.send_pyobj(
-                BatchStrOut(
+        out = BatchStrOut(
             rids=recv_obj.rids,
             finished_reasons=recv_obj.finished_reasons,
             output_strs=output_strs,
@@ -215,12 +225,13 @@ class DetokenizerManager:
             output_top_logprobs_idx=recv_obj.output_top_logprobs_idx,
             output_hidden_states=recv_obj.output_hidden_states,
         )
-            )
 
         # remove decodestatus for completed requests
         for rid in finished_reqs:
             self.decode_status.pop(rid)
 
+        return out
+
 
 class LimitedCapacityDict(OrderedDict):
     def __init__(self, capacity: int, *args, **kwargs):