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

sglang/srt/managers/controller/manager_multi.py

@@ -3,19 +3,17 @@ A controller that manages multiple data parallel workers.
 Each data parallel worker can manage multiple tensor parallel workers.
 """
 
-import asyncio
+import dataclasses
 import logging
-from concurrent.futures import ThreadPoolExecutor
+import multiprocessing
+import os
 from enum import Enum, auto
-from typing import Dict
 
+import numpy as np
 import zmq
-import zmq.asyncio
 
-from sglang.global_config import global_config
-from sglang.srt.managers.controller.dp_worker import (
-    DataParallelWorkerThread,
-    start_data_parallel_worker,
+from sglang.srt.managers.controller.manager_single import (
+    start_controller_process as start_controller_process_single,
 )
 from sglang.srt.managers.io_struct import (
     AbortReq,
@@ -23,12 +21,15 @@ from sglang.srt.managers.io_struct import (
     TokenizedGenerateReqInput,
 )
 from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import kill_parent_process
 from sglang.utils import get_exception_traceback
 
 logger = logging.getLogger("srt.controller")
 
 
 class LoadBalanceMethod(Enum):
+    """Load balance method."""
+
     ROUND_ROBIN = auto()
     SHORTEST_QUEUE = auto()
 
@@ -41,151 +42,161 @@ class LoadBalanceMethod(Enum):
             raise ValueError(f"Invalid load balance method: {method}") from exc
 
 
-class Controller:
+@dataclasses.dataclass
+class WorkerHandle:
+    """Store the handle of a data parallel worker."""
+
+    proc: multiprocessing.Process
+    queue: multiprocessing.Queue
+
+
+class ControllerMulti:
+    """A controller that manages multiple data parallel workers."""
+
     def __init__(
         self,
-        load_balance_method: str,
         server_args: ServerArgs,
         port_args: PortArgs,
         model_overide_args,
     ):
-        self.load_balance_method = LoadBalanceMethod.from_str(load_balance_method)
+        # Parse args
         self.server_args = server_args
         self.port_args = port_args
+        self.model_overide_args = model_overide_args
+        self.load_balance_method = LoadBalanceMethod.from_str(
+            server_args.load_balance_method
+        )
 
-        if self.load_balance_method == LoadBalanceMethod.ROUND_ROBIN:
-            self.round_robin_counter = 0
+        # Init communication
+        context = zmq.Context()
+        self.recv_from_tokenizer = context.socket(zmq.PULL)
+        self.recv_from_tokenizer.bind(f"tcp://127.0.0.1:{port_args.controller_port}")
 
-        self.dispatch_lookup = {
+        # Dispatch method
+        self.round_robin_counter = 0
+        dispatch_lookup = {
             LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler,
             LoadBalanceMethod.SHORTEST_QUEUE: self.shortest_queue_scheduler,
         }
-        self.dispatching = self.dispatch_lookup[self.load_balance_method]
-
-        # Init communication
-        context = zmq.asyncio.Context()
-        self.recv_from_tokenizer = context.socket(zmq.PULL)
-        self.recv_from_tokenizer.bind(f"tcp://127.0.0.1:{port_args.router_port}")
-
-        # Init status
-        self.recv_reqs = []
+        self.dispatching = dispatch_lookup[self.load_balance_method]
 
         # Start data parallel workers
-        self.workers: Dict[int, DataParallelWorkerThread] = {}
-        tp_size = server_args.tp_size
-
-        def start_dp_worker(i):
-            try:
-                gpu_ids = list(range(i * tp_size, (i + 1) * tp_size))
-                worker_thread = start_data_parallel_worker(
-                    server_args, port_args, model_overide_args, gpu_ids, i
-                )
-                self.workers[i] = worker_thread
-            except Exception:
-                logger.error(
-                    f"Failed to start local worker {i}\n{get_exception_traceback()}"
-                )
-
+        self.workers = []
         for i in range(server_args.dp_size):
-            start_dp_worker(i)
+            self.start_dp_worker(i)
 
-        # Parallel launch is slower, probably due to the disk bandwidth limitations.
-        # with ThreadPoolExecutor(server_args.dp_size) as executor:
-        #     executor.map(start_dp_worker, range(server_args.dp_size))
+    def start_dp_worker(self, dp_worker_id: int):
+        tp_size = self.server_args.tp_size
+
+        pipe_controller_reader, pipe_controller_writer = multiprocessing.Pipe(
+            duplex=False
+        )
 
-    def have_any_live_worker(self):
-        return any(worker_thread.liveness for worker_thread in self.workers.values())
+        gpu_ids = list(range(dp_worker_id * tp_size, (dp_worker_id + 1) * tp_size))
+        queue = multiprocessing.Queue()
+        proc = multiprocessing.Process(
+            target=start_controller_process_single,
+            args=(
+                self.server_args,
+                self.port_args,
+                pipe_controller_writer,
+                self.model_overide_args,
+                True,
+                gpu_ids,
+                dp_worker_id,
+                queue,
+            ),
+        )
+        proc.start()
 
-    def put_req_to_worker(self, worker_id, req):
-        self.workers[worker_id].request_queue.put(req)
+        controller_init_state = pipe_controller_reader.recv()
+        if controller_init_state != "init ok":
+            raise RuntimeError(
+                f"Initialization failed. controller_init_state: {controller_init_state}"
+            )
+        self.workers.append(
+            WorkerHandle(
+                proc=proc,
+                queue=queue,
+            )
+        )
 
-    async def round_robin_scheduler(self, input_requests):
-        available_workers = list(self.workers.keys())
+    def round_robin_scheduler(self, input_requests):
         for r in input_requests:
-            self.put_req_to_worker(available_workers[self.round_robin_counter], r)
+            self.workers[self.round_robin_counter].queue.put(r)
             self.round_robin_counter = (self.round_robin_counter + 1) % len(
-                available_workers
+                self.workers
             )
-        return
 
-    async def shortest_queue_scheduler(self, input_requests):
+    def shortest_queue_scheduler(self, input_requests):
         for r in input_requests:
-            worker = min(
-                self.workers, key=lambda w: self.workers[w].request_queue.qsize()
-            )
-            self.put_req_to_worker(worker, r)
-        return
-
-    async def remove_dead_workers(self):
-        for i in list(self.workers.keys()):
-            worker_thread = self.workers[i]
-            if not worker_thread.liveness:
-                worker_thread.join()
-                # move unsuccessful requests back to the queue
-                while not worker_thread.request_queue.empty():
-                    self.recv_reqs.append(worker_thread.request_queue.get())
-                del self.workers[i]
-                logger.info(f"Stale worker {i} removed")
-
-    async def loop_for_forward(self):
-        while True:
-            await self.remove_dead_workers()
+            queue_sizes = [worker.queue.qsize() for worker in self.workers]
+            wid = np.argmin(queue_sizes)
+            self.workers[wid].queue.put(r)
 
-            if self.have_any_live_worker():
-                next_step_input = list(self.recv_reqs)
-                self.recv_reqs = []
-                if next_step_input:
-                    await self.dispatching(next_step_input)
-            # else:
-            #    logger.error("There is no live worker.")
+    def loop_for_forward(self):
+        while True:
+            recv_reqs = self.recv_requests()
+            self.dispatching(recv_reqs)
 
-            await asyncio.sleep(global_config.wait_for_new_request_delay)
+    def recv_requests(self):
+        recv_reqs = []
 
-    async def loop_for_recv_requests(self):
         while True:
-            recv_req = await self.recv_from_tokenizer.recv_pyobj()
+            try:
+                recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
+            except zmq.ZMQError:
+                break
+
             if isinstance(recv_req, FlushCacheReq):
                 # TODO(lsyin): apply more specific flushCacheReq
-                for worker_thread in self.workers.values():
-                    worker_thread.request_queue.put(recv_req)
-            elif isinstance(recv_req, TokenizedGenerateReqInput):
-                self.recv_reqs.append(recv_req)
+                for worker in self.workers:
+                    worker.queue.put(recv_req)
             elif isinstance(recv_req, AbortReq):
                 in_queue = False
-                for i, req in enumerate(self.recv_reqs):
+                for i, req in enumerate(recv_reqs):
                     if req.rid == recv_req.rid:
-                        self.recv_reqs[i] = recv_req
+                        recv_reqs[i] = recv_req
                         in_queue = True
                         break
                 if not in_queue:
                     # Send abort req to all TP groups
-                    for worker in list(self.workers.keys()):
-                        self.put_req_to_worker(worker, recv_req)
+                    for worker in self.workers:
+                        worker.queue.put(recv_req)
+            elif isinstance(recv_req, TokenizedGenerateReqInput):
+                recv_reqs.append(recv_req)
             else:
                 logger.error(f"Invalid object: {recv_req}")
 
+        return recv_reqs
+
 
 def start_controller_process(
     server_args: ServerArgs,
     port_args: PortArgs,
     pipe_writer,
-    model_overide_args=None,
+    model_overide_args: dict,
 ):
+    """Start a controller process."""
+
     logging.basicConfig(
         level=getattr(logging, server_args.log_level.upper()),
         format="%(message)s",
     )
 
     try:
-        controller = Controller(
-            server_args.load_balance_method, server_args, port_args, model_overide_args
-        )
+        controller = ControllerMulti(server_args, port_args, model_overide_args)
     except Exception:
         pipe_writer.send(get_exception_traceback())
         raise
 
     pipe_writer.send("init ok")
-    loop = asyncio.get_event_loop()
-    asyncio.set_event_loop(loop)
-    loop.create_task(controller.loop_for_recv_requests())
-    loop.run_until_complete(controller.loop_for_forward())
+
+    try:
+        controller.loop_for_forward()
+    except Exception:
+        logger.error("Exception in ControllerMulti:\n" + get_exception_traceback())
+    finally:
+        for w in controller.workers:
+            os.kill(w.proc.pid, 9)
+        kill_parent_process()

sglang/srt/managers/controller/manager_single.py

@@ -1,102 +1,157 @@
 """A controller that manages a group of tensor parallel workers."""
 
-import asyncio
 import logging
-from concurrent.futures import ThreadPoolExecutor
+import multiprocessing
+import os
+from typing import List
 
-import uvloop
 import zmq
-import zmq.asyncio
 
-from sglang.global_config import global_config
-from sglang.srt.managers.controller.tp_worker import ModelTpClient
+from sglang.srt.managers.controller.tp_worker import (
+    ModelTpServer,
+    broadcast_recv_input,
+    launch_tp_servers,
+)
 from sglang.srt.server_args import PortArgs, ServerArgs
 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")
 
 
 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,
+        gpu_ids: List[int],
+        is_data_parallel_worker: bool,
+        dp_worker_id: int,
+        mp_queue: multiprocessing.Queue,
+    ):
+        # Parse args
+        self.tp_size = server_args.tp_size
+        self.is_dp_worker = is_data_parallel_worker
+        self.dp_worker_id = dp_worker_id
+        self.mp_queue = mp_queue
+
         # Init communication
-        context = zmq.asyncio.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}")
+        context = zmq.Context(2)
+
+        if not self.is_dp_worker:
+            self.recv_from_tokenizer = context.socket(zmq.PULL)
+            self.recv_from_tokenizer.bind(
+                f"tcp://127.0.0.1:{port_args.controller_port}"
+            )
 
         self.send_to_detokenizer = context.socket(zmq.PUSH)
         self.send_to_detokenizer.connect(
             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
+        # Launch other tp ranks
+        tp_size_local = server_args.tp_size // server_args.nnodes
+        self.tp_procs = []
+        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.nccl_ports[dp_worker_id],
+                model_overide_args,
+            )
+
+        # Launch tp rank 0
+        self.tp_server = ModelTpServer(
+            gpu_ids[0],
+            0,
+            server_args,
+            port_args.nccl_ports[dp_worker_id],
+            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)
+            if not self.is_dp_worker:
+                recv_reqs = self.recv_requests_from_zmq()
+            else:
+                recv_reqs = self.recv_requests_from_mp_queue()
+
+            if self.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_from_zmq(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)
+            except zmq.ZMQError:
+                break
+            recv_reqs.append(recv_req)
+
+        return recv_reqs
+
+    def recv_requests_from_mp_queue(self):
+        recv_reqs = []
+        while not self.mp_queue.empty():
+            recv_reqs.append(self.mp_queue.get())
+        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: multiprocessing.connection.Connection,
+    model_overide_args: dict,
+    is_data_parallel_worker: bool = False,
+    gpu_ids: List[int] = None,
+    dp_worker_id: int = None,
+    queue: multiprocessing.connection.Connection = None,
 ):
+    """Start a controller process."""
+
     logging.basicConfig(
         level=getattr(logging, server_args.log_level.upper()),
         format="%(message)s",
     )
 
-    try:
+    if not is_data_parallel_worker:
         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)],
+        gpu_ids = [i for _ in range(server_args.nnodes) for i in range(tp_size_local)]
+        dp_worker_id = 0
+        queue = None
+
+    try:
+        controller = ControllerSingle(
             server_args,
-            port_args.model_port_args[0],
+            port_args,
             model_overide_args,
+            gpu_ids,
+            is_data_parallel_worker,
+            dp_worker_id,
+            queue,
         )
-        controller = ControllerSingle(model_client, port_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()

sglang/srt/managers/controller/model_runner.py

@@ -9,14 +9,24 @@ from typing import Optional, Type
 
 import torch
 import torch.nn as nn
+from flashinfer import (
+    BatchDecodeWithPagedKVCacheWrapper,
+    BatchPrefillWithPagedKVCacheWrapper,
+    BatchPrefillWithRaggedKVCacheWrapper,
+)
+from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
 from vllm.config import DeviceConfig, LoadConfig
 from vllm.config import ModelConfig as VllmModelConfig
-from vllm.distributed import init_distributed_environment, initialize_model_parallel
+from vllm.distributed import (
+    get_tp_group,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
 from vllm.model_executor.model_loader import get_model
 from vllm.model_executor.models import ModelRegistry
 
 from sglang.global_config import global_config
-from sglang.srt.managers.controller.infer_batch import Batch, ForwardMode, InputMetadata, global_server_args_dict
+from sglang.srt.managers.controller.infer_batch import Batch, ForwardMode, InputMetadata
 from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
@@ -70,6 +80,7 @@ class ModelRunner:
             distributed_init_method=nccl_init_method,
         )
         initialize_model_parallel(tensor_model_parallel_size=self.tp_size)
+        self.tp_group = get_tp_group()
         total_gpu_memory = get_available_gpu_memory(
             self.gpu_id, distributed=self.tp_size > 1
         )
@@ -81,10 +92,6 @@ class ModelRunner:
                     "The memory capacity is unbalanced. Some GPUs may be occupied by other processes."
                 )
 
-        # Set some global args
-        global_server_args_dict["disable_flashinfer"] = server_args.disable_flashinfer
-        global_server_args_dict["attention_reduce_in_fp32"] = server_args.attention_reduce_in_fp32
-
         # Load the model and create memory pool
         self.load_model()
         self.init_memory_pool(total_gpu_memory)
@@ -116,6 +123,15 @@ class ModelRunner:
         if self.model_config.model_overide_args is not None:
             vllm_model_config.hf_config.update(self.model_config.model_overide_args)
 
+        if (
+            self.server_args.efficient_weight_load
+            and "llama" in self.server_args.model_path.lower()
+            and self.server_args.quantization == "fp8"
+        ):
+            from sglang.srt.model_loader.model_loader import get_model
+        else:
+            from vllm.model_executor.model_loader import get_model
+
         self.model = get_model(
             model_config=vllm_model_config,
             device_config=device_config,
@@ -161,7 +177,10 @@ class ModelRunner:
             )
 
         self.req_to_token_pool = ReqToTokenPool(
-            int(self.max_total_num_tokens / self.model_config.context_len * 256),
+            max(
+                int(self.max_total_num_tokens / self.model_config.context_len * 512),
+                2048,
+            ),
             self.model_config.context_len + 8,
         )
         self.token_to_kv_pool = TokenToKVPool(
@@ -192,13 +211,6 @@ class ModelRunner:
             self.flashinfer_decode_wrapper = None
             return
 
-        from flashinfer import (
-            BatchDecodeWithPagedKVCacheWrapper,
-            BatchPrefillWithPagedKVCacheWrapper,
-            BatchPrefillWithRaggedKVCacheWrapper,
-        )
-        from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
-
         if not _grouped_size_compiled_for_decode_kernels(
             self.model_config.num_attention_heads // self.tp_size,
             self.model_config.get_num_kv_heads(self.tp_size),
@@ -217,7 +229,9 @@ class ModelRunner:
             self.flashinfer_workspace_buffers[1], "NHD"
         )
         self.flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-            self.flashinfer_workspace_buffers[0], "NHD", use_tensor_cores=use_tensor_cores
+            self.flashinfer_workspace_buffers[0],
+            "NHD",
+            use_tensor_cores=use_tensor_cores,
         )
 
     def init_cuda_graphs(self):
@@ -228,9 +242,19 @@ class ModelRunner:
             return
 
         logger.info(f"[gpu_id={self.gpu_id}] Capture cuda graph begin.")
-        batch_size_list = [1, 2, 4] + [i * 8 for i in range(1, 16)]
-        self.cuda_graph_runner = CudaGraphRunner(self, max_batch_size_to_capture=max(batch_size_list))
-        self.cuda_graph_runner.capture(batch_size_list)
+        batch_size_list = [1, 2, 4] + [i * 8 for i in range(1, 17)]
+        self.cuda_graph_runner = CudaGraphRunner(
+            self, max_batch_size_to_capture=max(batch_size_list)
+        )
+        try:
+            self.cuda_graph_runner.capture(batch_size_list)
+        except RuntimeError as e:
+            raise Exception(
+                f"Capture cuda graph failed {e}. Possible solutions:\n"
+                f"1. disable cuda graph by --disable-cuda-graph\n"
+                f"2. set --mem-fraction-static to a smaller value\n"
+                f"Open an issue on GitHub with reproducible scripts if you need help.\n"
+            )
 
     @torch.inference_mode()
     def forward_decode(self, batch: Batch):

sglang/srt/managers/controller/radix_cache.py

@@ -82,12 +82,12 @@ class RadixCache:
 
         if self.disable:
             if del_in_memory_pool:
-                self.token_to_kv_pool.dec_refs(indices)
+                self.token_to_kv_pool.free(indices)
             else:
                 return torch.tensor([], dtype=torch.int64), self.root_node
 
         # Radix Cache takes one ref in memory pool
-        self.token_to_kv_pool.dec_refs(indices[last_uncached_pos:new_prefix_len])
+        self.token_to_kv_pool.free(indices[last_uncached_pos:new_prefix_len])
 
         if del_in_memory_pool:
             self.req_to_token_pool.free(req_pool_idx)
@@ -125,7 +125,8 @@ class RadixCache:
             if x.lock_ref > 0:
                 continue
 
-            num_evicted += evict_callback(x.value)
+            evict_callback(x.value)
+            num_evicted += len(x.value)
             self._delete_leaf(x)
 
             if len(x.parent.children) == 0:

sglang/srt/managers/controller/schedule_heuristic.py

@@ -13,6 +13,10 @@ class ScheduleHeuristic:
         max_total_num_tokens,
         tree_cache,
     ):
+        if tree_cache.disable and schedule_heuristic == "lpm":
+            # LMP is not meaningless when tree cache is disabled.
+            schedule_heuristic = "fcfs"
+
         self.schedule_heuristic = schedule_heuristic
         self.max_running_seqs = max_running_seqs
         self.max_prefill_num_tokens = max_prefill_num_tokens