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

sglang/srt/managers/controller/tp_worker.py

@@ -1,9 +1,11 @@
+"""A tensor parallel worker."""
+
 import asyncio
 import logging
 import time
 import warnings
 from concurrent.futures import ThreadPoolExecutor
-from typing import List
+from typing import List, Optional
 
 import rpyc
 import torch
@@ -13,23 +15,30 @@ from sglang.global_config import global_config
 from sglang.srt.constrained.fsm_cache import FSMCache
 from sglang.srt.constrained.jump_forward import JumpForwardCache
 from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
+from sglang.srt.managers.controller.infer_batch import (
+    FINISH_ABORT,
+    BaseFinishReason,
+    Batch,
+    ForwardMode,
+    Req,
+)
+from sglang.srt.managers.controller.model_runner import ModelRunner
+from sglang.srt.managers.controller.radix_cache import RadixCache
+from sglang.srt.managers.controller.schedule_heuristic import ScheduleHeuristic
 from sglang.srt.managers.io_struct import (
     AbortReq,
     BatchTokenIDOut,
     FlushCacheReq,
     TokenizedGenerateReqInput,
 )
-from sglang.srt.managers.controller.infer_batch import BaseFinishReason, Batch, FINISH_ABORT, ForwardMode, Req
-from sglang.srt.managers.controller.model_runner import ModelRunner
-from sglang.srt.managers.controller.radix_cache import RadixCache
-from sglang.srt.managers.controller.schedule_heuristic import ScheduleHeuristic
 from sglang.srt.model_config import ModelConfig
 from sglang.srt.server_args import ModelPortArgs, ServerArgs
 from sglang.srt.utils import (
+    connect_rpyc_service,
     get_int_token_logit_bias,
     is_multimodal_model,
     set_random_seed,
-    start_rpyc_process,
+    start_rpyc_service_process,
     suppress_other_loggers,
 )
 from sglang.utils import get_exception_traceback
@@ -88,16 +97,16 @@ class ModelTpServer:
                 trust_remote_code=server_args.trust_remote_code,
             )
         self.max_total_num_tokens = self.model_runner.max_total_num_tokens
-        self.max_prefill_tokens = max(
-            self.model_config.context_len,
-            (
-                min(self.max_total_num_tokens // 6, 65536)
-                if server_args.max_prefill_tokens is None
-                else server_args.max_prefill_tokens
-            ),
+        self.max_prefill_tokens = (
+            4096
+            if server_args.max_prefill_tokens is None
+            else server_args.max_prefill_tokens
+        )
+        self.max_running_requests = (
+            self.max_total_num_tokens // 2
+            if server_args.max_running_requests is None
+            else server_args.max_running_requests
         )
-        self.max_running_requests = (self.max_total_num_tokens // 2
-            if server_args.max_running_requests is None else server_args.max_running_requests)
         self.int_token_logit_bias = torch.tensor(
             get_int_token_logit_bias(self.tokenizer, self.model_config.vocab_size)
         )
@@ -108,7 +117,7 @@ class ModelTpServer:
             f"[gpu_id={self.gpu_id}] "
             f"max_total_num_tokens={self.max_total_num_tokens}, "
             f"max_prefill_tokens={self.max_prefill_tokens}, "
-            f"context_len={self.model_config.context_len}, "
+            f"context_len={self.model_config.context_len}"
         )
         if self.tp_rank == 0:
             logger.info(
@@ -242,7 +251,7 @@ class ModelTpServer:
                         self.running_batch = None
                         break
 
-                    if self.out_pyobjs and self.running_batch.reqs[0].stream:
+                    if self.out_pyobjs and self.running_batch.has_stream():
                         break
             else:
                 # Check the available size
@@ -271,13 +280,14 @@ class ModelTpServer:
                 (recv_req.image_hash >> 64) % self.model_config.vocab_size,
             ]
             req.image_size = recv_req.image_size
-            req.origin_input_ids, req.image_offset = (
-                self.model_runner.model.pad_input_ids(
-                    req.origin_input_ids_unpadded,
-                    req.pad_value,
-                    req.pixel_values.shape,
-                    req.image_size,
-                )
+            (
+                req.origin_input_ids,
+                req.image_offset,
+            ) = self.model_runner.model.pad_input_ids(
+                req.origin_input_ids_unpadded,
+                req.pad_value,
+                req.pixel_values.shape,
+                req.image_size,
             )
         req.sampling_params = recv_req.sampling_params
         req.return_logprob = recv_req.return_logprob
@@ -303,7 +313,7 @@ class ModelTpServer:
         )
         self.forward_queue.append(req)
 
-    def get_new_fill_batch(self):
+    def get_new_fill_batch(self) -> Optional[Batch]:
         if (
             self.running_batch is not None
             and len(self.running_batch.reqs) > self.max_running_requests
@@ -312,10 +322,7 @@ class ModelTpServer:
 
         # Compute matched prefix length
         for req in self.forward_queue:
-            assert (
-                len(req.output_ids) == 0
-            ), "The output ids should be empty when prefilling"
-            req.input_ids = req.origin_input_ids + req.prev_output_ids
+            req.input_ids = req.origin_input_ids + req.output_ids
             prefix_indices, last_node = self.tree_cache.match_prefix(req.input_ids)
             if req.return_logprob:
                 prefix_indices = prefix_indices[: req.logprob_start_len]
@@ -361,8 +368,11 @@ class ModelTpServer:
             if (
                 req.extend_input_len + req.max_new_tokens() + new_batch_total_tokens
                 < available_size
-                and req.extend_input_len + new_batch_input_tokens
-                < self.max_prefill_tokens
+                and (
+                    req.extend_input_len + new_batch_input_tokens
+                    <= self.max_prefill_tokens
+                    or len(can_run_list) == 0
+                )
             ):
                 delta = self.tree_cache.inc_lock_ref(req.last_node)
                 available_size += delta
@@ -401,7 +411,7 @@ class ModelTpServer:
                 self.tree_cache_metrics["hit"] / self.tree_cache_metrics["total"]
             )
             logger.info(
-                f"[gpu_id={self.gpu_id}] Prefil batch. "
+                f"[gpu_id={self.gpu_id}] Prefill batch. "
                 f"#new-seq: {len(can_run_list)}, "
                 f"#new-token: {new_batch_input_tokens}, "
                 f"#cached-token: {hit_tokens}, "
@@ -432,97 +442,93 @@ class ModelTpServer:
             self.model_config.vocab_size, self.int_token_logit_bias
         )
 
+        # Forward and sample the next tokens
         if batch.extend_num_tokens != 0:
-            # Forward
-            logits, (
-                prefill_token_logprobs,
-                normalized_prompt_logprobs,
-                prefill_top_logprobs,
-                decode_top_logprobs,
-                last_logprobs,
-            ) = self.model_runner.forward(batch, ForwardMode.EXTEND)
-            if prefill_token_logprobs is not None:
-                prefill_token_logprobs = prefill_token_logprobs.tolist()
-                normalized_prompt_logprobs = normalized_prompt_logprobs.tolist()
-
-            next_token_ids, _ = batch.sample(logits)
-
-            # Only transfer the selected logprobs of the next token to CPU to reduce overhead.
-            if last_logprobs is not None:
-                last_token_logprobs = last_logprobs[
-                    torch.arange(len(batch.reqs), device=next_token_ids.device),
+            output = self.model_runner.forward(batch, ForwardMode.EXTEND)
+            next_token_ids, _ = batch.sample(output.next_token_logits)
+
+            # Move logprobs to cpu
+            if output.next_token_logprobs is not None:
+                output.next_token_logprobs = output.next_token_logprobs[
+                    torch.arange(len(next_token_ids), device=next_token_ids.device),
                     next_token_ids,
                 ].tolist()
+                output.prefill_token_logprobs = output.prefill_token_logprobs.tolist()
+                output.normalized_prompt_logprobs = (
+                    output.normalized_prompt_logprobs.tolist()
+                )
 
             next_token_ids = next_token_ids.tolist()
         else:
             next_token_ids = [self.tokenizer.eos_token_id] * len(batch.reqs)
 
-        # Check finish condition
+        # Check finish conditions
         pt = 0
         for i, req in enumerate(batch.reqs):
             req.completion_tokens_wo_jump_forward += 1
-            req.output_ids = [next_token_ids[i]]
+            req.output_ids.append(next_token_ids[i])
             req.check_finished()
 
             if req.return_logprob:
-                if req.normalized_prompt_logprob is None:
-                    req.normalized_prompt_logprob = normalized_prompt_logprobs[i]
-
-                if req.prefill_token_logprobs is None:
-                    # If logprob_start_len > 0, then first logprob_start_len prompt tokens will be ignored.
-                    req.prefill_token_logprobs = list(
-                        zip(
-                            prefill_token_logprobs[pt : pt + req.extend_input_len - 1],
-                            req.input_ids[-req.extend_input_len + 1 :],
-                        )
-                    )
-                    if req.logprob_start_len == 0:
-                        req.prefill_token_logprobs = [
-                            (None, req.input_ids[0])
-                        ] + req.prefill_token_logprobs
-
-                if req.last_update_decode_tokens != 0:
-                    req.decode_token_logprobs.extend(
-                        list(
-                            zip(
-                                prefill_token_logprobs[
-                                    pt
-                                    + req.extend_input_len
-                                    - req.last_update_decode_tokens : pt
-                                    + req.extend_input_len
-                                    - 1
-                                ],
-                                req.input_ids[-req.last_update_decode_tokens + 1 :],
-                            )
-                        )
-                    )
+                self.add_logprob_return_values(i, req, pt, next_token_ids, output)
+                pt += req.extend_input_len
 
-                req.decode_token_logprobs.append(
-                    (last_token_logprobs[i], next_token_ids[i])
-                )
+        self.handle_finished_requests(batch)
 
-            if req.top_logprobs_num > 0:
-                if req.prefill_top_logprobs is None:
-                    req.prefill_top_logprobs = prefill_top_logprobs[i]
-                    if req.logprob_start_len == 0:
-                        req.prefill_top_logprobs = [None] + req.prefill_top_logprobs
+    def add_logprob_return_values(self, i, req, pt, next_token_ids, output):
+        if req.normalized_prompt_logprob is None:
+            req.normalized_prompt_logprob = output.normalized_prompt_logprobs[i]
 
-                if req.last_update_decode_tokens != 0:
-                    req.decode_top_logprobs.extend(
-                        prefill_top_logprobs[i][-req.last_update_decode_tokens + 1 :]
+        if req.prefill_token_logprobs is None:
+            # If logprob_start_len > 0, then first logprob_start_len prompt tokens will be ignored.
+            req.prefill_token_logprobs = list(
+                zip(
+                    output.prefill_token_logprobs[pt : pt + req.extend_input_len - 1],
+                    req.input_ids[-req.extend_input_len + 1 :],
+                )
+            )
+            if req.logprob_start_len == 0:
+                req.prefill_token_logprobs = [
+                    (None, req.input_ids[0])
+                ] + req.prefill_token_logprobs
+
+        if req.last_update_decode_tokens != 0:
+            req.decode_token_logprobs.extend(
+                list(
+                    zip(
+                        output.prefill_token_logprobs[
+                            pt
+                            + req.extend_input_len
+                            - req.last_update_decode_tokens : pt
+                            + req.extend_input_len
+                            - 1
+                        ],
+                        req.input_ids[-req.last_update_decode_tokens + 1 :],
                     )
-                req.decode_top_logprobs.append(decode_top_logprobs[i])
+                )
+            )
 
-            pt += req.extend_input_len
+        req.decode_token_logprobs.append(
+            (output.next_token_logprobs[i], next_token_ids[i])
+        )
 
-        self.handle_finished_requests(batch)
+        if req.top_logprobs_num > 0:
+            if req.prefill_top_logprobs is None:
+                req.prefill_top_logprobs = output.prefill_top_logprobs[i]
+                if req.logprob_start_len == 0:
+                    req.prefill_top_logprobs = [None] + req.prefill_top_logprobs
+
+            if req.last_update_decode_tokens != 0:
+                req.decode_top_logprobs.extend(
+                    output.prefill_top_logprobs[i][-req.last_update_decode_tokens + 1 :]
+                )
+            req.decode_top_logprobs.append(output.decode_top_logprobs[i])
 
     def cache_filled_batch(self, batch: Batch):
         req_pool_indices_cpu = batch.req_pool_indices.cpu().numpy()
         for i, req in enumerate(batch.reqs):
             new_prefix_indices, new_last_node = self.tree_cache.cache_req(
-                token_ids=tuple(req.input_ids + req.output_ids)[:-1],
+                token_ids=tuple(req.origin_input_ids + req.output_ids)[:-1],
                 last_uncached_pos=len(req.prefix_indices),
                 req_pool_idx=req_pool_indices_cpu[i],
                 del_in_memory_pool=False,
@@ -531,7 +537,7 @@ class ModelTpServer:
             req.prefix_indices, req.last_node = new_prefix_indices, new_last_node
 
     def forward_decode_batch(self, batch: Batch):
-        # check if decode out of memory
+        # Check if decode out of memory
         if not batch.check_decode_mem():
             old_ratio = self.new_token_ratio
             self.new_token_ratio = min(old_ratio + self.new_token_ratio_recovery, 1.0)
@@ -550,9 +556,8 @@ class ModelTpServer:
             )
 
         if not self.disable_regex_jump_forward:
-            # check for jump-forward
+            # Check for jump-forward
             jump_forward_reqs = batch.check_for_jump_forward(self.model_runner)
-
             self.forward_queue.extend(jump_forward_reqs)
             if batch.is_empty():
                 return
@@ -561,23 +566,19 @@ class ModelTpServer:
         self.decode_forward_ct = (self.decode_forward_ct + 1) % (1 << 30)
         batch.prepare_for_decode()
 
-        # Forward
-        logits, (
-            _,
-            _,
-            _,
-            decode_top_logprobs,
-            last_logprobs,
-        ) = self.model_runner.forward(batch, ForwardMode.DECODE)
-        next_token_ids, _ = batch.sample(logits)
-        next_token_ids = next_token_ids.tolist()
+        # Forward and sample the next tokens
+        output = self.model_runner.forward(batch, ForwardMode.DECODE)
+        next_token_ids, _ = batch.sample(output.next_token_logits)
 
-        # Only batch transfer the selected logprobs of the next token to CPU to reduce overhead.
-        if last_logprobs is not None:
-            new_token_logprobs = last_logprobs[
-                torch.arange(len(batch.reqs)), next_token_ids
+        # Move logprobs to cpu
+        if output.next_token_logprobs is not None:
+            next_token_logprobs = output.next_token_logprobs[
+                torch.arange(len(next_token_ids), device=next_token_ids.device),
+                next_token_ids,
             ].tolist()
 
+        next_token_ids = next_token_ids.tolist()
+
         # Check finish condition
         for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
             req.completion_tokens_wo_jump_forward += 1
@@ -585,17 +586,19 @@ class ModelTpServer:
             req.check_finished()
 
             if req.return_logprob:
-                req.decode_token_logprobs.append((new_token_logprobs[i], next_token_id))
-
-            if req.top_logprobs_num > 0:
-                req.decode_top_logprobs.append(decode_top_logprobs[i])
+                req.decode_token_logprobs.append(
+                    (next_token_logprobs[i], next_token_id)
+                )
+                if req.top_logprobs_num > 0:
+                    req.decode_top_logprobs.append(output.decode_top_logprobs[i])
 
         self.handle_finished_requests(batch)
 
     def handle_finished_requests(self, batch: Batch):
         output_rids = []
-        prev_output_strs = []
-        output_tokens = []
+        decoded_texts = []
+        surr_output_ids = []
+        read_output_ids = []
         output_skip_special_tokens = []
         output_spaces_between_special_tokens = []
         output_meta_info = []
@@ -618,8 +621,10 @@ class ModelTpServer:
                 )
             ):
                 output_rids.append(req.rid)
-                prev_output_strs.append(req.prev_output_str)
-                output_tokens.append(req.output_ids)
+                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)
                 output_skip_special_tokens.append(
                     req.sampling_params.skip_special_tokens
                 )
@@ -629,7 +634,7 @@ class ModelTpServer:
 
                 meta_info = {
                     "prompt_tokens": len(req.origin_input_ids),
-                    "completion_tokens": len(req.prev_output_ids) + len(req.output_ids),
+                    "completion_tokens": len(req.output_ids),
                     "completion_tokens_wo_jump_forward": req.completion_tokens_wo_jump_forward,
                     "finish_reason": str(req.finished_reason),
                 }
@@ -655,8 +660,9 @@ class ModelTpServer:
             self.out_pyobjs.append(
                 BatchTokenIDOut(
                     output_rids,
-                    prev_output_strs,
-                    output_tokens,
+                    decoded_texts,
+                    surr_output_ids,
+                    read_output_ids,
                     output_skip_special_tokens,
                     output_spaces_between_special_tokens,
                     output_meta_info,
@@ -671,7 +677,7 @@ class ModelTpServer:
             for i in finished_indices:
                 req = batch.reqs[i]
                 self.tree_cache.cache_req(
-                    token_ids=tuple(req.input_ids + req.output_ids)[:-1],
+                    token_ids=tuple(req.origin_input_ids + req.output_ids)[:-1],
                     last_uncached_pos=len(req.prefix_indices),
                     req_pool_idx=req_pool_indices_cpu[i],
                 )
@@ -758,12 +764,28 @@ class ModelTpClient:
         else:
             with ThreadPoolExecutor(self.tp_size) as executor:
                 # Launch model processes
-                rets = executor.map(
-                    lambda args: start_rpyc_process(*args),
-                    [(ModelTpService, p) for p in model_port_args.model_tp_ports],
+                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)
                 )
-                self.model_services = [x[0] for x in rets]
-                self.procs = [x[1] for x in rets]
 
                 # Init model
                 def init_model(i):
@@ -775,7 +797,7 @@ class ModelTpClient:
                         model_overide_args,
                     )
 
-                self.model_servers = executor.map(init_model, range(self.tp_size))
+                self.model_servers = list(executor.map(init_model, range(self.tp_size)))
 
             # Wrap functions
             def async_wrap(func_name):
@@ -788,4 +810,4 @@ class ModelTpClient:
 
                 return _func
 
-            self.step = async_wrap("step")
+            self.step = async_wrap("step")

sglang/srt/managers/detokenizer_manager.py

@@ -1,3 +1,5 @@
+"""DetokenizerManager is a process that detokenizes the token ids."""
+
 import asyncio
 import inspect
 
@@ -6,10 +8,10 @@ import zmq
 import zmq.asyncio
 
 from sglang.srt.hf_transformers_utils import get_tokenizer
+from sglang.srt.managers.controller.infer_batch import FINISH_MATCHED_STR
 from sglang.srt.managers.io_struct import BatchStrOut, BatchTokenIDOut
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.utils import get_exception_traceback, graceful_registry
-from sglang.srt.managers.controller.infer_batch import FINISH_MATCHED_STR
+from sglang.utils import find_printable_text, get_exception_traceback, graceful_registry
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
@@ -38,30 +40,26 @@ class DetokenizerManager:
             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,
+            surr_texts = self.tokenizer.batch_decode(
+                recv_obj.surr_output_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,
                 skip_special_tokens=recv_obj.skip_special_tokens[0],
-                spaces_between_special_tokens=recv_obj.spaces_between_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]
+            output_strs = []
+            for i in range(len(recv_obj.rids)):
+                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)
 
                 if isinstance(recv_obj.finished_reason[i], FINISH_MATCHED_STR):
                     pos = output_strs[i].find(recv_obj.finished_reason[i].matched)
@@ -71,7 +69,7 @@ class DetokenizerManager:
             self.send_to_tokenizer.send_pyobj(
                 BatchStrOut(
                     rids=recv_obj.rids,
-                    output_str=output_strs,
+                    output_strs=output_strs,
                     meta_info=recv_obj.meta_info,
                     finished_reason=recv_obj.finished_reason,
                 )

sglang/srt/managers/io_struct.py

@@ -1,9 +1,14 @@
+"""
+The definition of objects transfered between different
+processes (TokenizerManager, DetokenizerManager, Controller).
+"""
+
 import uuid
 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
+from sglang.srt.sampling_params import SamplingParams
 
 
 @dataclass
@@ -30,7 +35,6 @@ class GenerateReqInput:
     stream: bool = False
 
     def post_init(self):
-
         if (self.text is None and self.input_ids is None) or (
             self.text is not None and self.input_ids is not None
         ):
@@ -106,17 +110,19 @@ class TokenizedGenerateReqInput:
 @dataclass
 class BatchTokenIDOut:
     rids: List[str]
-    prev_output_strs: List[str]
-    output_tokens: List[List[int]]
+    decoded_texts: List[str]
+    surr_output_ids: List[List[int]]
+    read_output_ids: List[List[int]]
     skip_special_tokens: List[bool]
     spaces_between_special_tokens: List[bool]
     meta_info: List[Dict]
     finished_reason: List[BaseFinishReason]
 
+
 @dataclass
 class BatchStrOut:
     rids: List[str]
-    output_str: List[str]
+    output_strs: List[str]
     meta_info: List[Dict]
     finished_reason: List[BaseFinishReason]