[API] add get memory pool size (#1760)

Co-authored-by: Byron Hsu <byronhsu1230@gmail.com>
This commit is contained in:
Ying Sheng
2024-10-23 00:02:29 -07:00
committed by GitHub
parent ad4125d1a9
commit 2fce449b1c
6 changed files with 50 additions and 0 deletions

View File

@@ -27,6 +27,7 @@ from sglang.srt.managers.io_struct import (
BatchEmbeddingOut,
BatchStrOut,
BatchTokenIDOut,
GetMemPoolSizeReqOutput,
UpdateWeightReqOutput,
)
from sglang.srt.managers.schedule_batch import FINISH_MATCHED_STR, FINISH_MATCHED_TOKEN
@@ -111,6 +112,9 @@ class DetokenizerManager:
# If it is a weight update request, no detokenization is needed.
self.send_to_tokenizer.send_pyobj(recv_obj)
continue
elif isinstance(recv_obj, GetMemPoolSizeReqOutput):
self.send_to_tokenizer.send_pyobj(recv_obj)
continue
elif self.tokenizer is None:
# If the tokenizer is skipped, no detokenization is needed
self.send_to_tokenizer.send_pyobj(recv_obj)