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

sglang/srt/managers/controller/infer_batch.py

@@ -15,10 +15,16 @@ from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
 
 INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
 
+# Store some global server args
+global_server_args_dict = {}
+
 
 class ForwardMode(IntEnum):
+    # Prefill a new sequence. This is deprecated now. "EXTEND" covers this case.
     PREFILL = auto()
+    # Extend a sequence. The KV cache of the first part of the sequence is already computed (e.g., system prompt).
     EXTEND = auto()
+    # Decode one token.
     DECODE = auto()
 
 
@@ -66,7 +72,10 @@ class FINISH_ABORT(BaseFinishReason):
 
 
 class Req:
+    """Store all inforamtion of a request."""
+
     def __init__(self, rid, origin_input_text, origin_input_ids):
+        # Input and output info
         self.rid = rid
         self.origin_input_text = origin_input_text
         self.origin_input_ids_unpadded = origin_input_ids  # Before image padding
@@ -74,7 +83,7 @@ class Req:
         self.output_ids = []  # Each decode stage's output ids
         self.input_ids = None  # input_ids = origin_input_ids + output_ids
 
-        # For incremental decode
+        # For incremental decoding
         self.decoded_text = ""
         self.surr_offset = None  # Surrounding offset to defeat the cleanup algorithm
         self.read_offset = None
@@ -89,20 +98,19 @@ class Req:
         self.image_offset = 0
         self.pad_value = None
 
+        # Prefix info
+        self.extend_input_len = 0
+        self.prefix_indices = []
+        self.last_node = None
+
         # Sampling parameters
         self.sampling_params = None
         self.stream = False
 
-        self.tokenizer = None
-
         # Check finish
+        self.tokenizer = None
         self.finished_reason = None
 
-        # Prefix info
-        self.extend_input_len = 0
-        self.prefix_indices = []
-        self.last_node = None
-
         # Logprobs
         self.return_logprob = False
         self.logprob_start_len = 0
@@ -166,9 +174,6 @@ class Req:
 
         return False, ""
 
-    def max_new_tokens(self):
-        return self.sampling_params.max_new_tokens
-
     def check_finished(self):
         if self.finished():
             return
@@ -252,35 +257,36 @@ class Req:
 
 @dataclass
 class Batch:
+    """Store all inforamtion of a batch."""
+
+    # Request, memory pool, and cache
     reqs: List[Req]
     req_to_token_pool: ReqToTokenPool
     token_to_kv_pool: TokenToKVPool
     tree_cache: RadixCache
 
-    # batched arguments to model runner
+    # Batched arguments to model runner
     input_ids: torch.Tensor = None
     req_pool_indices: torch.Tensor = None
     seq_lens: torch.Tensor = None
     prefix_lens: torch.Tensor = None
     position_ids_offsets: torch.Tensor = None
     out_cache_loc: torch.Tensor = None
-    out_cache_cont_start: torch.Tensor = None
-    out_cache_cont_end: torch.Tensor = None
 
-    # for processing logprobs
+    # For processing logprobs
     return_logprob: bool = False
     top_logprobs_nums: List[int] = None
 
-    # for multimodal
+    # For multimodal
     pixel_values: List[torch.Tensor] = None
     image_sizes: List[List[int]] = None
     image_offsets: List[int] = None
 
-    # other arguments for control
+    # Other arguments for control
     output_ids: torch.Tensor = None
     extend_num_tokens: int = None
 
-    # batched sampling params
+    # Batched sampling params
     temperatures: torch.Tensor = None
     top_ps: torch.Tensor = None
     top_ks: torch.Tensor = None
@@ -303,8 +309,8 @@ 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 whether batch has at least 1 streaming request
         return any(r.stream for r in self.reqs)
 
     def prepare_for_extend(self, vocab_size: int, int_token_logit_bias: torch.Tensor):
@@ -338,12 +344,12 @@ class Batch:
 
         position_ids_offsets = torch.zeros((bs,), dtype=torch.int32, device=device)
 
-        # Alloc mem
+        # Allocate memory
         seq_lens, prefix_lens = np.array(seq_lens), np.array(prefix_lens)
         extend_num_tokens = seq_lens.sum() - prefix_lens.sum()
         out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
         if out_cache_loc is None:
-            self.tree_cache.evict(extend_num_tokens, self.token_to_kv_pool.dec_refs)
+            self.tree_cache.evict(extend_num_tokens, self.token_to_kv_pool.free)
             out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
 
             if out_cache_loc is None:
@@ -413,7 +419,7 @@ class Batch:
         if self.token_to_kv_pool.available_size() >= bs:
             return True
 
-        self.tree_cache.evict(bs, self.token_to_kv_pool.dec_refs)
+        self.tree_cache.evict(bs, self.token_to_kv_pool.free)
 
         if self.token_to_kv_pool.available_size() >= bs:
             return True
@@ -444,7 +450,7 @@ class Batch:
             token_indices = self.req_to_token_pool.req_to_token[
                 req_pool_indices_cpu[idx]
             ][last_uncached_pos : seq_lens_cpu[idx]]
-            self.token_to_kv_pool.dec_refs(token_indices)
+            self.token_to_kv_pool.free(token_indices)
 
             # release the last node
             self.tree_cache.dec_lock_ref(req.last_node)
@@ -555,21 +561,12 @@ class Batch:
 
         # Alloc mem
         bs = len(self.reqs)
-        alloc_res = self.token_to_kv_pool.alloc_contiguous(bs)
-        if alloc_res is None:
-            self.out_cache_loc = self.token_to_kv_pool.alloc(bs)
+        self.out_cache_loc = self.token_to_kv_pool.alloc(bs)
 
-            if self.out_cache_loc is None:
-                print("Decode out of memory. This should never happen.")
-                self.tree_cache.pretty_print()
-                exit()
-
-            self.out_cache_cont_start = None
-            self.out_cache_cont_end = None
-        else:
-            self.out_cache_loc = alloc_res[0]
-            self.out_cache_cont_start = alloc_res[1]
-            self.out_cache_cont_end = alloc_res[2]
+        if self.out_cache_loc is None:
+            print("Decode out of memory. This should never happen.")
+            self.tree_cache.pretty_print()
+            exit()
 
         self.req_to_token_pool.req_to_token[
             self.req_pool_indices, self.seq_lens - 1
@@ -583,7 +580,7 @@ class Batch:
         self.req_pool_indices = self.req_pool_indices[new_indices]
         self.prefix_lens = None
         self.position_ids_offsets = self.position_ids_offsets[new_indices]
-        self.out_cache_loc = self.out_cache_cont_start = self.out_cache_cont_end = None
+        self.out_cache_loc = None
         self.top_logprobs_nums = [self.top_logprobs_nums[i] for i in unfinished_indices]
         self.return_logprob = any(req.return_logprob for req in self.reqs)
 
@@ -596,8 +593,7 @@ class Batch:
             "logit_bias",
         ]:
             self_val = getattr(self, item, None)
-            # logit_bias can be None
-            if self_val is not None:
+            if self_val is not None:  # logit_bias can be None
                 setattr(self, item, self_val[new_indices])
 
     def merge(self, other: "Batch"):
@@ -611,7 +607,7 @@ class Batch:
         self.position_ids_offsets = torch.concat(
             [self.position_ids_offsets, other.position_ids_offsets]
         )
-        self.out_cache_loc = self.out_cache_cont_start = self.out_cache_cont_end = None
+        self.out_cache_loc = None
         self.top_logprobs_nums.extend(other.top_logprobs_nums)
         self.return_logprob = any(req.return_logprob for req in self.reqs)
 
@@ -664,7 +660,13 @@ class Batch:
         # TODO(lmzheng): apply penalty
         probs = torch.softmax(logits, dim=-1)
         probs_sort, probs_idx = _top_p_top_k(probs, self.top_ps, self.top_ks)
-        sampled_index = torch.multinomial(probs_sort, num_samples=1)
+        try:
+            sampled_index = torch.multinomial(probs_sort, num_samples=1)
+        except RuntimeError as e:
+            warnings.warn(f"Ignore errors in sampling: {e}")
+            sampled_index = torch.ones(
+                probs_sort.shape[:-1] + (1,), dtype=torch.int64, device=probs.device
+            )
         batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(
             -1
         )
@@ -692,3 +694,215 @@ def _top_p_top_k(probs: torch.Tensor, top_ps: torch.Tensor, top_ks: torch.Tensor
     ] = 0.0
     probs_sort.div_(probs_sort.max(dim=-1, keepdim=True)[0])
     return probs_sort, probs_idx
+
+
+@dataclass
+class InputMetadata:
+    """Store all inforamtion of a forward pass."""
+
+    forward_mode: ForwardMode
+    batch_size: int
+    total_num_tokens: int
+    req_pool_indices: torch.Tensor
+    seq_lens: torch.Tensor
+    positions: torch.Tensor
+    req_to_token_pool: ReqToTokenPool
+    token_to_kv_pool: TokenToKVPool
+
+    # For extend
+    extend_seq_lens: torch.Tensor
+    extend_start_loc: torch.Tensor
+    extend_no_prefix: bool
+
+    # Output location of the KV cache
+    out_cache_loc: torch.Tensor = None
+
+    # Output options
+    return_logprob: bool = False
+    top_logprobs_nums: List[int] = None
+
+    # Trition attention backend
+    triton_max_seq_len: int = 0
+    triton_max_extend_len: int = 0
+    triton_start_loc: torch.Tensor = None
+    triton_prefix_lens: torch.Tensor = None
+
+    # FlashInfer attention backend
+    flashinfer_prefill_wrapper_ragged: "BatchPrefillWithRaggedKVCacheWrapper" = None
+    flashinfer_prefill_wrapper_paged: "BatchPrefillWithPagedKVCacheWrapper" = None
+    flashinfer_decode_wrapper: "BatchDecodeWithPagedKVCacheWrapper" = None
+
+    @classmethod
+    def create(
+        cls,
+        model_runner,
+        forward_mode,
+        req_pool_indices,
+        seq_lens,
+        prefix_lens,
+        position_ids_offsets,
+        out_cache_loc,
+        top_logprobs_nums=None,
+        return_logprob=False,
+        skip_flashinfer_init=False,
+    ):
+        if not skip_flashinfer_init and not model_runner.server_args.disable_flashinfer:
+            init_flashinfer_args(
+                forward_mode,
+                model_runner,
+                req_pool_indices,
+                seq_lens,
+                prefix_lens,
+                model_runner.flashinfer_decode_wrapper,
+            )
+
+        batch_size = len(req_pool_indices)
+
+        if forward_mode == ForwardMode.DECODE:
+            positions = ((seq_lens - 1) + position_ids_offsets).to(torch.int64)
+            extend_seq_lens = extend_start_loc = extend_no_prefix = None
+            if not model_runner.server_args.disable_flashinfer:
+                # This variable is not needed in this case,
+                # we do not compute it to make it compatbile with cuda graph.
+                total_num_tokens = None
+            else:
+                total_num_tokens = int(torch.sum(seq_lens))
+        else:
+            seq_lens_cpu = seq_lens.cpu().numpy()
+            prefix_lens_cpu = prefix_lens.cpu().numpy()
+            position_ids_offsets_cpu = position_ids_offsets.cpu().numpy()
+            positions = torch.tensor(
+                np.concatenate(
+                    [
+                        np.arange(
+                            prefix_lens_cpu[i] + position_ids_offsets_cpu[i],
+                            seq_lens_cpu[i] + position_ids_offsets_cpu[i],
+                        )
+                        for i in range(batch_size)
+                    ],
+                    axis=0,
+                ),
+                device="cuda",
+            )
+            extend_seq_lens = seq_lens - prefix_lens
+            extend_start_loc = torch.zeros_like(seq_lens)
+            extend_start_loc[1:] = torch.cumsum(extend_seq_lens[:-1], dim=0)
+            extend_no_prefix = torch.all(prefix_lens == 0)
+            total_num_tokens = int(torch.sum(seq_lens))
+
+        ret = cls(
+            forward_mode=forward_mode,
+            batch_size=batch_size,
+            total_num_tokens=total_num_tokens,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            positions=positions,
+            req_to_token_pool=model_runner.req_to_token_pool,
+            token_to_kv_pool=model_runner.token_to_kv_pool,
+            out_cache_loc=out_cache_loc,
+            extend_seq_lens=extend_seq_lens,
+            extend_start_loc=extend_start_loc,
+            extend_no_prefix=extend_no_prefix,
+            return_logprob=return_logprob,
+            top_logprobs_nums=top_logprobs_nums,
+            flashinfer_prefill_wrapper_ragged=model_runner.flashinfer_prefill_wrapper_ragged,
+            flashinfer_prefill_wrapper_paged=model_runner.flashinfer_prefill_wrapper_paged,
+            flashinfer_decode_wrapper=model_runner.flashinfer_decode_wrapper,
+        )
+
+        if model_runner.server_args.disable_flashinfer:
+            (
+                ret.triton_max_seq_len,
+                ret.triton_max_extend_len,
+                ret.triton_start_loc,
+                ret.triton_prefix_lens,
+            ) = init_triton_args(forward_mode, seq_lens, prefix_lens)
+
+        return ret
+
+
+def init_flashinfer_args(
+    forward_mode,
+    model_runner,
+    req_pool_indices,
+    seq_lens,
+    prefix_lens,
+    flashinfer_decode_wrapper,
+):
+    num_qo_heads = model_runner.model_config.num_attention_heads // model_runner.tp_size
+    num_kv_heads = model_runner.model_config.get_num_kv_heads(model_runner.tp_size)
+    head_dim = model_runner.model_config.head_dim
+    batch_size = len(req_pool_indices)
+
+    if forward_mode == ForwardMode.DECODE:
+        paged_kernel_lens = seq_lens
+    else:
+        paged_kernel_lens = prefix_lens
+
+    kv_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
+    kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
+    req_pool_indices_cpu = req_pool_indices.cpu().numpy()
+    paged_kernel_lens_cpu = paged_kernel_lens.cpu().numpy()
+    kv_indices = torch.cat(
+        [
+            model_runner.req_to_token_pool.req_to_token[
+                req_pool_indices_cpu[i], : paged_kernel_lens_cpu[i]
+            ]
+            for i in range(batch_size)
+        ],
+        dim=0,
+    ).contiguous()
+    kv_last_page_len = torch.ones((batch_size,), dtype=torch.int32, device="cuda")
+
+    if forward_mode == ForwardMode.DECODE:
+        flashinfer_decode_wrapper.end_forward()
+        flashinfer_decode_wrapper.begin_forward(
+            kv_indptr,
+            kv_indices,
+            kv_last_page_len,
+            num_qo_heads,
+            num_kv_heads,
+            head_dim,
+            1,
+        )
+    else:
+        # extend part
+        qo_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
+        qo_indptr[1:] = torch.cumsum(seq_lens - prefix_lens, dim=0)
+
+        model_runner.flashinfer_prefill_wrapper_ragged.end_forward()
+        model_runner.flashinfer_prefill_wrapper_ragged.begin_forward(
+            qo_indptr,
+            qo_indptr,
+            num_qo_heads,
+            num_kv_heads,
+            head_dim,
+        )
+
+        # cached part
+        model_runner.flashinfer_prefill_wrapper_paged.end_forward()
+        model_runner.flashinfer_prefill_wrapper_paged.begin_forward(
+            qo_indptr,
+            kv_indptr,
+            kv_indices,
+            kv_last_page_len,
+            num_qo_heads,
+            num_kv_heads,
+            head_dim,
+            1,
+        )
+
+
+def init_triton_args(forward_mode, seq_lens, prefix_lens):
+    batch_size = len(seq_lens)
+    max_seq_len = int(torch.max(seq_lens))
+    start_loc = torch.zeros((batch_size,), dtype=torch.int32, device="cuda")
+    start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
+
+    if forward_mode == ForwardMode.DECODE:
+        max_extend_len = None
+    else:
+        extend_seq_lens = seq_lens - prefix_lens
+        max_extend_len = int(torch.max(extend_seq_lens))
+
+    return max_seq_len, max_extend_len, start_loc, prefix_lens

sglang/srt/managers/controller/manager_multi.py

@@ -42,6 +42,8 @@ class LoadBalanceMethod(Enum):
 
 
 class Controller:
+    """A controller that manages multiple data parallel workers."""
+
     def __init__(
         self,
         load_balance_method: str,
@@ -183,9 +185,11 @@ def start_controller_process(
     except Exception:
         pipe_writer.send(get_exception_traceback())
         raise
-
     pipe_writer.send("init ok")
-    loop = asyncio.get_event_loop()
+
+    loop = asyncio.new_event_loop()
+    loop.set_default_executor(ThreadPoolExecutor(max_workers=256))
+
     asyncio.set_event_loop(loop)
     loop.create_task(controller.loop_for_recv_requests())
     loop.run_until_complete(controller.loop_for_forward())

sglang/srt/managers/controller/manager_single.py

@@ -1,28 +1,104 @@
 """A controller that manages a group of tensor parallel workers."""
 
-import asyncio
+import multiprocessing
 import logging
-from concurrent.futures import ThreadPoolExecutor
+import os
+import pickle
 
-import uvloop
+import torch
+import torch.distributed as dist
 import zmq
 import zmq.asyncio
 
-from sglang.global_config import global_config
-from sglang.srt.managers.controller.tp_worker import ModelTpClient
-from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.managers.controller.tp_worker import ModelTpServer
+from sglang.srt.server_args import PortArgs, ServerArgs, ModelPortArgs
 from sglang.srt.utils import kill_parent_process
 from sglang.utils import get_exception_traceback
 
-asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 logger = logging.getLogger("srt.controller")
 
 
+def run_tp_server(
+    gpu_id: int,
+    tp_rank: int,
+    server_args: ServerArgs,
+    model_port_args: ModelPortArgs,
+    model_overide_args: dict,
+):
+    """Run a tp server."""
+    try:
+        model_server = ModelTpServer(
+            gpu_id,
+            tp_rank,
+            server_args,
+            model_port_args,
+            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,
+                      model_port_args, model_overide_args):
+    """Launch multiple tp servers."""
+    procs = []
+    for i in tp_rank_range:
+        proc = multiprocessing.Process(target=run_tp_server, args=(
+            gpu_ids[i], i, server_args, model_port_args, model_overide_args
+        ))
+        proc.start()
+        procs.append(proc)
+
+    return procs
+
+
+def broadcast_recv_input(data, rank, dist_group):
+    """Broadcast inputs from rank=0 to all other ranks with torch.dist backend."""
+
+    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:
+            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
+
+
 class ControllerSingle:
-    def __init__(self, model_client: ModelTpClient, port_args: PortArgs):
+    """A controller that manages a group of tensor parallel workers."""
+
+    def __init__(self, server_args: ServerArgs, port_args: PortArgs, model_overide_args: dict):
+        # Parse args
+        self.server_args = server_args
+
         # Init communication
-        context = zmq.asyncio.Context(2)
+        context = zmq.Context(2)
         self.recv_from_tokenizer = context.socket(zmq.PULL)
         self.recv_from_tokenizer.bind(f"tcp://127.0.0.1:{port_args.router_port}")
 
@@ -31,44 +107,52 @@ class ControllerSingle:
             f"tcp://127.0.0.1:{port_args.detokenizer_port}"
         )
 
-        # Init status
-        self.model_client = model_client
-        self.recv_reqs = []
-
-        # Init some configs
-        self.request_dependency_delay = global_config.request_dependency_delay
+        # Init model server
+        tp_size_local = server_args.tp_size // server_args.nnodes
+        gpu_ids = [i for _ in range(server_args.nnodes) for i in range(tp_size_local)]
+
+        # Launch other tp ranks
+        if tp_size_local > 1:
+            tp_rank_range = range(1, tp_size_local)
+            self.tp_procs = launch_tp_servers(
+                gpu_ids, tp_rank_range, server_args,
+                port_args.model_port_args[0], model_overide_args)
+
+        # Launch tp rank 0
+        self.tp_server = ModelTpServer(
+            gpu_ids[0],
+            0,
+            server_args,
+            port_args.model_port_args[0],
+            model_overide_args,
+        )
+        self.tp_cpu_group = self.tp_server.model_runner.tp_group.cpu_group
 
-    async def loop_for_forward(self):
+    def loop_for_forward(self):
         while True:
-            next_step_input = list(self.recv_reqs)
-            self.recv_reqs = []
-            out_pyobjs = await self.model_client.step(next_step_input)
+            recv_reqs = self.recv_requests()
+
+            if self.server_args.tp_size > 1:
+                broadcast_recv_input(recv_reqs, 0, self.tp_cpu_group)
+
+            out_pyobjs = self.tp_server.exposed_step(recv_reqs)
 
             for obj in out_pyobjs:
                 self.send_to_detokenizer.send_pyobj(obj)
 
-            # async sleep for receiving the subsequent request and avoiding cache miss
-            slept = False
-            if len(out_pyobjs) != 0:
-                has_finished = any(
-                    [obj.finished_reason is not None for obj in out_pyobjs]
-                )
-                if has_finished:
-                    if self.request_dependency_delay > 0:
-                        slept = True
-                        await asyncio.sleep(self.request_dependency_delay)
-
-            if not slept:
-                await asyncio.sleep(global_config.wait_for_new_request_delay)
-
-    async def loop_for_recv_requests(self):
+    def recv_requests(self):
+        recv_reqs = []
         while True:
-            recv_req = await self.recv_from_tokenizer.recv_pyobj()
-            self.recv_reqs.append(recv_req)
+            try:
+                recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
+                recv_reqs.append(recv_req)
+            except zmq.ZMQError:
+                break
+        return recv_reqs
 
 
 def start_controller_process(
-    server_args: ServerArgs, port_args: PortArgs, pipe_writer, model_overide_args
+    server_args: ServerArgs, port_args: PortArgs, pipe_writer, model_overide_args: dict
 ):
     logging.basicConfig(
         level=getattr(logging, server_args.log_level.upper()),
@@ -76,27 +160,18 @@ def start_controller_process(
     )
 
     try:
-        tp_size_local = server_args.tp_size // server_args.nnodes
-        model_client = ModelTpClient(
-            [i for _ in range(server_args.nnodes) for i in range(tp_size_local)],
-            server_args,
-            port_args.model_port_args[0],
-            model_overide_args,
-        )
-        controller = ControllerSingle(model_client, port_args)
+        controller = ControllerSingle(server_args, port_args, model_overide_args)
     except Exception:
         pipe_writer.send(get_exception_traceback())
         raise
 
     pipe_writer.send("init ok")
 
-    loop = asyncio.new_event_loop()
-    loop.set_default_executor(ThreadPoolExecutor(max_workers=256))
-    asyncio.set_event_loop(loop)
-    loop.create_task(controller.loop_for_recv_requests())
     try:
-        loop.run_until_complete(controller.loop_for_forward())
+        controller.loop_for_forward()
     except Exception:
         logger.error("Exception in ControllerSingle:\n" + get_exception_traceback())
     finally:
+        for t in controller.tp_procs:
+            os.kill(t.pid, 9)
         kill_parent_process()