sglang 0.1.21__py3-none-any.whl → 0.1.22__py3-none-any.whl

sglang/srt/managers/controller/tp_worker.py

@@ -1,15 +1,14 @@
 """A tensor parallel worker."""
 
-import asyncio
 import logging
+import multiprocessing
+import pickle
 import time
 import warnings
-from concurrent.futures import ThreadPoolExecutor
 from typing import List, Optional
 
-import rpyc
 import torch
-from rpyc.utils.classic import obtain
+import torch.distributed as dist
 
 from sglang.global_config import global_config
 from sglang.srt.constrained.fsm_cache import FSMCache
@@ -32,13 +31,11 @@ from sglang.srt.managers.io_struct import (
     TokenizedGenerateReqInput,
 )
 from sglang.srt.model_config import ModelConfig
-from sglang.srt.server_args import ModelPortArgs, ServerArgs
+from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
-    connect_rpyc_service,
     get_int_token_logit_bias,
     is_multimodal_model,
     set_random_seed,
-    start_rpyc_service_process,
     suppress_other_loggers,
 )
 from sglang.utils import get_exception_traceback
@@ -52,10 +49,9 @@ class ModelTpServer:
         gpu_id: int,
         tp_rank: int,
         server_args: ServerArgs,
-        model_port_args: ModelPortArgs,
+        nccl_port: int,
         model_overide_args: dict,
     ):
-        server_args, model_port_args = obtain(server_args), obtain(model_port_args)
         suppress_other_loggers()
 
         # Copy arguments
@@ -79,7 +75,7 @@ class ModelTpServer:
             gpu_id=gpu_id,
             tp_rank=tp_rank,
             tp_size=server_args.tp_size,
-            nccl_port=model_port_args.nccl_port,
+            nccl_port=nccl_port,
             server_args=server_args,
         )
 
@@ -178,9 +174,6 @@ class ModelTpServer:
         self.new_token_ratio_recovery = global_config.new_token_ratio_recovery
 
     def exposed_step(self, recv_reqs):
-        if not isinstance(recv_reqs, list):
-            recv_reqs = obtain(recv_reqs)
-
         try:
             # Recv requests
             for recv_req in recv_reqs:
@@ -228,23 +221,7 @@ class ModelTpServer:
 
                     # Print stats
                     if self.tp_rank == 0 and self.decode_forward_ct % 40 == 0:
-                        num_used = self.max_total_num_tokens - (
-                            self.token_to_kv_pool.available_size()
-                            + self.tree_cache.evictable_size()
-                        )
-                        throughput = self.num_generated_tokens / (
-                            time.time() - self.last_stats_tic
-                        )
-                        self.num_generated_tokens = 0
-                        self.last_stats_tic = time.time()
-                        logger.info(
-                            f"[gpu_id={self.gpu_id}] Decode batch. "
-                            f"#running-req: {len(self.running_batch.reqs)}, "
-                            f"#token: {num_used}, "
-                            f"token usage: {num_used / self.max_total_num_tokens:.2f}, "
-                            f"gen throughput (token/s): {throughput:.2f}, "
-                            f"#queue-req: {len(self.forward_queue)}"
-                        )
+                        self.print_stats()
 
                     if self.running_batch.is_empty():
                         self.running_batch = None
@@ -253,17 +230,34 @@ class ModelTpServer:
                     if self.out_pyobjs and self.running_batch.has_stream():
                         break
             else:
-                # Check the available size
-                available_size = (
-                    self.token_to_kv_pool.available_size()
-                    + self.tree_cache.evictable_size()
-                )
-                if available_size != self.max_total_num_tokens:
-                    warnings.warn(
-                        "Warning: "
-                        f"available_size={available_size}, max_total_num_tokens={self.max_total_num_tokens}\n"
-                        "KV cache pool leak detected!"
-                    )
+                self.check_memory()
+
+    def print_stats(self):
+        num_used = self.max_total_num_tokens - (
+            self.token_to_kv_pool.available_size() + self.tree_cache.evictable_size()
+        )
+        throughput = self.num_generated_tokens / (time.time() - self.last_stats_tic)
+        self.num_generated_tokens = 0
+        self.last_stats_tic = time.time()
+        logger.info(
+            f"[gpu_id={self.gpu_id}] Decode batch. "
+            f"#running-req: {len(self.running_batch.reqs)}, "
+            f"#token: {num_used}, "
+            f"token usage: {num_used / self.max_total_num_tokens:.2f}, "
+            f"gen throughput (token/s): {throughput:.2f}, "
+            f"#queue-req: {len(self.forward_queue)}"
+        )
+
+    def check_memory(self):
+        available_size = (
+            self.token_to_kv_pool.available_size() + self.tree_cache.evictable_size()
+        )
+        if available_size != self.max_total_num_tokens:
+            warnings.warn(
+                "Warning: "
+                f"available_size={available_size}, max_total_num_tokens={self.max_total_num_tokens}\n"
+                "KV cache pool leak detected!"
+            )
 
     def handle_generate_request(
         self,
@@ -310,6 +304,12 @@ class ModelTpServer:
             self.model_config.context_len - 1 - len(req.origin_input_ids),
             self.max_total_num_tokens - 128 - len(req.origin_input_ids),
         )
+        if req.sampling_params.max_new_tokens < 0:
+            req.origin_input_ids = req.origin_input_ids[
+                : self.max_total_num_tokens - 128
+            ]
+            logger.error("Request longer than memory pool size, truncated!!!")
+
         self.forward_queue.append(req)
 
     def get_new_prefill_batch(self) -> Optional[Batch]:
@@ -343,7 +343,8 @@ class ModelTpServer:
         if self.running_batch:
             available_size -= sum(
                 [
-                    (r.sampling_params.max_new_tokens - len(r.output_ids)) * self.new_token_ratio
+                    (r.sampling_params.max_new_tokens - len(r.output_ids))
+                    * self.new_token_ratio
                     for r in self.running_batch.reqs
                 ]
             )
@@ -365,7 +366,9 @@ class ModelTpServer:
                     req.image_offset += 1
 
             if (
-                req.extend_input_len + req.sampling_params.max_new_tokens + new_batch_total_tokens
+                req.extend_input_len
+                + req.sampling_params.max_new_tokens
+                + new_batch_total_tokens
                 < available_size
                 and (
                     req.extend_input_len + new_batch_input_tokens
@@ -377,7 +380,9 @@ class ModelTpServer:
                 available_size += delta
 
                 if not (
-                    req.extend_input_len + req.sampling_params.max_new_tokens + new_batch_total_tokens
+                    req.extend_input_len
+                    + req.sampling_params.max_new_tokens
+                    + new_batch_total_tokens
                     < available_size
                 ):
                     # Undo locking
@@ -419,12 +424,6 @@ class ModelTpServer:
                 f"#running-req: {running_bs}, "
                 f"#queue-req: {len(self.forward_queue) - len(can_run_list)}"
             )
-            # logger.debug(
-            #    f"fsm_cache_hit_rate: {100.0 * self.regex_fsm_cache.get_cache_hit_rate():.2f}%. "
-            #    f"fsm_cache_avg_init_time: {self.regex_fsm_cache.get_avg_init_time():.2f}s. "
-            #    f"ff_cache_hit_rate: {100.0 * self.jump_forward_cache.get_cache_hit_rate():.2f}%. "
-            #    f"ff_cache_avg_init_time: {self.jump_forward_cache.get_avg_init_time():.2f}s. "
-            # )
 
         # Return the new batch
         new_batch = Batch.init_new(
@@ -445,7 +444,7 @@ class ModelTpServer:
         # Forward and sample the next tokens
         if batch.extend_num_tokens != 0:
             output = self.model_runner.forward(batch, ForwardMode.EXTEND)
-            next_token_ids, _ = batch.sample(output.next_token_logits)
+            next_token_ids = batch.sample(output.next_token_logits)
 
             # Move logprobs to cpu
             if output.next_token_logprobs is not None:
@@ -568,7 +567,7 @@ class ModelTpServer:
 
         # Forward and sample the next tokens
         output = self.model_runner.forward(batch, ForwardMode.DECODE)
-        next_token_ids, _ = batch.sample(output.next_token_logits)
+        next_token_ids = batch.sample(output.next_token_logits)
 
         # Move logprobs to cpu
         if output.next_token_logprobs is not None:
@@ -596,9 +595,10 @@ class ModelTpServer:
 
     def handle_finished_requests(self, batch: Batch):
         output_rids = []
+        output_vids = []
         decoded_texts = []
-        surr_output_ids = []
-        read_output_ids = []
+        output_read_ids = []
+        output_read_offsets = []
         output_skip_special_tokens = []
         output_spaces_between_special_tokens = []
         output_meta_info = []
@@ -621,10 +621,11 @@ class ModelTpServer:
                 )
             ):
                 output_rids.append(req.rid)
+                output_vids.append(req.vid)
                 decoded_texts.append(req.decoded_text)
-                surr_ids, read_ids, _ = req.init_detokenize_incrementally()
-                surr_output_ids.append(surr_ids)
-                read_output_ids.append(read_ids)
+                read_ids, read_offset = req.init_incremental_detokenize()
+                output_read_ids.append(read_ids)
+                output_read_offsets.append(read_offset)
                 output_skip_special_tokens.append(
                     req.sampling_params.skip_special_tokens
                 )
@@ -660,9 +661,10 @@ class ModelTpServer:
             self.out_pyobjs.append(
                 BatchTokenIDOut(
                     output_rids,
+                    output_vids,
                     decoded_texts,
-                    surr_output_ids,
-                    read_output_ids,
+                    output_read_ids,
+                    output_read_offsets,
                     output_skip_special_tokens,
                     output_spaces_between_special_tokens,
                     output_meta_info,
@@ -727,87 +729,74 @@ class ModelTpServer:
                     break
 
 
-class ModelTpService(rpyc.Service):
-    exposed_ModelTpServer = ModelTpServer
-
-
-class ModelTpClient:
-    def __init__(
-        self,
-        gpu_ids: List[int],
-        server_args: ServerArgs,
-        model_port_args: ModelPortArgs,
-        model_overide_args,
-    ):
-        server_args, model_port_args = obtain(server_args), obtain(model_port_args)
-        self.tp_size = server_args.tp_size
+def run_tp_server(
+    gpu_id: int,
+    tp_rank: int,
+    server_args: ServerArgs,
+    nccl_port: int,
+    model_overide_args: dict,
+):
+    """Run a tensor parallel server."""
+    try:
+        model_server = ModelTpServer(
+            gpu_id,
+            tp_rank,
+            server_args,
+            nccl_port,
+            model_overide_args,
+        )
+        tp_cpu_group = model_server.model_runner.tp_group.cpu_group
+
+        while True:
+            recv_reqs = broadcast_recv_input(None, tp_rank, tp_cpu_group)
+            model_server.exposed_step(recv_reqs)
+    except Exception:
+        logger.error("Exception in run_tp_server:\n" + get_exception_traceback())
+        raise
+
+
+def launch_tp_servers(
+    gpu_ids, tp_rank_range, server_args, nccl_port, model_overide_args
+):
+    """Launch multiple tensor parallel servers."""
+    procs = []
+    for i in tp_rank_range:
+        proc = multiprocessing.Process(
+            target=run_tp_server,
+            args=(gpu_ids[i], i, server_args, nccl_port, model_overide_args),
+        )
+        proc.start()
+        procs.append(proc)
 
-        if self.tp_size * server_args.dp_size == 1:
-            # Init model
-            assert len(gpu_ids) == 1
-            self.model_server = ModelTpService().exposed_ModelTpServer(
-                gpu_ids[0],
-                0,
-                server_args,
-                model_port_args,
-                model_overide_args,
-            )
+    return procs
 
-            # Wrap functions
-            def async_wrap(f):
-                async def _func(*args, **kwargs):
-                    return f(*args, **kwargs)
 
-                return _func
+def broadcast_recv_input(data, rank, dist_group):
+    """Broadcast inputs from rank=0 to all other ranks with torch.dist backend."""
 
-            self.step = async_wrap(self.model_server.exposed_step)
+    if rank == 0:
+        if len(data) == 0:
+            tensor_size = torch.tensor([0], dtype=torch.long)
+            dist.broadcast(tensor_size, src=0, group=dist_group)
         else:
-            with ThreadPoolExecutor(self.tp_size) as executor:
-                # Launch model processes
-                if server_args.nnodes == 1:
-                    self.procs = list(
-                        executor.map(
-                            lambda args: start_rpyc_service_process(*args),
-                            [
-                                (ModelTpService, p)
-                                for p in model_port_args.model_tp_ports
-                            ],
-                        )
-                    )
-                    addrs = [("localhost", p) for p in model_port_args.model_tp_ports]
-                else:
-                    addrs = [
-                        (ip, port)
-                        for ip, port in zip(
-                            model_port_args.model_tp_ips, model_port_args.model_tp_ports
-                        )
-                    ]
-
-                self.model_services = list(
-                    executor.map(lambda args: connect_rpyc_service(*args), addrs)
-                )
-
-                # Init model
-                def init_model(i):
-                    return self.model_services[i].ModelTpServer(
-                        gpu_ids[i],
-                        i,
-                        server_args,
-                        model_port_args,
-                        model_overide_args,
-                    )
-
-                self.model_servers = list(executor.map(init_model, range(self.tp_size)))
-
-            # Wrap functions
-            def async_wrap(func_name):
-                fs = [rpyc.async_(getattr(m, func_name)) for m in self.model_servers]
-
-                async def _func(*args, **kwargs):
-                    tasks = [f(*args, **kwargs) for f in fs]
-                    await asyncio.gather(*[asyncio.to_thread(t.wait) for t in tasks])
-                    return obtain(tasks[0].value)
-
-                return _func
-
-            self.step = async_wrap("step")
+            serialized_data = pickle.dumps(data)
+            size = len(serialized_data)
+            tensor_data = torch.ByteTensor(list(serialized_data))
+            tensor_size = torch.tensor([size], dtype=torch.long)
+
+            dist.broadcast(tensor_size, src=0, group=dist_group)
+            dist.broadcast(tensor_data, src=0, group=dist_group)
+    else:
+        tensor_size = torch.tensor([0], dtype=torch.long)
+        dist.broadcast(tensor_size, src=0, group=dist_group)
+        size = tensor_size.item()
+
+        if size == 0:
+            return []
+
+        tensor_data = torch.empty(size, dtype=torch.uint8)
+        dist.broadcast(tensor_data, src=0, group=dist_group)
+
+        serialized_data = bytes(tensor_data.tolist())
+        data = pickle.loads(serialized_data)
+        return data

sglang/srt/managers/detokenizer_manager.py

@@ -1,7 +1,9 @@
 """DetokenizerManager is a process that detokenizes the token ids."""
 
 import asyncio
+import dataclasses
 import inspect
+from typing import List
 
 import uvloop
 import zmq
@@ -16,6 +18,15 @@ from sglang.utils import find_printable_text, get_exception_traceback, graceful_
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
+@dataclasses.dataclass
+class DecodeStatus:
+    vid: int
+    decoded_text: str
+    decode_ids: List[int]
+    surr_offset: int
+    read_offset: int
+
+
 class DetokenizerManager:
     def __init__(
         self,
@@ -35,19 +46,43 @@ class DetokenizerManager:
             trust_remote_code=server_args.trust_remote_code,
         )
 
+        self.decode_status = {}
+
     async def handle_loop(self):
         while True:
             recv_obj: BatchTokenIDOut = await self.recv_from_router.recv_pyobj()
             assert isinstance(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:
+                    s = self.decode_status[rid]
+                    s.decode_ids = recv_obj.decode_ids[i]
+
+                read_ids.append(s.decode_ids[s.surr_offset :])
+                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(
-                recv_obj.surr_output_ids,
+                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(
-                recv_obj.read_output_ids,
+                read_ids,
                 skip_special_tokens=recv_obj.skip_special_tokens[0],
                 spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
             )
@@ -55,11 +90,20 @@ class DetokenizerManager:
             # Trim stop str
             # TODO(lmzheng): handle the case where multiple stop strs are hit
             output_strs = []
-            for i in range(len(recv_obj.rids)):
+            for i in range(bs):
+                s = self.decode_status[recv_obj.rids[i]]
                 new_text = read_texts[i][len(surr_texts[i]) :]
                 if recv_obj.finished_reason[i] is None:
-                    new_text = find_printable_text(new_text)
-                output_strs.append(recv_obj.decoded_texts[i] + new_text)
+                    # 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)
+
+                output_strs.append(s.decoded_text + new_text)
 
                 if isinstance(recv_obj.finished_reason[i], FINISH_MATCHED_STR):
                     pos = output_strs[i].find(recv_obj.finished_reason[i].matched)

sglang/srt/managers/io_struct.py

@@ -13,25 +13,26 @@ from sglang.srt.sampling_params import SamplingParams
 
 @dataclass
 class GenerateReqInput:
-    # The input prompt
+    # The input prompt. It can be a single prompt or a batch of prompts.
     text: Optional[Union[List[str], str]] = None
-    # The token ids for text; one can either specify text or input_ids
+    # 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
+    # The image input. It can be a file name, a url, or base64 encoded string.
+    # See also python/sglang/srt/utils.py:load_image.
     image_data: Optional[Union[List[str], str]] = None
-    # The sampling_params
+    # The sampling_params.
     sampling_params: Union[List[Dict], Dict] = None
-    # The request id
+    # The request id.
     rid: Optional[Union[List[str], str]] = None
-    # Whether to return logprobs
+    # Whether to return logprobs.
     return_logprob: Optional[Union[List[bool], bool]] = None
-    # The start location of the prompt for return_logprob
+    # The start location of the prompt for return_logprob.
     logprob_start_len: Optional[Union[List[int], int]] = None
-    # The number of top logprobs to return
+    # The number of top logprobs to return.
     top_logprobs_num: Optional[Union[List[int], int]] = None
-    # Whether to detokenize tokens in logprobs
+    # Whether to detokenize tokens in logprobs.
     return_text_in_logprobs: bool = False
-    # Whether to stream output
+    # Whether to stream output.
     stream: bool = False
 
     def post_init(self):
@@ -39,11 +40,13 @@ class GenerateReqInput:
             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)
+        if self.sampling_params.get("n", 1) != 1:
+            is_single = False
         else:
-            is_single = isinstance(self.input_ids[0], int)
+            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:
@@ -58,7 +61,22 @@ class GenerateReqInput:
             if self.top_logprobs_num is None:
                 self.top_logprobs_num = 0
         else:
-            num = len(self.text) if self.text is not None else len(self.input_ids)
+
+            parallel_sample_num = self.sampling_params.get("n", 1)
+
+            if parallel_sample_num != 1:
+                # parallel sampling +1 represents the original prefill stage
+                num = parallel_sample_num + 1
+                if isinstance(self.text, List):
+                    ## suppot batch operation
+                    self.batch_size = len(self.text)
+                    num = num * len(self.text)
+                else:
+                    self.batch_size = 1
+            else:
+                ## support select operation
+                num = len(self.text) if self.text is not None else len(self.input_ids)
+                self.batch_size = num
 
             if self.image_data is None:
                 self.image_data = [None] * num
@@ -110,9 +128,10 @@ class TokenizedGenerateReqInput:
 @dataclass
 class BatchTokenIDOut:
     rids: List[str]
+    vids: List[int]
     decoded_texts: List[str]
-    surr_output_ids: List[List[int]]
-    read_output_ids: List[List[int]]
+    decode_ids: List[int]
+    read_offsets: List[int]
     skip_special_tokens: List[bool]
     spaces_between_special_tokens: List[bool]
     meta_info: List[Dict]