sglang 0.1.21__py3-none-any.whl → 0.1.24__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,155 +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)
+
+    def start_dp_worker(self, dp_worker_id: int):
+        tp_size = self.server_args.tp_size
 
-        # 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))
+        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.new_event_loop()
-    loop.set_default_executor(ThreadPoolExecutor(max_workers=256))
+    pipe_writer.send("init ok")
 
-    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,138 +1,88 @@
 """A controller that manages a group of tensor parallel workers."""
 
-import multiprocessing
 import logging
+import multiprocessing
 import os
-import pickle
+from typing import List
 
-import torch
-import torch.distributed as dist
 import zmq
-import zmq.asyncio
 
-from sglang.srt.managers.controller.tp_worker import ModelTpServer
-from sglang.srt.server_args import PortArgs, ServerArgs, ModelPortArgs
+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
 
-
 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:
     """A controller that manages a group of tensor parallel workers."""
 
-    def __init__(self, server_args: ServerArgs, port_args: PortArgs, model_overide_args: dict):
+    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.server_args = server_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.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}")
+
+        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 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
+        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.model_port_args[0], model_overide_args)
+                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.model_port_args[0],
+            port_args.nccl_ports[dp_worker_id],
             model_overide_args,
         )
         self.tp_cpu_group = self.tp_server.model_runner.tp_group.cpu_group
 
     def loop_for_forward(self):
         while True:
-            recv_reqs = self.recv_requests()
+            if not self.is_dp_worker:
+                recv_reqs = self.recv_requests_from_zmq()
+            else:
+                recv_reqs = self.recv_requests_from_mp_queue()
 
-            if self.server_args.tp_size > 1:
+            if self.tp_size > 1:
                 broadcast_recv_input(recv_reqs, 0, self.tp_cpu_group)
 
             out_pyobjs = self.tp_server.exposed_step(recv_reqs)
@@ -140,27 +90,57 @@ class ControllerSingle:
             for obj in out_pyobjs:
                 self.send_to_detokenizer.send_pyobj(obj)
 
-    def recv_requests(self):
+    def recv_requests_from_zmq(self):
         recv_reqs = []
         while True:
             try:
                 recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
-                recv_reqs.append(recv_req)
             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: dict
+    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",
     )
 
+    if not is_data_parallel_worker:
+        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)]
+        dp_worker_id = 0
+        queue = None
+
     try:
-        controller = ControllerSingle(server_args, port_args, model_overide_args)
+        controller = ControllerSingle(
+            server_args,
+            port_args,
+            model_overide_args,
+            gpu_ids,
+            is_data_parallel_worker,
+            dp_worker_id,
+            queue,
+        )
     except Exception:
         pipe_writer.send(get_exception_traceback())
         raise

sglang/srt/managers/controller/model_runner.py

@@ -9,19 +9,23 @@ 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, get_tp_group
-from vllm.model_executor.model_loader import get_model
+from vllm.distributed import (
+    get_tp_group,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
 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 (
@@ -87,12 +91,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)
@@ -124,6 +122,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,
@@ -169,7 +176,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(
@@ -200,13 +210,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),
@@ -237,12 +240,24 @@ class ModelRunner:
             self.cuda_graph_runner = None
             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)]
+        logger.info(
+            f"[gpu_id={self.gpu_id}] Capture cuda graph begin. This can take up to several minutes."
+        )
+        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)
+            self,
+            max_batch_size_to_capture=max(batch_size_list),
+            use_torch_compile=self.server_args.enable_torch_compile,
         )
-        self.cuda_graph_runner.capture(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/schedule_heuristic.py

@@ -14,7 +14,7 @@ class ScheduleHeuristic:
         tree_cache,
     ):
         if tree_cache.disable and schedule_heuristic == "lpm":
-            # LMP is not meaningless when tree cache is disabled.
+            # LMP is meaningless when the tree cache is disabled.
             schedule_heuristic = "fcfs"
 
         self.schedule_heuristic = schedule_heuristic
@@ -28,11 +28,16 @@ class ScheduleHeuristic:
             # longest prefix match
             forward_queue.sort(key=lambda x: -len(x.prefix_indices))
             return forward_queue
+        elif self.schedule_heuristic == "fcfs":
+            # first come first serve
+            return forward_queue
+        elif self.schedule_heuristic == "lof":
+            # longest output first
+            forward_queue.sort(key=lambda x: -x.sampling_params.max_new_tokens)
+            return forward_queue
         elif self.schedule_heuristic == "random":
             random.shuffle(forward_queue)
             return forward_queue
-        elif self.schedule_heuristic == "fcfs":
-            return forward_queue
         elif self.schedule_heuristic == "dfs-weight":
             last_node_to_reqs = defaultdict(list)
             for req in forward_queue: