sglang 0.1.16__py3-none-any.whl → 0.1.18__py3-none-any.whl

sglang/srt/managers/{router → controller}/infer_batch.py

@@ -1,3 +1,6 @@
+"""Meta data for requests and batches"""
+
+import warnings
 from dataclasses import dataclass
 from enum import IntEnum, auto
 from typing import List
@@ -5,9 +8,13 @@ from typing import List
 import numpy as np
 import torch
 
-from sglang.srt.managers.router.radix_cache import RadixCache
+from sglang.srt.constrained import RegexGuide
+from sglang.srt.constrained.jump_forward import JumpForwardMap
+from sglang.srt.managers.controller.radix_cache import RadixCache
 from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
 
+INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
+
 
 class ForwardMode(IntEnum):
     PREFILL = auto()
@@ -15,33 +22,62 @@ class ForwardMode(IntEnum):
     DECODE = auto()
 
 
-class FinishReason(IntEnum):
-    EOS_TOKEN = auto()
-    LENGTH = auto()
-    STOP_STR = auto()
+class BaseFinishReason:
+    def __init__(self, is_error: bool = False):
+        self.is_error = is_error
+
+    def __str__(self):
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class FINISH_MATCHED_TOKEN(BaseFinishReason):
+    def __init__(self, matched: int | List[int]):
+        super().__init__()
+        self.matched = matched
+
+    def __str__(self) -> str:
+        return f"FINISH_MATCHED_TOKEN: {self.matched}"
+
+
+class FINISH_LENGTH(BaseFinishReason):
+    def __init__(self, length: int):
+        super().__init__()
+        self.length = length
+
+    def __str__(self) -> str:
+        return f"FINISH_LENGTH: {self.length}"
+
+
+class FINISH_MATCHED_STR(BaseFinishReason):
+    def __init__(self, matched: str):
+        super().__init__()
+        self.matched = matched
+
+    def __str__(self) -> str:
+        return f"FINISH_MATCHED_STR: {self.matched}"
 
-    @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"
-        else:
-            return None
+
+class FINISH_ABORT(BaseFinishReason):
+    def __init__(self):
+        super().__init__(is_error=True)
+
+    def __str__(self) -> str:
+        return "FINISH_ABORT"
 
 
 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.output_ids = []
+        self.origin_input_text = origin_input_text
+        self.origin_input_ids_unpadded = origin_input_ids  # Before image padding
+        self.origin_input_ids = origin_input_ids
+        self.output_ids = []  # Each decode stage's output ids
+        self.input_ids = None  # input_ids = origin_input_ids + 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)
+        # For incremental decode
+        self.decoded_text = ""
+        self.surr_offset = None  # Surrounding offset to defeat the cleanup algorithm
+        self.read_offset = None
 
         # The number of decoded tokens for token usage report. Note that
         # this does not include the jump forward tokens.
@@ -57,12 +93,12 @@ class Req:
         self.sampling_params = None
         self.stream = False
 
-        # Check finish
         self.tokenizer = None
-        self.finished = False
-        self.finish_reason = None
-        self.hit_stop_str = None
 
+        # Check finish
+        self.finished_reason = None
+
+        # Prefix info
         self.extend_input_len = 0
         self.prefix_indices = []
         self.last_node = None
@@ -73,80 +109,81 @@ 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 = ""
+        self.regex_fsm: RegexGuide = None
+        self.regex_fsm_state: int = 0
+        self.jump_forward_map: JumpForwardMap = None
+
+    # whether request reached finished condition
+    def finished(self) -> bool:
+        return self.finished_reason is not None
+
+    # Based on https://github.com/vllm-project/vllm/blob/7a64d24aad69e4d2548aa0bf528d9fe63428ab01/vllm/transformers_utils/detokenizer.py#L194-L313
+    def init_detokenize_incrementally(self):
+        first_iter = self.surr_offset is None or self.read_offset is None
+
+        if first_iter:
+            self.read_offset = len(self.origin_input_ids_unpadded)
+            self.surr_offset = max(
+                self.read_offset - INIT_INCREMENTAL_DETOKENIZATION_OFFSET, 0
+            )
 
-    def max_new_tokens(self):
-        return self.sampling_params.max_new_tokens
+        all_ids = self.origin_input_ids_unpadded + self.output_ids
+        surr_ids = all_ids[self.surr_offset : self.read_offset]
+        read_ids = all_ids[self.surr_offset :]
 
-    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])
-        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)
-            )
+        return surr_ids, read_ids, len(all_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))
+    def detokenize_incrementally(self, inplace: bool = True):
+        surr_ids, read_ids, num_all_tokens = self.init_detokenize_incrementally()
 
-        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
+        surr_text = self.tokenizer.decode(
+            surr_ids,
+            skip_special_tokens=self.sampling_params.skip_special_tokens,
+            spaces_between_special_tokens=self.sampling_params.spaces_between_special_tokens,
         )
-        self.regex_fsm_state = next_state
-        self.output_and_jump_forward_str = (
-            self.output_and_jump_forward_str + old_output_str + jump_forward_str
+        new_text = self.tokenizer.decode(
+            read_ids,
+            skip_special_tokens=self.sampling_params.skip_special_tokens,
+            spaces_between_special_tokens=self.sampling_params.spaces_between_special_tokens,
         )
 
-        # print(f"Output and jump forward str:\n{self.output_and_jump_forward_str}")
-        # print("*" * 100)
+        if len(new_text) > len(surr_text) and not new_text.endswith("�"):
+            new_text = new_text[len(surr_text) :]
+            if inplace:
+                self.decoded_text += new_text
+                self.surr_offset = self.read_offset
+                self.read_offset = num_all_tokens
+
+            return True, new_text
+
+        return False, ""
+
+    def max_new_tokens(self):
+        return self.sampling_params.max_new_tokens
 
     def check_finished(self):
-        if self.finished:
+        if self.finished():
             return
 
         if len(self.output_ids) >= self.sampling_params.max_new_tokens:
-            self.finished = True
-            self.finish_reason = FinishReason.LENGTH
+            self.finished_reason = FINISH_LENGTH(len(self.output_ids))
             return
 
         if (
             self.output_ids[-1] == self.tokenizer.eos_token_id
-            and self.sampling_params.ignore_eos == False
+            and not self.sampling_params.ignore_eos
         ):
-            self.finished = True
-            self.finish_reason = FinishReason.EOS_TOKEN
+            self.finished_reason = FINISH_MATCHED_TOKEN(
+                matched=self.tokenizer.eos_token_id
+            )
             return
 
         if len(self.sampling_params.stop_strs) > 0:
@@ -155,14 +192,62 @@ class Req:
             )
 
             for stop_str in self.sampling_params.stop_strs:
-                if stop_str in tail_str:
-                    self.finished = True
-                    self.finish_reason = FinishReason.STOP_STR
-                    self.hit_stop_str = stop_str
+                if stop_str in tail_str or stop_str in self.decoded_text:
+                    self.finished_reason = FINISH_MATCHED_STR(matched=stop_str)
                     return
 
+    def jump_forward_and_retokenize(self, jump_forward_str, next_state):
+        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.decoded_text + jump_forward_str
+        all_ids = self.tokenizer.encode(all_text)
+        prompt_tokens = len(self.origin_input_ids_unpadded)
+
+        if all_ids[prompt_tokens - 1] != self.origin_input_ids_unpadded[-1]:
+            # TODO(lsyin): fix token fusion
+            warnings.warn(
+                "Token fusion between input and output, try to avoid this by removing the space at the end of the input."
+            )
+            return False
+
+        old_output_ids = self.output_ids
+        self.output_ids = all_ids[prompt_tokens:]
+        self.decoded_text = self.decoded_text + jump_forward_str
+        self.surr_offset = prompt_tokens
+        self.read_offset = len(all_ids)
+
+        # NOTE: A trick to reduce the surrouding tokens decoding overhead
+        for i in range(0, INIT_INCREMENTAL_DETOKENIZATION_OFFSET):
+            surr_text_ = self.tokenizer.decode(
+                all_ids[self.read_offset - i : self.read_offset]
+            )
+            if not surr_text_.endswith("�"):
+                self.surr_offset = self.read_offset - i
+                break
+
+        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_output_ids):
+                if old_id == self.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.output_ids) - k
+
+        return True
+
     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
@@ -218,6 +303,10 @@ class Batch:
     def is_empty(self):
         return len(self.reqs) == 0
 
+    # whether batch has at least 1 streaming request
+    def has_stream(self) -> bool:
+        return any(r.stream for r in self.reqs)
+
     def prepare_for_extend(self, vocab_size: int, int_token_logit_bias: torch.Tensor):
         device = "cuda"
         bs = len(self.reqs)
@@ -333,8 +422,12 @@ 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)),
+            key=lambda i: (
+                len(self.reqs[i].output_ids),
+                -len(self.reqs[i].origin_input_ids),
+            ),
             reverse=True,
         )
 
@@ -353,18 +446,22 @@ 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)
+
             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))]
 
@@ -372,18 +469,54 @@ class Batch:
 
         for i, req in enumerate(self.reqs):
             if req.jump_forward_map is not None:
-                res = req.jump_forward_map.jump_forward(req.regex_fsm_state)
-                if res is not None:
-                    jump_forward_str, next_state = res
-                    if len(jump_forward_str) <= 1:
+                jump_forward_bytes = req.jump_forward_map.jump_forward_byte(
+                    req.regex_fsm_state
+                )
+                if jump_forward_bytes is not None and len(jump_forward_bytes) > 1:
+                    suffix_bytes = []
+                    continuation_range = range(0x80, 0xC0)
+                    cur_state = req.regex_fsm_state
+                    while (
+                        len(jump_forward_bytes)
+                        and jump_forward_bytes[0][0] in continuation_range
+                    ):
+                        # continuation bytes
+                        byte_edge = jump_forward_bytes.pop(0)
+                        suffix_bytes.append(byte_edge[0])
+                        cur_state = byte_edge[1]
+
+                    suffix_tokens = [f"<0x{hex(b)[2:].upper()}>" for b in suffix_bytes]
+                    suffix_ids = req.tokenizer.convert_tokens_to_ids(suffix_tokens)
+
+                    # Current ids, for cache and revert
+                    cur_all_ids = tuple(req.origin_input_ids + req.output_ids)[:-1]
+                    cur_output_ids = req.output_ids
+
+                    req.output_ids.extend(suffix_ids)
+                    decode_res, new_text = req.detokenize_incrementally(inplace=False)
+                    if not decode_res:
+                        req.output_ids = cur_output_ids
                         continue
 
-                    if req_pool_indices_cpu is None:
-                        req_pool_indices_cpu = self.req_pool_indices.tolist()
+                    (
+                        jump_forward_str,
+                        next_state,
+                    ) = req.jump_forward_map.jump_forward_symbol(cur_state)
+
+                    # Make the incrementally decoded text part of jump_forward_str
+                    # so that the UTF-8 will not corrupt
+                    jump_forward_str = new_text + jump_forward_str
+                    if not req.jump_forward_and_retokenize(
+                        jump_forward_str, next_state
+                    ):
+                        req.output_ids = cur_output_ids
+                        continue
 
                     # insert the old request into tree_cache
+                    if req_pool_indices_cpu is None:
+                        req_pool_indices_cpu = self.req_pool_indices.tolist()
                     self.tree_cache.cache_req(
-                        token_ids=tuple(req.input_ids + req.output_ids)[:-1],
+                        token_ids=cur_all_ids,
                         last_uncached_pos=len(req.prefix_indices),
                         req_pool_idx=req_pool_indices_cpu[i],
                     )
@@ -391,8 +524,17 @@ class Batch:
                     # 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)
@@ -515,7 +657,7 @@ class Batch:
                 if req.regex_fsm is not None:
                     allowed_mask.zero_()
                     allowed_mask[
-                        req.regex_fsm.allowed_token_ids(req.regex_fsm_state)
+                        req.regex_fsm.get_next_instruction(req.regex_fsm_state).tokens
                     ] = 1
                     logits[i].masked_fill_(~allowed_mask, float("-inf"))
 
@@ -534,7 +676,7 @@ class Batch:
             batch_next_token_ids_cpu = batch_next_token_ids.cpu().numpy()
             for i, req in enumerate(self.reqs):
                 if req.regex_fsm is not None:
-                    req.regex_fsm_state = req.regex_fsm.next_state(
+                    req.regex_fsm_state = req.regex_fsm.get_next_state(
                         req.regex_fsm_state, batch_next_token_ids_cpu[i]
                     )
 

sglang/srt/managers/controller/manager_multi.py

@@ -0,0 +1,191 @@
+"""
+A controller that manages multiple data parallel workers.
+Each data parallel worker can manage multiple tensor parallel workers.
+"""
+
+import asyncio
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from enum import Enum, auto
+from typing import Dict
+
+import zmq
+import zmq.asyncio
+
+from sglang.global_config import global_config
+from sglang.srt.managers.controller.dp_worker import (
+    DataParallelWorkerThread,
+    start_data_parallel_worker,
+)
+from sglang.srt.managers.io_struct import (
+    AbortReq,
+    FlushCacheReq,
+    TokenizedGenerateReqInput,
+)
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.utils import get_exception_traceback
+
+logger = logging.getLogger("srt.controller")
+
+
+class LoadBalanceMethod(Enum):
+    ROUND_ROBIN = auto()
+    SHORTEST_QUEUE = auto()
+
+    @classmethod
+    def from_str(cls, method: str):
+        method = method.upper()
+        try:
+            return cls[method]
+        except KeyError as exc:
+            raise ValueError(f"Invalid load balance method: {method}") from exc
+
+
+class Controller:
+    def __init__(
+        self,
+        load_balance_method: str,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        model_overide_args,
+    ):
+        self.load_balance_method = LoadBalanceMethod.from_str(load_balance_method)
+        self.server_args = server_args
+        self.port_args = port_args
+
+        if self.load_balance_method == LoadBalanceMethod.ROUND_ROBIN:
+            self.round_robin_counter = 0
+
+        self.dispatch_lookup = {
+            LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler,
+            LoadBalanceMethod.SHORTEST_QUEUE: self.shortest_queue_scheduler,
+        }
+        self.dispatching = self.dispatch_lookup[self.load_balance_method]
+
+        # Init communication
+        context = zmq.asyncio.Context()
+        self.recv_from_tokenizer = context.socket(zmq.PULL)
+        self.recv_from_tokenizer.bind(f"tcp://127.0.0.1:{port_args.router_port}")
+
+        # Init status
+        self.recv_reqs = []
+
+        # Start data parallel workers
+        self.workers: Dict[int, DataParallelWorkerThread] = {}
+        tp_size = server_args.tp_size
+
+        def start_dp_worker(i):
+            try:
+                gpu_ids = list(range(i * tp_size, (i + 1) * tp_size))
+                worker_thread = start_data_parallel_worker(
+                    server_args, port_args, model_overide_args, gpu_ids, i
+                )
+                self.workers[i] = worker_thread
+            except Exception:
+                logger.error(
+                    f"Failed to start local worker {i}\n{get_exception_traceback()}"
+                )
+
+        for i in range(server_args.dp_size):
+            start_dp_worker(i)
+
+        # Parallel launch is slower, probably due to the disk bandwidth limitations.
+        # with ThreadPoolExecutor(server_args.dp_size) as executor:
+        #     executor.map(start_dp_worker, range(server_args.dp_size))
+
+    def have_any_live_worker(self):
+        return any(worker_thread.liveness for worker_thread in self.workers.values())
+
+    def put_req_to_worker(self, worker_id, req):
+        self.workers[worker_id].request_queue.put(req)
+
+    async def round_robin_scheduler(self, input_requests):
+        available_workers = list(self.workers.keys())
+        for r in input_requests:
+            self.put_req_to_worker(available_workers[self.round_robin_counter], r)
+            self.round_robin_counter = (self.round_robin_counter + 1) % len(
+                available_workers
+            )
+        return
+
+    async def shortest_queue_scheduler(self, input_requests):
+        for r in input_requests:
+            worker = min(
+                self.workers, key=lambda w: self.workers[w].request_queue.qsize()
+            )
+            self.put_req_to_worker(worker, r)
+        return
+
+    async def remove_dead_workers(self):
+        for i in list(self.workers.keys()):
+            worker_thread = self.workers[i]
+            if not worker_thread.liveness:
+                worker_thread.join()
+                # move unsuccessful requests back to the queue
+                while not worker_thread.request_queue.empty():
+                    self.recv_reqs.append(worker_thread.request_queue.get())
+                del self.workers[i]
+                logger.info(f"Stale worker {i} removed")
+
+    async def loop_for_forward(self):
+        while True:
+            await self.remove_dead_workers()
+
+            if self.have_any_live_worker():
+                next_step_input = list(self.recv_reqs)
+                self.recv_reqs = []
+                if next_step_input:
+                    await self.dispatching(next_step_input)
+            # else:
+            #    logger.error("There is no live worker.")
+
+            await asyncio.sleep(global_config.wait_for_new_request_delay)
+
+    async def loop_for_recv_requests(self):
+        while True:
+            recv_req = await self.recv_from_tokenizer.recv_pyobj()
+            if isinstance(recv_req, FlushCacheReq):
+                # TODO(lsyin): apply more specific flushCacheReq
+                for worker_thread in self.workers.values():
+                    worker_thread.request_queue.put(recv_req)
+            elif isinstance(recv_req, TokenizedGenerateReqInput):
+                self.recv_reqs.append(recv_req)
+            elif isinstance(recv_req, AbortReq):
+                in_queue = False
+                for i, req in enumerate(self.recv_reqs):
+                    if req.rid == recv_req.rid:
+                        self.recv_reqs[i] = recv_req
+                        in_queue = True
+                        break
+                if not in_queue:
+                    # Send abort req to all TP groups
+                    for worker in list(self.workers.keys()):
+                        self.put_req_to_worker(worker, recv_req)
+            else:
+                logger.error(f"Invalid object: {recv_req}")
+
+
+def start_controller_process(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+    pipe_writer,
+    model_overide_args=None,
+):
+    logging.basicConfig(
+        level=getattr(logging, server_args.log_level.upper()),
+        format="%(message)s",
+    )
+
+    try:
+        controller = Controller(
+            server_args.load_balance_method, server_args, port_args, model_overide_args
+        )
+    except Exception:
+        pipe_writer.send(get_exception_traceback())
+        raise
+
+    pipe_writer.send("init ok")
+    loop = asyncio.get_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.create_task(controller.loop_for_recv_requests())
+    loop.run_until_complete(controller.loop_for_forward())