Skip to content

GitLab

Unverified Commit 90876940 authored by fzyzcjy's avatar fzyzcjy Committed by GitHub
Browse files

Improve: Use TypeBasedDispatcher in DetokenizerManager (#3117)

parent a3339d8c
Hide whitespace changes
Inline Side-by-side
Showing with 112 additions and 101 deletions
python/sglang/srt/managers/detokenizer_manager.py
View file @ 90876940
...@@ -32,7 +32,11 @@ from sglang.srt.managers.io_struct import ( ...@@ -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: ...@@ -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,115 +122,115 @@ class DetokenizerManager: ...@@ -111,115 +122,115 @@ 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)
if isinstance(recv_obj, BatchEmbeddingOut): self.send_to_tokenizer.send_pyobj(output)
# If it is embedding model, no detokenization is needed.
self.send_to_tokenizer.send_pyobj(recv_obj) def handle_batch_embedding_out(self, recv_obj: BatchEmbeddingOut):
continue # If it is embedding model, no detokenization is needed.
return recv_obj
def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOut):
bs = len(recv_obj.rids)
# Initialize decode status
read_ids, surr_ids = [], []
for i in range(bs):
rid = recv_obj.rids[i]
vid = recv_obj.vids[i]
if rid not in self.decode_status or self.decode_status[rid].vid != vid:
s = DecodeStatus(
vid=vid,
decoded_text=recv_obj.decoded_texts[i],
decode_ids=recv_obj.decode_ids[i],
surr_offset=0,
read_offset=recv_obj.read_offsets[i],
)
self.decode_status[rid] = s
else: else:
assert isinstance(recv_obj, BatchTokenIDOut) s = self.decode_status[rid]
s.decode_ids = recv_obj.decode_ids[i]
bs = len(recv_obj.rids)
read_ids.append(
# Initialize decode status self.trim_matched_stop(
read_ids, surr_ids = [], [] s.decode_ids[s.surr_offset :],
for i in range(bs): recv_obj.finished_reasons[i],
rid = recv_obj.rids[i] recv_obj.no_stop_trim[i],
vid = recv_obj.vids[i]
if rid not in self.decode_status or self.decode_status[rid].vid != vid:
s = DecodeStatus(
vid=vid,
decoded_text=recv_obj.decoded_texts[i],
decode_ids=recv_obj.decode_ids[i],
surr_offset=0,
read_offset=recv_obj.read_offsets[i],
)
self.decode_status[rid] = s
else:
s = self.decode_status[rid]
s.decode_ids = recv_obj.decode_ids[i]
read_ids.append(
self.trim_matched_stop(
s.decode_ids[s.surr_offset :],
recv_obj.finished_reasons[i],
recv_obj.no_stop_trim[i],
)
) )
surr_ids.append(s.decode_ids[s.surr_offset : s.read_offset])
# TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request
surr_texts = self.tokenizer.batch_decode(
surr_ids,
skip_special_tokens=recv_obj.skip_special_tokens[0],
spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
)
read_texts = self.tokenizer.batch_decode(
read_ids,
skip_special_tokens=recv_obj.skip_special_tokens[0],
spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
) )
surr_ids.append(s.decode_ids[s.surr_offset : s.read_offset])
# Incremental decoding # TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request
output_strs = [] surr_texts = self.tokenizer.batch_decode(
finished_reqs = [] surr_ids,
for i in range(bs): skip_special_tokens=recv_obj.skip_special_tokens[0],
try: spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
s = self.decode_status[recv_obj.rids[i]] )
except KeyError: read_texts = self.tokenizer.batch_decode(
raise RuntimeError( read_ids,
f"Decode status not found for request {recv_obj.rids[i]}. " skip_special_tokens=recv_obj.skip_special_tokens[0],
"It may be due to the request being evicted from the decode status due to memory pressure. " spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
"Please increase the maximum number of requests by setting " )
"the SGLANG_DETOKENIZER_MAX_STATES environment variable to a bigger value than the default value. "
f"The current value is {DETOKENIZER_MAX_STATES}. " # Incremental decoding
"For more details, see: https://github.com/sgl-project/sglang/issues/2812" output_strs = []
) finished_reqs = []
new_text = read_texts[i][len(surr_texts[i]) :] for i in range(bs):
if recv_obj.finished_reasons[i] is None: try:
# Streaming chunk: update the decode status s = self.decode_status[recv_obj.rids[i]]
if len(new_text) > 0 and not new_text.endswith("�"): except KeyError:
s.decoded_text = s.decoded_text + new_text raise RuntimeError(
s.surr_offset = s.read_offset f"Decode status not found for request {recv_obj.rids[i]}. "
s.read_offset = len(s.decode_ids) "It may be due to the request being evicted from the decode status due to memory pressure. "
new_text = "" "Please increase the maximum number of requests by setting "
else: "the SGLANG_DETOKENIZER_MAX_STATES environment variable to a bigger value than the default value. "
new_text = find_printable_text(new_text) f"The current value is {DETOKENIZER_MAX_STATES}. "
else: "For more details, see: https://github.com/sgl-project/sglang/issues/2812"
finished_reqs.append(recv_obj.rids[i])
output_strs.append(
self.trim_matched_stop(
s.decoded_text + new_text,
recv_obj.finished_reasons[i],
recv_obj.no_stop_trim[i],
)
) )
new_text = read_texts[i][len(surr_texts[i]) :]
if recv_obj.finished_reasons[i] is None:
# Streaming chunk: update the decode status
if len(new_text) > 0 and not new_text.endswith("�"):
s.decoded_text = s.decoded_text + new_text
s.surr_offset = s.read_offset
s.read_offset = len(s.decode_ids)
new_text = ""
else:
new_text = find_printable_text(new_text)
else:
finished_reqs.append(recv_obj.rids[i])
self.send_to_tokenizer.send_pyobj( output_strs.append(
BatchStrOut( self.trim_matched_stop(
rids=recv_obj.rids, s.decoded_text + new_text,
finished_reasons=recv_obj.finished_reasons, recv_obj.finished_reasons[i],
output_strs=output_strs, recv_obj.no_stop_trim[i],
prompt_tokens=recv_obj.prompt_tokens,
completion_tokens=recv_obj.completion_tokens,
cached_tokens=recv_obj.cached_tokens,
spec_verify_ct=recv_obj.spec_verify_ct,
input_token_logprobs_val=recv_obj.input_token_logprobs_val,
input_token_logprobs_idx=recv_obj.input_token_logprobs_idx,
output_token_logprobs_val=recv_obj.output_token_logprobs_val,
output_token_logprobs_idx=recv_obj.output_token_logprobs_idx,
input_top_logprobs_val=recv_obj.input_top_logprobs_val,
input_top_logprobs_idx=recv_obj.input_top_logprobs_idx,
output_top_logprobs_val=recv_obj.output_top_logprobs_val,
output_top_logprobs_idx=recv_obj.output_top_logprobs_idx,
output_hidden_states=recv_obj.output_hidden_states,
) )
) )
# remove decodestatus for completed requests out = BatchStrOut(
for rid in finished_reqs: rids=recv_obj.rids,
self.decode_status.pop(rid) finished_reasons=recv_obj.finished_reasons,
output_strs=output_strs,
prompt_tokens=recv_obj.prompt_tokens,
completion_tokens=recv_obj.completion_tokens,
cached_tokens=recv_obj.cached_tokens,
spec_verify_ct=recv_obj.spec_verify_ct,
input_token_logprobs_val=recv_obj.input_token_logprobs_val,
input_token_logprobs_idx=recv_obj.input_token_logprobs_idx,
output_token_logprobs_val=recv_obj.output_token_logprobs_val,
output_token_logprobs_idx=recv_obj.output_token_logprobs_idx,
input_top_logprobs_val=recv_obj.input_top_logprobs_val,
input_top_logprobs_idx=recv_obj.input_top_logprobs_idx,
output_top_logprobs_val=recv_obj.output_top_logprobs_val,
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): class LimitedCapacityDict(OrderedDict):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment