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

sglang/srt/managers/{router/model_rpc.py → controller/tp_worker.py}

@@ -1,61 +1,68 @@
+"""A tensor parallel worker."""
+
 import asyncio
 import logging
-import multiprocessing
 import time
 import warnings
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import List, Optional
 
 import rpyc
 import torch
 from rpyc.utils.classic import obtain
-from rpyc.utils.server import ThreadedServer
-
-try:
-    from vllm.logger import _default_handler as vllm_default_logger
-except ImportError:
-    from vllm.logger import logger as vllm_default_logger
 
+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.router.infer_batch import Batch, ForwardMode, Req, FinishReason
-from sglang.srt.managers.router.model_runner import ModelRunner
-from sglang.srt.managers.router.radix_cache import RadixCache
-from sglang.srt.managers.router.scheduler import Scheduler
 from sglang.srt.model_config import ModelConfig
-from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.server_args import ModelPortArgs, ServerArgs
 from sglang.srt.utils import (
-    get_exception_traceback,
     get_int_token_logit_bias,
     is_multimodal_model,
     set_random_seed,
+    start_rpyc_service_process,
+    connect_rpyc_service,
+    suppress_other_loggers,
 )
+from sglang.utils import get_exception_traceback
 
-
-logger = logging.getLogger("model_rpc")
-vllm_default_logger.setLevel(logging.WARN)
-logging.getLogger("vllm.utils").setLevel(logging.WARN)
+logger = logging.getLogger("srt.tp_worker")
 
 
-class ModelRpcServer:
+class ModelTpServer:
     def __init__(
         self,
+        gpu_id: int,
         tp_rank: int,
         server_args: ServerArgs,
-        port_args: PortArgs,
-        model_overide_args: Optional[dict] = None,
+        model_port_args: ModelPortArgs,
+        model_overide_args,
     ):
-        server_args, port_args = [obtain(x) for x in [server_args, port_args]]
+        server_args, model_port_args = obtain(server_args), obtain(model_port_args)
+        suppress_other_loggers()
 
         # Copy arguments
+        self.gpu_id = gpu_id
         self.tp_rank = tp_rank
         self.tp_size = server_args.tp_size
+        self.dp_size = server_args.dp_size
         self.schedule_heuristic = server_args.schedule_heuristic
         self.disable_regex_jump_forward = server_args.disable_regex_jump_forward
 
@@ -66,23 +73,16 @@ class ModelRpcServer:
             context_length=server_args.context_length,
             model_overide_args=model_overide_args,
         )
-
-        # For model end global settings
-        server_args_dict = {
-            "enable_flashinfer": server_args.enable_flashinfer,
-            "attention_reduce_in_fp32": server_args.attention_reduce_in_fp32,
-        }
-
         self.model_runner = ModelRunner(
             model_config=self.model_config,
             mem_fraction_static=server_args.mem_fraction_static,
+            gpu_id=gpu_id,
             tp_rank=tp_rank,
             tp_size=server_args.tp_size,
-            nccl_port=port_args.nccl_port,
-            load_format=server_args.load_format,
-            trust_remote_code=server_args.trust_remote_code,
-            server_args_dict=server_args_dict,
+            nccl_port=model_port_args.nccl_port,
+            server_args=server_args,
         )
+
         if is_multimodal_model(server_args.model_path):
             self.processor = get_processor(
                 server_args.tokenizer_path,
@@ -96,28 +96,34 @@ class ModelRpcServer:
                 tokenizer_mode=server_args.tokenizer_mode,
                 trust_remote_code=server_args.trust_remote_code,
             )
-        self.max_total_num_token = self.model_runner.max_total_num_token
-        self.max_num_running_seq = self.max_total_num_token // 2
-        self.max_prefill_num_token = max(
-            self.model_config.context_len,
-            (
-                self.max_total_num_token // 6
-                if server_args.max_prefill_num_token is None
-                else server_args.max_prefill_num_token
-            ),
+        self.max_total_num_tokens = self.model_runner.max_total_num_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.int_token_logit_bias = torch.tensor(
             get_int_token_logit_bias(self.tokenizer, self.model_config.vocab_size)
         )
         set_random_seed(server_args.random_seed)
+
+        # Print info
         logger.info(
-            f"Rank {self.tp_rank}: "
-            f"max_total_num_token={self.max_total_num_token}, "
-            f"max_prefill_num_token={self.max_prefill_num_token}, "
-            f"context_len={self.model_config.context_len}, "
+            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}"
         )
         if self.tp_rank == 0:
-            logger.info(f"server_args: {server_args.print_mode_args()}")
+            logger.info(
+                f"[gpu_id={self.gpu_id}] "
+                f"server_args: {server_args.print_mode_args()}"
+            )
 
         # Init cache
         self.tree_cache = RadixCache(
@@ -126,11 +132,11 @@ class ModelRpcServer:
             disable=server_args.disable_radix_cache,
         )
         self.tree_cache_metrics = {"total": 0, "hit": 0}
-        self.scheduler = Scheduler(
+        self.scheduler = ScheduleHeuristic(
             self.schedule_heuristic,
-            self.max_num_running_seq,
-            self.max_prefill_num_token,
-            self.max_total_num_token,
+            self.max_running_requests,
+            self.max_prefill_tokens,
+            self.max_total_num_tokens,
             self.tree_cache,
         )
         self.req_to_token_pool = self.model_runner.req_to_token_pool
@@ -156,30 +162,23 @@ class ModelRpcServer:
         self.jump_forward_cache = JumpForwardCache()
 
         # Init new token estimation
-        self.new_token_ratio = min(0.4 * server_args.schedule_conservativeness, 1.0)
-        self.min_new_token_ratio = min(0.2 * server_args.schedule_conservativeness, 1.0)
-        self.new_token_ratio_step = (0.0001, 0.05)  # (down, up)
-
-    def flush_cache(self):
-        if len(self.forward_queue) == 0 and (
-            self.running_batch is None or len(self.running_batch.reqs) == 0
-        ):
-            self.tree_cache.reset()
-            self.tree_cache_metrics = {"total": 0, "hit": 0}
-            self.regex_fsm_cache.reset()
-            self.req_to_token_pool.clear()
-            self.token_to_kv_pool.clear()
-            torch.cuda.empty_cache()
-            logger.info("Cache flushed successfully!")
-        else:
-            warnings.warn(
-                f"Cache not flushed because there are pending requests. "
-                f"#queue-req: {len(self.forward_queue)}, "
-                f"#running-req: {0 if self.running_batch is None else len(self.running_batch.reqs)}"
-            )
+        assert (
+            server_args.schedule_conservativeness >= 0
+        ), "Invalid schedule_conservativeness"
+        self.new_token_ratio = min(
+            global_config.base_new_token_ratio * server_args.schedule_conservativeness,
+            1.0,
+        )
+        self.min_new_token_ratio = min(
+            global_config.base_min_new_token_ratio
+            * server_args.schedule_conservativeness,
+            1.0,
+        )
+        self.new_token_ratio_decay = global_config.new_token_ratio_decay
+        self.new_token_ratio_recovery = global_config.new_token_ratio_recovery
 
     def exposed_step(self, recv_reqs):
-        if self.tp_size != 1:
+        if self.tp_size * self.dp_size != 1:
             recv_reqs = obtain(recv_reqs)
 
         try:
@@ -189,13 +188,16 @@ class ModelRpcServer:
                     self.handle_generate_request(recv_req)
                 elif isinstance(recv_req, FlushCacheReq):
                     self.flush_cache()
+                elif isinstance(recv_req, AbortReq):
+                    self.abort_request(recv_req)
                 else:
                     raise ValueError(f"Invalid request: {recv_req}")
 
             # Forward
             self.forward_step()
         except Exception:
-            logger.error("Exception in ModelRpcClient:\n" + get_exception_traceback())
+            logger.error("Exception in ModelTpServer:\n" + get_exception_traceback())
+            raise
 
         # Return results
         ret = self.out_pyobjs
@@ -207,9 +209,8 @@ class ModelRpcServer:
         new_batch = self.get_new_fill_batch()
 
         if new_batch is not None:
-            # Run new fill batch
+            # Run a new fill batch
             self.forward_fill_batch(new_batch)
-
             self.cache_filled_batch(new_batch)
 
             if not new_batch.is_empty():
@@ -225,39 +226,43 @@ class ModelRpcServer:
                     self.num_generated_tokens += len(self.running_batch.reqs)
                     self.forward_decode_batch(self.running_batch)
 
-                    if self.running_batch.is_empty():
-                        self.running_batch = None
-                        break
-
-                    if self.out_pyobjs and self.running_batch.reqs[0].stream:
-                        break
-
-                    if self.running_batch is not None and self.tp_rank == 0:
+                    # Print stats
+                    if self.tp_rank == 0:
                         if self.decode_forward_ct % 40 == 0:
-                            num_used = self.max_total_num_token - (
+                            num_used = self.max_total_num_tokens - (
                                 self.token_to_kv_pool.available_size()
                                 + self.tree_cache.evictable_size()
                             )
-                            throuhgput = self.num_generated_tokens / (time.time() - self.last_stats_tic)
+                            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_token:.2f}, "
-                                f"gen throughput (token/s): {throuhgput:.2f}, "
+                                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)}"
                             )
+
+                    if self.running_batch.is_empty():
+                        self.running_batch = None
+                        break
+
+                    if self.out_pyobjs and self.running_batch.has_stream():
+                        break
             else:
-                # check the available size
+                # 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_token:
+                if available_size != self.max_total_num_tokens:
                     warnings.warn(
                         "Warning: "
-                        f"available_size={available_size}, max_total_num_token={self.max_total_num_token}\n"
+                        f"available_size={available_size}, max_total_num_tokens={self.max_total_num_tokens}\n"
                         "KV cache pool leak detected!"
                     )
 
@@ -275,8 +280,14 @@ class ModelRpcServer:
                 (recv_req.image_hash >> 64) % self.model_config.vocab_size,
             ]
             req.image_size = recv_req.image_size
-            req.input_ids, req.image_offset = self.model_runner.model.pad_input_ids(
-                req.input_ids, 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
@@ -293,23 +304,25 @@ class ModelRpcServer:
                     req.sampling_params.regex
                 )
 
-        # Truncate long prompts
-        req.input_ids = req.input_ids[: self.model_config.context_len - 1]
+        # Truncate prompts that are too long
+        req.origin_input_ids = req.origin_input_ids[: self.model_config.context_len - 1]
         req.sampling_params.max_new_tokens = min(
             req.sampling_params.max_new_tokens,
-            self.model_config.context_len - 1 - len(req.input_ids),
-            self.max_total_num_token - 128 - len(req.input_ids),
+            self.model_config.context_len - 1 - len(req.origin_input_ids),
+            self.max_total_num_tokens - 128 - len(req.origin_input_ids),
         )
         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_num_running_seq
+            and len(self.running_batch.reqs) > self.max_running_requests
         ):
             return None
 
+        # Compute matched prefix length
         for req in self.forward_queue:
+            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]
@@ -337,7 +350,7 @@ class ModelRpcServer:
             )
 
         for req in self.forward_queue:
-            if req.return_logprob:
+            if req.return_logprob and req.normalized_prompt_logprob is None:
                 # Need at least two tokens to compute normalized logprob
                 if req.extend_input_len < 2:
                     delta = 2 - req.extend_input_len
@@ -355,8 +368,9 @@ class ModelRpcServer:
             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_num_token
+                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
@@ -381,6 +395,7 @@ class ModelRpcServer:
         if len(can_run_list) == 0:
             return None
 
+        # Print stats
         if self.tp_rank == 0:
             running_req = (
                 0 if self.running_batch is None else len(self.running_batch.reqs)
@@ -394,20 +409,22 @@ class ModelRpcServer:
                 self.tree_cache_metrics["hit"] / self.tree_cache_metrics["total"]
             )
             logger.info(
-                f"new fill batch. #seq: {len(can_run_list)}. "
-                f"#cached_token: {hit_tokens}. "
-                f"#new_token: {new_batch_input_tokens}. "
-                f"#remaining_req: {len(self.forward_queue) - len(can_run_list)}. "
-                f"#running_req: {running_req}. "
-                f"tree_cache_hit_rate: {100.0 * tree_cache_hit_rate:.2f}%."
+                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}, "
+                f"cache hit rate: {100.0 * tree_cache_hit_rate:.2f}%, "
+                f"#running-req: {running_req}, "
+                f"#queue-req: {len(self.forward_queue) - len(can_run_list)}"
             )
-            #logger.debug(
+            # 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(
             can_run_list,
             self.req_to_token_pool,
@@ -423,73 +440,91 @@ class ModelRpcServer:
             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),
-                        next_token_ids].tolist()
-                )
+            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:
-                req.normalized_prompt_logprob = normalized_prompt_logprobs[i]
+                self.add_logprob_return_values(i, req, pt, next_token_ids, output)
+                pt += req.extend_input_len
+
+        self.handle_finished_requests(batch)
+
+    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 logprob_start_len > 0, then first logprob_start_len prompt tokens will be ignored.
-                req.prefill_token_logprobs = list(
+        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(
-                        prefill_token_logprobs[pt : pt + req.extend_input_len - 1],
-                        req.input_ids[-req.extend_input_len + 1 :],
+                        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 :],
                     )
                 )
-                if req.logprob_start_len == 0:
-                    req.prefill_token_logprobs = [
-                        (None, req.input_ids[0])
-                    ] + req.prefill_token_logprobs
-                req.decode_token_logprobs = [
-                    (last_token_logprobs[i], next_token_ids[i])
-                ]
+            )
+
+        req.decode_token_logprobs.append(
+            (output.next_token_logprobs[i], next_token_ids[i])
+        )
 
-            if req.top_logprobs_num > 0:
-                req.prefill_top_logprobs = prefill_top_logprobs[i]
+        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
-                req.decode_top_logprobs = [decode_top_logprobs[i]]
 
-            pt += req.extend_input_len
-
-        self.handle_finished_requests(batch)
+            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().tolist()
+        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,
@@ -498,10 +533,10 @@ class ModelRpcServer:
             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_step[1], 1.0)
+            self.new_token_ratio = min(old_ratio + self.new_token_ratio_recovery, 1.0)
 
             retracted_reqs = batch.retract_decode()
             logger.info(
@@ -512,27 +547,13 @@ class ModelRpcServer:
             self.forward_queue.extend(retracted_reqs)
         else:
             self.new_token_ratio = max(
-                self.new_token_ratio - self.new_token_ratio_step[0],
+                self.new_token_ratio - self.new_token_ratio_decay,
                 self.min_new_token_ratio,
             )
 
         if not self.disable_regex_jump_forward:
-            # check for jump-forward
-            jump_forward_reqs = batch.check_for_jump_forward()
-
-            # check for image jump-forward
-            for req in jump_forward_reqs:
-                if req.pixel_values is not None:
-                    (
-                        req.input_ids,
-                        req.image_offset,
-                    ) = self.model_runner.model.pad_input_ids(
-                        req.input_ids,
-                        req.pad_value,
-                        req.pixel_values.shape,
-                        req.image_size,
-                    )
-
+            # 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
@@ -541,23 +562,19 @@ class ModelRpcServer:
         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
@@ -565,31 +582,30 @@ class ModelRpcServer:
             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 = []
-        output_tokens = []
-        output_and_jump_forward_strs = []
-        output_hit_stop_str = []
+        decoded_texts = []
+        surr_output_ids = []
+        read_output_ids = []
         output_skip_special_tokens = []
         output_spaces_between_special_tokens = []
         output_meta_info = []
-        output_finished = []
+        output_finished_reason: List[BaseFinishReason] = []
         finished_indices = []
         unfinished_indices = []
         for i, req in enumerate(batch.reqs):
-            if req.finished:
+            if req.finished():
                 finished_indices.append(i)
             else:
                 unfinished_indices.append(i)
 
-            if req.finished or (
+            if req.finished() or (
                 (
                     req.stream
                     and (
@@ -599,9 +615,10 @@ class ModelRpcServer:
                 )
             ):
                 output_rids.append(req.rid)
-                output_tokens.append(req.output_ids)
-                output_and_jump_forward_strs.append(req.output_and_jump_forward_str)
-                output_hit_stop_str.append(req.hit_stop_str)
+                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
                 )
@@ -610,13 +627,10 @@ class ModelRpcServer:
                 )
 
                 meta_info = {
-                    "prompt_tokens": req.prompt_tokens,
-                    "completion_tokens": len(req.input_ids)
-                    + len(req.output_ids)
-                    - req.prompt_tokens,
+                    "prompt_tokens": len(req.origin_input_ids),
+                    "completion_tokens": len(req.output_ids),
                     "completion_tokens_wo_jump_forward": req.completion_tokens_wo_jump_forward,
-                    "finish_reason": FinishReason.to_str(req.finish_reason),
-                    "hit_stop_str": req.hit_stop_str,
+                    "finish_reason": str(req.finished_reason),
                 }
                 if req.return_logprob:
                     (
@@ -633,20 +647,20 @@ class ModelRpcServer:
                         req.normalized_prompt_logprob,
                     )
                 output_meta_info.append(meta_info)
-                output_finished.append(req.finished)
+                output_finished_reason.append(req.finished_reason)
 
         # Send to detokenizer
         if output_rids:
             self.out_pyobjs.append(
                 BatchTokenIDOut(
                     output_rids,
-                    output_tokens,
-                    output_and_jump_forward_strs,
-                    output_hit_stop_str,
+                    decoded_texts,
+                    surr_output_ids,
+                    read_output_ids,
                     output_skip_special_tokens,
                     output_spaces_between_special_tokens,
                     output_meta_info,
-                    output_finished,
+                    output_finished_reason,
                 )
             )
 
@@ -657,7 +671,7 @@ class ModelRpcServer:
             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],
                 )
@@ -670,21 +684,67 @@ class ModelRpcServer:
             else:
                 batch.reqs = []
 
+    def flush_cache(self):
+        if len(self.forward_queue) == 0 and (
+            self.running_batch is None or len(self.running_batch.reqs) == 0
+        ):
+            self.tree_cache.reset()
+            self.tree_cache_metrics = {"total": 0, "hit": 0}
+            self.regex_fsm_cache.reset()
+            self.req_to_token_pool.clear()
+            self.token_to_kv_pool.clear()
+            torch.cuda.empty_cache()
+            logger.info("Cache flushed successfully!")
+        else:
+            warnings.warn(
+                f"Cache not flushed because there are pending requests. "
+                f"#queue-req: {len(self.forward_queue)}, "
+                f"#running-req: {0 if self.running_batch is None else len(self.running_batch.reqs)}"
+            )
+
+    def abort_request(self, recv_req):
+        # Delete requests in the waiting queue
+        to_del = None
+        for i, req in enumerate(self.forward_queue):
+            if req.rid == recv_req.rid:
+                to_del = i
+                break
+
+        if to_del is not None:
+            del self.forward_queue[to_del]
+
+        # Delete requests in the running batch
+        if self.running_batch:
+            for req in self.running_batch.reqs:
+                if req.rid == recv_req.rid:
+                    req.finished_reason = FINISH_ABORT()
+                    break
 
-class ModelRpcService(rpyc.Service):
-    exposed_ModelRpcServer = ModelRpcServer
 
+class ModelTpService(rpyc.Service):
+    exposed_ModelTpServer = ModelTpServer
 
-class ModelRpcClient:
+
+class ModelTpClient:
     def __init__(
-        self, server_args: ServerArgs, port_args: PortArgs, model_overide_args
+        self,
+        gpu_ids: List[int],
+        server_args: ServerArgs,
+        model_port_args: ModelPortArgs,
+        model_overide_args,
     ):
-        tp_size = server_args.tp_size
+        server_args, model_port_args = obtain(server_args), obtain(model_port_args)
+        self.tp_size = server_args.tp_size
 
-        if tp_size == 1:
+        if self.tp_size * server_args.dp_size == 1:
             # Init model
-            self.model_server = ModelRpcService().exposed_ModelRpcServer(
-                0, server_args, port_args, model_overide_args
+            assert len(gpu_ids) == 1
+            self.model_server = ModelTpService().exposed_ModelTpServer(
+                0,
+                gpu_ids[0],
+                server_args,
+                model_port_args,
+                model_overide_args,
             )
 
             # Wrap functions
@@ -696,19 +756,31 @@ class ModelRpcClient:
 
             self.step = async_wrap(self.model_server.exposed_step)
         else:
-            with ThreadPoolExecutor(tp_size) as executor:
+            with ThreadPoolExecutor(self.tp_size) as executor:
                 # Launch model processes
-                rets = executor.map(start_model_process, port_args.model_rpc_ports)
-                self.remote_services = [x[0] for x in rets]
-                self.procs = [x[1] for x in rets]
+                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.remote_services[i].ModelRpcServer(
-                        i, server_args, port_args, model_overide_args
+                    return self.model_services[i].ModelTpServer(
+                        gpu_ids[i],
+                        i,
+                        server_args,
+                        model_port_args,
+                        model_overide_args,
                     )
 
-                self.model_servers = executor.map(init_model, range(tp_size))
+                self.model_servers = list(executor.map(init_model, range(self.tp_size)))
 
             # Wrap functions
             def async_wrap(func_name):
@@ -722,44 +794,3 @@ class ModelRpcClient:
                 return _func
 
             self.step = async_wrap("step")
-
-
-def _init_service(port):
-    t = ThreadedServer(
-        ModelRpcService(),
-        port=port,
-        protocol_config={
-            "allow_public_attrs": True,
-            "allow_pickle": True,
-            "sync_request_timeout": 1800,
-        },
-    )
-    t.start()
-
-
-def start_model_process(port):
-    proc = multiprocessing.Process(target=_init_service, args=(port,))
-    proc.start()
-    time.sleep(1)
-
-    repeat_count = 0
-    while repeat_count < 20:
-        try:
-            con = rpyc.connect(
-                "localhost",
-                port,
-                config={
-                    "allow_public_attrs": True,
-                    "allow_pickle": True,
-                    "sync_request_timeout": 1800,
-                },
-            )
-            break
-        except ConnectionRefusedError:
-            time.sleep(1)
-        repeat_count += 1
-    if repeat_count == 20:
-        raise RuntimeError("init rpc env error!")
-
-    assert proc.is_alive()
-    return con.root, proc