sglang 0.1.16__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,6 +3,7 @@ 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
@@ -27,14 +28,13 @@ 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):
 
-        if self.text is None:
-            assert self.input_ids is not None, "Either text or input_ids should be provided"
-        else:
-            assert self.input_ids is None, "Either text or input_ids should be provided"
+        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)
@@ -69,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
@@ -105,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
@@ -127,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,6 +19,7 @@ class FinishReason(IntEnum):
     EOS_TOKEN = auto()
     LENGTH = auto()
     STOP_STR = auto()
+    ABORT = auto()
 
     @staticmethod
     def to_str(reason):
@@ -28,20 +29,22 @@ class FinishReason(IntEnum):
             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.
@@ -63,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
@@ -73,70 +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
-        if self.input_text is None:
-            # TODO(lmzheng): This can be wrong. Check with Liangsheng.
-            self.input_text = self.tokenizer.decode(self.input_ids)
-        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))
-
-        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
-        )
+        return (" " if first_token.startswith("▁") else "") + self.tokenizer.decode(ids)
 
-        # 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
@@ -155,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
@@ -333,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,
@@ -353,18 +376,27 @@ class Batch:
             ][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
+
+            # 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))]
 
@@ -394,6 +426,18 @@ class Batch:
                     # 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

@@ -8,7 +8,7 @@ import zmq.asyncio
 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())