sglang 0.1.15__py3-none-any.whl → 0.1.17__py3-none-any.whl

sglang/srt/managers/detokenizer_manager.py

@@ -1,4 +1,5 @@
 import asyncio
+import inspect
 
 import uvloop
 import zmq
@@ -7,7 +8,8 @@ import zmq.asyncio
 from sglang.srt.hf_transformers_utils import get_tokenizer
 from sglang.srt.managers.io_struct import BatchStrOut, BatchTokenIDOut
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import get_exception_traceback
+from sglang.utils import get_exception_traceback, graceful_registry
+from sglang.srt.managers.controller.infer_batch import FINISH_MATCHED_STR
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
@@ -33,51 +35,47 @@ class DetokenizerManager:
 
     async def handle_loop(self):
         while True:
-            recv_obj = await self.recv_from_router.recv_pyobj()
-
-            if isinstance(recv_obj, BatchTokenIDOut):
-                output_tokens = recv_obj.output_tokens
-
-                # TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request
-                output_strs = self.tokenizer.batch_decode(
-                    output_tokens,
-                    skip_special_tokens=recv_obj.skip_special_tokens[0],
-                    spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[
-                        0
-                    ],
-                )
-
-                # Trim stop str
-                # TODO(lmzheng): handle the case where multiple stop strs are hit
-                for i in range(len(output_strs)):
-                    if recv_obj.hit_stop_str[i] is not None:
-                        pos = output_strs[i].find(recv_obj.hit_stop_str[i])
-                        if pos != -1:
-                            output_strs[i] = output_strs[i][:pos]
-
-                    if len(output_tokens[i]) > 0:
-                        first_token = self.tokenizer.convert_ids_to_tokens(
-                            int(output_tokens[i][0])
-                        )
-                        if not isinstance(first_token, str):
-                            first_token = first_token.decode("utf-8", errors="ignore")
-                        if first_token.startswith("▁"):
-                            output_strs[i] = " " + output_strs[i]
-
-                    output_strs[i] = (
-                        recv_obj.output_and_jump_forward_strs[i] + output_strs[i]
-                    )
-
-                self.send_to_tokenizer.send_pyobj(
-                    BatchStrOut(
-                        recv_obj.rids,
-                        output_strs,
-                        recv_obj.meta_info,
-                        recv_obj.finished,
+            recv_obj: BatchTokenIDOut = await self.recv_from_router.recv_pyobj()
+            assert isinstance(recv_obj, BatchTokenIDOut)
+
+            output_tokens = recv_obj.output_tokens
+
+            # TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request
+            output_strs = self.tokenizer.batch_decode(
+                output_tokens,
+                skip_special_tokens=recv_obj.skip_special_tokens[0],
+                spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[
+                    0
+                ],
+            )
+
+            # Trim stop str
+            # TODO(lmzheng): handle the case where multiple stop strs are hit
+            for i in range(len(output_strs)):
+                if len(output_tokens[i]) > 0:
+                    first_token = self.tokenizer.convert_ids_to_tokens(
+                        int(output_tokens[i][0])
                     )
+                    if not isinstance(first_token, str):
+                        first_token = first_token.decode("utf-8", errors="ignore")
+                    if first_token.startswith("▁"):
+                        output_strs[i] = " " + output_strs[i]
+
+                output_strs[i] = recv_obj.prev_output_strs[i] + output_strs[i]
+
+                if isinstance(recv_obj.finished_reason[i], FINISH_MATCHED_STR):
+                    pos = output_strs[i].find(recv_obj.finished_reason[i].matched)
+                    if pos != -1:
+                        output_strs[i] = output_strs[i][:pos]
+
+            self.send_to_tokenizer.send_pyobj(
+                BatchStrOut(
+                    rids=recv_obj.rids,
+                    output_str=output_strs,
+                    meta_info=recv_obj.meta_info,
+                    finished_reason=recv_obj.finished_reason,
                 )
-            else:
-                raise ValueError(f"Invalid object: {recv_obj}")
+            )
 
 
 def start_detokenizer_process(
@@ -85,9 +83,11 @@ def start_detokenizer_process(
     port_args: PortArgs,
     pipe_writer,
 ):
+    graceful_registry(inspect.currentframe().f_code.co_name)
+
     try:
         manager = DetokenizerManager(server_args, port_args)
-    except Exception as e:
+    except Exception:
         pipe_writer.send(get_exception_traceback())
         raise
     pipe_writer.send("init ok")

sglang/srt/managers/io_struct.py

@@ -3,12 +3,15 @@ from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
 from sglang.srt.sampling_params import SamplingParams
+from sglang.srt.managers.controller.infer_batch import BaseFinishReason
 
 
 @dataclass
 class GenerateReqInput:
     # The input prompt
-    text: Union[List[str], str]
+    text: Optional[Union[List[str], str]] = None
+    # The token ids for text; one can either specify text or input_ids
+    input_ids: Optional[Union[List[List[int]], List[int]]] = None
     # The image input
     image_data: Optional[Union[List[str], str]] = None
     # The sampling_params
@@ -25,10 +28,19 @@ class GenerateReqInput:
     return_text_in_logprobs: bool = False
     # Whether to stream output
     stream: bool = False
-    # TODO: make all parameters a Union[List[T], T] to allow for batched requests
 
     def post_init(self):
-        is_single = isinstance(self.text, str)
+
+        if (self.text is None and self.input_ids is None) or (
+            self.text is not None and self.input_ids is not None
+        ):
+            raise ValueError("Either text or input_ids should be provided.")
+
+        if self.text is not None:
+            is_single = isinstance(self.text, str)
+        else:
+            is_single = isinstance(self.input_ids[0], int)
+        self.is_single = is_single
 
         if is_single:
             if self.sampling_params is None:
@@ -42,7 +54,7 @@ class GenerateReqInput:
             if self.top_logprobs_num is None:
                 self.top_logprobs_num = 0
         else:
-            num = len(self.text)
+            num = len(self.text) if self.text is not None else len(self.input_ids)
 
             if self.image_data is None:
                 self.image_data = [None] * num
@@ -57,7 +69,8 @@ class GenerateReqInput:
             if self.rid is None:
                 self.rid = [uuid.uuid4().hex for _ in range(num)]
             else:
-                assert isinstance(self.rid, list)
+                if not isinstance(self.rid, list):
+                    raise ValueError("The rid should be a list.")
 
             if self.return_logprob is None:
                 self.return_logprob = [False] * num
@@ -93,21 +106,19 @@ class TokenizedGenerateReqInput:
 @dataclass
 class BatchTokenIDOut:
     rids: List[str]
+    prev_output_strs: List[str]
     output_tokens: List[List[int]]
-    output_and_jump_forward_strs: List[str]
-    hit_stop_str: List[Optional[str]]
     skip_special_tokens: List[bool]
     spaces_between_special_tokens: List[bool]
     meta_info: List[Dict]
-    finished: List[bool]
-
+    finished_reason: List[BaseFinishReason]
 
 @dataclass
 class BatchStrOut:
     rids: List[str]
     output_str: List[str]
     meta_info: List[Dict]
-    finished: List[bool]
+    finished_reason: List[BaseFinishReason]
 
 
 @dataclass
@@ -115,6 +126,11 @@ class FlushCacheReq:
     pass
 
 
+@dataclass
+class AbortReq:
+    rid: str
+
+
 @dataclass
 class DetokenizeReqInput:
     input_ids: List[int]

sglang/srt/managers/router/infer_batch.py

@@ -19,18 +19,32 @@ class FinishReason(IntEnum):
     EOS_TOKEN = auto()
     LENGTH = auto()
     STOP_STR = auto()
+    ABORT = auto()
+
+    @staticmethod
+    def to_str(reason):
+        if reason == FinishReason.EOS_TOKEN:
+            return None
+        elif reason == FinishReason.LENGTH:
+            return "length"
+        elif reason == FinishReason.STOP_STR:
+            return "stop"
+        elif reason == FinishReason.ABORT:
+            return "abort"
+        else:
+            return None
 
 
 class Req:
-    def __init__(self, rid, input_text, input_ids):
+    def __init__(self, rid, origin_input_text, origin_input_ids):
         self.rid = rid
-        self.input_text = input_text
-        self.input_ids = input_ids
+        self.origin_input_text = origin_input_text
+        self.origin_input_ids = origin_input_ids
+        self.origin_input_ids_unpadded = origin_input_ids  # before image padding
+        self.prev_output_str = ""
+        self.prev_output_ids = []
         self.output_ids = []
-
-        # Since jump forward may retokenize the prompt with partial outputs,
-        # we maintain the original prompt length to report the correct usage.
-        self.prompt_tokens = len(input_ids)
+        self.input_ids = None  # input_ids = origin_input_ids + prev_output_ids
 
         # The number of decoded tokens for token usage report. Note that
         # this does not include the jump forward tokens.
@@ -52,6 +66,7 @@ class Req:
         self.finish_reason = None
         self.hit_stop_str = None
 
+        # Prefix info
         self.extend_input_len = 0
         self.prefix_indices = []
         self.last_node = None
@@ -62,67 +77,36 @@ class Req:
         self.top_logprobs_num = 0
         self.normalized_prompt_logprob = None
         self.prefill_token_logprobs = None
-        self.decode_token_logprobs = None
         self.prefill_top_logprobs = None
-        self.decode_top_logprobs = None
+        self.decode_token_logprobs = []
+        self.decode_top_logprobs = []
+        # The tokens is prefilled but need to be considered as decode tokens
+        # and should be updated for the decode logprobs
+        self.last_update_decode_tokens = 0
 
         # Constrained decoding
         self.regex_fsm = None
         self.regex_fsm_state = 0
         self.jump_forward_map = None
-        self.output_and_jump_forward_str = ""
-
-    def max_new_tokens(self):
-        return self.sampling_params.max_new_tokens
 
-    def jump_forward_and_retokenize(self, jump_forward_str, next_state):
-        old_output_str = self.tokenizer.decode(self.output_ids)
-        # FIXME: This logic does not really solve the problem of determining whether
-        # there should be a leading space.
-        first_token = self.tokenizer.convert_ids_to_tokens(self.output_ids[0])
+    def partial_decode(self, ids):
+        first_token = self.tokenizer.convert_ids_to_tokens(ids[0])
         first_token = (
             first_token.decode() if isinstance(first_token, bytes) else first_token
         )
-        if first_token.startswith("▁"):
-            old_output_str = " " + old_output_str
-        new_input_string = (
-            self.input_text
-            + self.output_and_jump_forward_str
-            + old_output_str
-            + jump_forward_str
-        )
-        new_input_ids = self.tokenizer.encode(new_input_string)
-        if self.pixel_values is not None:
-            # NOTE: This is a hack because the old input_ids contains the image padding
-            jump_forward_tokens_len = len(self.tokenizer.encode(jump_forward_str))
-        else:
-            jump_forward_tokens_len = (
-                len(new_input_ids) - len(self.input_ids) - len(self.output_ids)
-            )
-
-        # print("=" * 100)
-        # print(f"Catch jump forward:\n{jump_forward_str}")
-        # print(self.tokenizer.convert_ids_to_tokens(self.input_ids))
-        # print(self.tokenizer.convert_ids_to_tokens(new_input_ids))
+        return (" " if first_token.startswith("▁") else "") + self.tokenizer.decode(ids)
 
-        self.input_ids = new_input_ids
-        self.output_ids = []
-        self.sampling_params.max_new_tokens = max(
-            self.sampling_params.max_new_tokens - jump_forward_tokens_len, 0
-        )
-        self.regex_fsm_state = next_state
-        self.output_and_jump_forward_str = (
-            self.output_and_jump_forward_str + old_output_str + jump_forward_str
-        )
-
-        # print(f"Output and jump forward str:\n{self.output_and_jump_forward_str}")
-        # print("*" * 100)
+    def max_new_tokens(self):
+        return self.sampling_params.max_new_tokens
 
     def check_finished(self):
         if self.finished:
             return
 
-        if len(self.output_ids) >= self.sampling_params.max_new_tokens:
+        if (
+            len(self.prev_output_ids) + len(self.output_ids)
+            >= self.sampling_params.max_new_tokens
+        ):
             self.finished = True
             self.finish_reason = FinishReason.LENGTH
             return
@@ -141,14 +125,66 @@ class Req:
             )
 
             for stop_str in self.sampling_params.stop_strs:
-                if stop_str in tail_str:
+                # FIXME: (minor) try incremental match in prev_output_str
+                if stop_str in tail_str or stop_str in self.prev_output_str:
                     self.finished = True
                     self.finish_reason = FinishReason.STOP_STR
                     self.hit_stop_str = stop_str
                     return
 
+    def jump_forward_and_retokenize(self, jump_forward_str, next_state):
+        # FIXME: This logic does not really solve the problem of determining whether
+        # there should be a leading space.
+        cur_output_str = self.partial_decode(self.output_ids)
+
+        # TODO(lsyin): apply re-tokenize only for decode tokens so that we do not need origin_input_text anymore
+        if self.origin_input_text is None:
+            # Recovering text can only use unpadded ids
+            self.origin_input_text = self.tokenizer.decode(
+                self.origin_input_ids_unpadded
+            )
+
+        all_text = (
+            self.origin_input_text
+            + self.prev_output_str
+            + cur_output_str
+            + jump_forward_str
+        )
+        all_ids = self.tokenizer.encode(all_text)
+        prompt_tokens = len(self.origin_input_ids_unpadded)
+        self.origin_input_ids = all_ids[:prompt_tokens]
+        self.origin_input_ids_unpadded = self.origin_input_ids
+        # NOTE: the output ids may not strictly correspond to the output text
+        old_prev_output_ids = self.prev_output_ids
+        self.prev_output_ids = all_ids[prompt_tokens:]
+        self.prev_output_str = self.prev_output_str + cur_output_str + jump_forward_str
+        self.output_ids = []
+
+        self.regex_fsm_state = next_state
+
+        if self.return_logprob:
+            # For fast-forward part's logprobs
+            k = 0
+            for i, old_id in enumerate(old_prev_output_ids):
+                if old_id == self.prev_output_ids[i]:
+                    k = k + 1
+                else:
+                    break
+            self.decode_token_logprobs = self.decode_token_logprobs[:k]
+            self.decode_top_logprobs = self.decode_top_logprobs[:k]
+            self.logprob_start_len = prompt_tokens + k
+            self.last_update_decode_tokens = len(self.prev_output_ids) - k
+
+        # print("=" * 100)
+        # print(f"Catch jump forward:\n{jump_forward_str}")
+        # print(self.tokenizer.convert_ids_to_tokens(self.input_ids))
+        # print(self.tokenizer.convert_ids_to_tokens(new_input_ids))
+
+        # print(f"Output and jump forward str:\n{self.output_and_jump_forward_str}")
+        # print("*" * 100)
+
     def __repr__(self):
-        return f"rid(n={self.rid}, " f"input_ids={self.input_ids}, "
+        return f"rid(n={self.rid}, " f"input_ids={self.origin_input_ids}, "
 
 
 @dataclass
@@ -319,6 +355,7 @@ class Batch:
 
     def retract_decode(self):
         sorted_indices = [i for i in range(len(self.reqs))]
+        # TODO(lsyin): improve the priority of retraction
         sorted_indices.sort(
             key=lambda i: (len(self.reqs[i].output_ids), -len(self.reqs[i].input_ids)),
             reverse=True,
@@ -332,25 +369,34 @@ class Batch:
             req = self.reqs[idx]
             retracted_reqs.append(req)
 
-            self.tree_cache.dec_ref_counter(req.last_node)
+            # TODO: apply more fine-grained retraction
+            last_uncached_pos = len(req.prefix_indices)
+            token_indices = self.req_to_token_pool.req_to_token[
+                req_pool_indices_cpu[idx]
+            ][last_uncached_pos : seq_lens_cpu[idx]]
+            self.token_to_kv_pool.dec_refs(token_indices)
+
+            # release the last node
+            self.tree_cache.dec_lock_ref(req.last_node)
+
+            cur_output_str = req.partial_decode(req.output_ids)
+            req.prev_output_str = req.prev_output_str + cur_output_str
+            req.prev_output_ids.extend(req.output_ids)
+
             req.prefix_indices = None
             req.last_node = None
             req.extend_input_len = 0
             req.output_ids = []
-            req.regex_fsm_state = 0
-
-            # TODO: apply more fine-grained retraction
 
-            token_indices = self.req_to_token_pool.req_to_token[
-                req_pool_indices_cpu[idx]
-            ][: seq_lens_cpu[idx]]
-            self.token_to_kv_pool.dec_refs(token_indices)
+            # For incremental logprobs
+            req.last_update_decode_tokens = 0
+            req.logprob_start_len = 10**9
 
         self.filter_batch(sorted_indices)
 
         return retracted_reqs
 
-    def check_for_jump_forward(self):
+    def check_for_jump_forward(self, model_runner):
         jump_forward_reqs = []
         filter_indices = [i for i in range(len(self.reqs))]
 
@@ -364,24 +410,34 @@ class Batch:
                     if len(jump_forward_str) <= 1:
                         continue
 
-                    # insert the old request into tree_cache
-                    token_ids_in_memory = tuple(req.input_ids + req.output_ids)[:-1]
                     if req_pool_indices_cpu is None:
                         req_pool_indices_cpu = self.req_pool_indices.tolist()
-                    req_pool_idx = req_pool_indices_cpu[i]
-                    indices = self.req_to_token_pool.req_to_token[
-                        req_pool_idx, : len(token_ids_in_memory)
-                    ]
-                    prefix_len = self.tree_cache.insert(
-                        token_ids_in_memory, indices.clone()
+
+                    # insert the old request into tree_cache
+                    self.tree_cache.cache_req(
+                        token_ids=tuple(req.input_ids + req.output_ids)[:-1],
+                        last_uncached_pos=len(req.prefix_indices),
+                        req_pool_idx=req_pool_indices_cpu[i],
                     )
-                    self.token_to_kv_pool.dec_refs(indices[:prefix_len])
-                    self.req_to_token_pool.free(req_pool_idx)
-                    self.tree_cache.dec_ref_counter(req.last_node)
+
+                    # unlock the last node
+                    self.tree_cache.dec_lock_ref(req.last_node)
 
                     # jump-forward
                     req.jump_forward_and_retokenize(jump_forward_str, next_state)
 
+                    # re-applying image padding
+                    if req.pixel_values is not None:
+                        (
+                            req.origin_input_ids,
+                            req.image_offset,
+                        ) = model_runner.model.pad_input_ids(
+                            req.origin_input_ids_unpadded,
+                            req.pad_value,
+                            req.pixel_values.shape,
+                            req.image_size,
+                        )
+
                     jump_forward_reqs.append(req)
                     filter_indices.remove(i)
 

sglang/srt/managers/router/manager.py

@@ -5,10 +5,10 @@ import uvloop
 import zmq
 import zmq.asyncio
 
-from sglang.srt.backend_config import GLOBAL_BACKEND_CONFIG
+from sglang.global_config import global_config
 from sglang.srt.managers.router.model_rpc import ModelRpcClient
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import get_exception_traceback
+from sglang.utils import get_exception_traceback
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
@@ -30,7 +30,7 @@ class RouterManager:
         self.recv_reqs = []
 
         # Init some configs
-        self.extend_dependency_time = GLOBAL_BACKEND_CONFIG.extend_dependency_time
+        self.request_dependency_time = global_config.request_dependency_time
 
     async def loop_for_forward(self):
         while True:
@@ -46,9 +46,9 @@ class RouterManager:
             if len(out_pyobjs) != 0:
                 has_finished = any([obj.finished for obj in out_pyobjs])
                 if has_finished:
-                    if self.extend_dependency_time > 0:
+                    if self.request_dependency_time > 0:
                         slept = True
-                        await asyncio.sleep(self.extend_dependency_time)
+                        await asyncio.sleep(self.request_dependency_time)
 
             if not slept:
                 await asyncio.sleep(0.0006)
@@ -60,9 +60,7 @@ class RouterManager:
 
 
 def start_router_process(
-    server_args: ServerArgs,
-    port_args: PortArgs,
-    pipe_writer,
+    server_args: ServerArgs, port_args: PortArgs, pipe_writer, model_overide_args
 ):
     logging.basicConfig(
         level=getattr(logging, server_args.log_level.upper()),
@@ -70,7 +68,7 @@ def start_router_process(
     )
 
     try:
-        model_client = ModelRpcClient(server_args, port_args)
+        model_client = ModelRpcClient(server_args, port_args, model_overide_args)
         router = RouterManager(model_client, port_args)
     except Exception:
         pipe_writer.send(get_exception_traceback())