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

sglang/srt/server_args.py

@@ -29,11 +29,11 @@ class ServerArgs:
     max_prefill_tokens: Optional[int] = None
     max_running_requests: Optional[int] = None
     schedule_heuristic: str = "lpm"
-    schedule_conservativeness: float = 0.8
+    schedule_conservativeness: float = 1.0
 
     # Other runtime options
     tp_size: int = 1
-    stream_interval: int = 8
+    stream_interval: int = 1
     random_seed: Optional[int] = None
 
     # Logging
@@ -55,8 +55,10 @@ class ServerArgs:
     disable_regex_jump_forward: bool = False
     disable_cuda_graph: bool = False
     disable_disk_cache: bool = False
+    enable_torch_compile: bool = False
     attention_reduce_in_fp32: bool = False
     enable_p2p_check: bool = False
+    efficient_weight_load: bool = False
 
     # Distributed args
     nccl_init_addr: Optional[str] = None
@@ -68,15 +70,15 @@ class ServerArgs:
             self.tokenizer_path = self.model_path
         if self.mem_fraction_static is None:
             if self.tp_size >= 16:
-                self.mem_fraction_static = 0.74
+                self.mem_fraction_static = 0.80
             elif self.tp_size >= 8:
-                self.mem_fraction_static = 0.78
+                self.mem_fraction_static = 0.84
             elif self.tp_size >= 4:
-                self.mem_fraction_static = 0.82
+                self.mem_fraction_static = 0.86
             elif self.tp_size >= 2:
-                self.mem_fraction_static = 0.85
-            else:
                 self.mem_fraction_static = 0.88
+            else:
+                self.mem_fraction_static = 0.89
         if isinstance(self.additional_ports, int):
             self.additional_ports = [self.additional_ports]
         elif self.additional_ports is None:
@@ -166,6 +168,15 @@ class ServerArgs:
             "--quantization",
             type=str,
             default=ServerArgs.quantization,
+            choices=[
+                "awq",
+                "fp8",
+                "gptq",
+                "marlin",
+                "gptq_marlin",
+                "squeezellm",
+                "bitsandbytes",
+            ],
             help="The quantization method.",
         )
         parser.add_argument(
@@ -243,13 +254,13 @@ class ServerArgs:
         parser.add_argument(
             "--show-time-cost",
             action="store_true",
-            help="Show time cost of custom marks",
+            help="Show time cost of custom marks.",
         )
         parser.add_argument(
             "--api-key",
             type=str,
             default=ServerArgs.api_key,
-            help="Set API key of the server",
+            help="Set API key of the server.",
         )
 
         # Data parallelism
@@ -285,17 +296,17 @@ class ServerArgs:
         parser.add_argument(
             "--disable-flashinfer",
             action="store_true",
-            help="Disable flashinfer inference kernels",
+            help="Disable flashinfer inference kernels.",
         )
         parser.add_argument(
             "--disable-radix-cache",
             action="store_true",
-            help="Disable RadixAttention",
+            help="Disable RadixAttention for prefix caching.",
         )
         parser.add_argument(
             "--disable-regex-jump-forward",
             action="store_true",
-            help="Disable regex jump-forward",
+            help="Disable regex jump-forward.",
         )
         parser.add_argument(
             "--disable-cuda-graph",
@@ -307,6 +318,11 @@ class ServerArgs:
             action="store_true",
             help="Disable disk cache to avoid possible crashes related to file system or high concurrency.",
         )
+        parser.add_argument(
+            "--enable-torch-compile",
+            action="store_true",
+            help="Optimize the model with torch.compile, experimental feature.",
+        )
         parser.add_argument(
             "--attention-reduce-in-fp32",
             action="store_true",
@@ -318,6 +334,11 @@ class ServerArgs:
             action="store_true",
             help="Enable P2P check for GPU access, otherwise the p2p access is allowed by default.",
         )
+        parser.add_argument(
+            "--efficient-weight-load",
+            action="store_true",
+            help="Turn on memory efficient weight loading with quantization (quantize per layer during loading).",
+        )
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace):
@@ -337,16 +358,9 @@ class ServerArgs:
         )
 
 
-@dataclasses.dataclass
-class ModelPortArgs:
-    nccl_port: int
-    model_tp_ips: List[str]
-    model_tp_ports: List[int]
-
-
 @dataclasses.dataclass
 class PortArgs:
     tokenizer_port: int
-    router_port: int
+    controller_port: int
     detokenizer_port: int
-    model_port_args: List[ModelPortArgs]
+    nccl_ports: List[int]

sglang/srt/utils.py

@@ -3,9 +3,9 @@
 import base64
 import fcntl
 import logging
-import multiprocessing
 import os
 import random
+import resource
 import socket
 import struct
 import time
@@ -16,12 +16,11 @@ from typing import List, Optional
 import numpy as np
 import psutil
 import requests
-import rpyc
 import torch
+import torch.distributed as dist
 import triton
 from fastapi.responses import JSONResponse
 from packaging import version as pkg_version
-from rpyc.utils.server import ThreadedServer
 from starlette.middleware.base import BaseHTTPMiddleware
 
 logger = logging.getLogger(__name__)
@@ -148,7 +147,6 @@ def is_port_available(port):
 def allocate_init_ports(
     port: Optional[int] = None,
     additional_ports: Optional[List[int]] = None,
-    tp_size: int = 1,
     dp_size: int = 1,
 ):
     """Allocate ports for all connections."""
@@ -160,8 +158,8 @@ def allocate_init_ports(
     ret_ports = list(set(x for x in ret_ports if is_port_available(x)))
     cur_port = ret_ports[-1] + 1 if len(ret_ports) > 0 else 10000
 
-    # HTTP + Tokenizer + Controller + Detokenizer + dp_size * (nccl + tp_size)
-    num_ports_needed = 4 + dp_size * (1 + tp_size)
+    # HTTP + Tokenizer + Controller + Detokenizer + dp_size * 1 (nccl)
+    num_ports_needed = 4 + dp_size
     while len(ret_ports) < num_ports_needed:
         if cur_port not in ret_ports and is_port_available(cur_port):
             ret_ports.append(cur_port)
@@ -188,71 +186,6 @@ def get_int_token_logit_bias(tokenizer, vocab_size):
     return logit_bias
 
 
-def wrap_kernel_launcher(kernel):
-    """A faster launcher for triton kernels."""
-    if int(triton.__version__.split(".")[0]) >= 3:
-        return None
-
-    gpu_id = torch.cuda.current_device()
-    kernels = kernel.cache[gpu_id].values()
-    kernel = next(iter(kernels))
-
-    # Different trition versions use different low-level names
-    if hasattr(kernel, "cu_function"):
-        kfunction = kernel.cu_function
-    else:
-        kfunction = kernel.function
-
-    if hasattr(kernel, "c_wrapper"):
-        run = kernel.c_wrapper
-    else:
-        run = kernel.run
-
-    add_cluster_dim = True
-
-    def ret_func(grid, num_warps, *args):
-        nonlocal add_cluster_dim
-
-        try:
-            if add_cluster_dim:
-                run(
-                    grid[0],
-                    grid[1],
-                    grid[2],
-                    num_warps,
-                    1,
-                    1,
-                    1,
-                    1,
-                    kernel.shared,
-                    0,
-                    kfunction,
-                    None,
-                    None,
-                    kernel,
-                    *args,
-                )
-            else:
-                run(
-                    grid[0],
-                    grid[1],
-                    grid[2],
-                    num_warps,
-                    kernel.shared,
-                    0,
-                    kfunction,
-                    None,
-                    None,
-                    kernel,
-                    *args,
-                )
-        except TypeError:
-            add_cluster_dim = not add_cluster_dim
-            ret_func(grid, num_warps, *args)
-
-    return ret_func
-
-
 def is_multimodal_model(model):
     from sglang.srt.model_config import ModelConfig
 
@@ -371,49 +304,6 @@ def load_image(image_file):
     return image, image_size
 
 
-def connect_rpyc_service(host, port):
-    repeat_count = 0
-    while repeat_count < 20:
-        try:
-            con = rpyc.connect(
-                host,
-                port,
-                config={
-                    "allow_public_attrs": True,
-                    "allow_pickle": True,
-                    "sync_request_timeout": 3600,
-                },
-            )
-            break
-        except ConnectionRefusedError as e:
-            time.sleep(1)
-        repeat_count += 1
-    if repeat_count == 20:
-        raise RuntimeError(f"Connect rpyc error: {e}")
-
-    return con.root
-
-
-def start_rpyc_service(service: rpyc.Service, port: int):
-    t = ThreadedServer(
-        service=service,
-        port=port,
-        protocol_config={
-            "allow_public_attrs": True,
-            "allow_pickle": True,
-            "sync_request_timeout": 3600,
-        },
-    )
-    t.logger.setLevel(logging.WARN)
-    t.start()
-
-
-def start_rpyc_service_process(service: rpyc.Service, port: int):
-    proc = multiprocessing.Process(target=start_rpyc_service, args=(service, port))
-    proc.start()
-    return proc
-
-
 def suppress_other_loggers():
     from vllm.logger import logger as vllm_default_logger
 
@@ -422,6 +312,9 @@ def suppress_other_loggers():
     logging.getLogger("vllm.distributed.device_communicators.pynccl").setLevel(
         logging.WARN
     )
+    logging.getLogger("vllm.distributed.device_communicators.shm_broadcast").setLevel(
+        logging.WARN
+    )
     logging.getLogger("vllm.selector").setLevel(logging.WARN)
     logging.getLogger("vllm.utils").setLevel(logging.WARN)
 
@@ -445,7 +338,7 @@ def kill_parent_process():
     """Kill the parent process and all children of the parent process."""
     current_process = psutil.Process()
     parent_process = current_process.parent()
-    children = current_process.children(recursive=True)
+    children = parent_process.children(recursive=True)
     for child in children:
         if child.pid != current_process.pid:
             os.kill(child.pid, 9)
@@ -521,6 +414,52 @@ def monkey_patch_vllm_dummy_weight_loader():
     setattr(DummyModelLoader, "load_model", load_model)
 
 
+vllm_all_gather_backup = None
+
+
+def monkey_patch_vllm_all_gather(reverse: bool = False):
+    """Monkey patch all-gather to remove in-place operations."""
+    from torch.distributed import _functional_collectives as funcol
+    from vllm.distributed.parallel_state import GroupCoordinator
+
+    global vllm_all_gather_backup
+    if vllm_all_gather_backup is None:
+        vllm_all_gather_backup = GroupCoordinator.all_gather
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        world_size = self.world_size
+        # Bypass the function if we are using only 1 GPU.
+        if world_size == 1:
+            return input_
+        assert (
+            -input_.dim() <= dim < input_.dim()
+        ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+        input_size = input_.size()
+        # Allocate output tensor.
+        output_tensor = torch.empty(
+            (world_size,) + input_size, dtype=input_.dtype, device=input_.device
+        )
+
+        output_tensor = funcol.all_gather_tensor(
+            input_, gather_dim=0, group=self.device_group
+        ).view((world_size,) + input_size)
+
+        # Reshape
+        output_tensor = output_tensor.movedim(0, dim)
+        output_tensor = output_tensor.reshape(
+            input_size[:dim] + (world_size * input_size[dim],) + input_size[dim + 1 :]
+        )
+        return output_tensor
+
+    if reverse:
+        setattr(GroupCoordinator, "all_gather", vllm_all_gather_backup)
+    else:
+        setattr(GroupCoordinator, "all_gather", all_gather)
+
+
 API_KEY_HEADER_NAME = "X-API-Key"
 
 
@@ -559,7 +498,6 @@ def get_ip_address(ifname):
 
 def send_addrs_to_rank_0(model_port_args, server_args):
     assert server_args.node_rank != 0 and server_args.dp_size == 1
-    import torch.distributed as dist
 
     ifname = os.environ.get(
         "SGLANG_SOCKET_IFNAME", os.environ.get("NCCL_SOCKET_IFNAME", "eth0")
@@ -591,7 +529,6 @@ def send_addrs_to_rank_0(model_port_args, server_args):
 
 def receive_addrs(model_port_args, server_args):
     assert server_args.node_rank == 0 and server_args.dp_size == 1
-    import torch.distributed as dist
 
     ifname = os.environ.get(
         "SGLANG_SOCKET_IFNAME", os.environ.get("NCCL_SOCKET_IFNAME", "eth0")
@@ -624,3 +561,14 @@ def receive_addrs(model_port_args, server_args):
 
     dist.barrier()
     dist.destroy_process_group()
+
+
+def set_ulimit(target_soft_limit=65535):
+    resource_type = resource.RLIMIT_NOFILE
+    current_soft, current_hard = resource.getrlimit(resource_type)
+
+    if current_soft < target_soft_limit:
+        try:
+            resource.setrlimit(resource_type, (target_soft_limit, current_hard))
+        except ValueError as e:
+            logger.warn(f"Fail to set RLIMIT_NOFILE: {e}")

sglang/test/test_conversation.py

@@ -1,5 +1,5 @@
 from sglang.srt.conversation import generate_chat_conv
-from sglang.srt.managers.openai_protocol import (
+from sglang.srt.managers.openai_api.protocol import (
     ChatCompletionMessageContentImagePart,
     ChatCompletionMessageContentImageURL,
     ChatCompletionMessageContentTextPart,

sglang/test/test_openai_protocol.py

@@ -1,4 +1,4 @@
-from sglang.srt.managers.openai_protocol import (
+from sglang.srt.managers.openai_api.protocol import (
     ChatCompletionMessageContentImagePart,
     ChatCompletionMessageContentImageURL,
     ChatCompletionMessageContentTextPart,

sglang/test/test_programs.py

@@ -306,7 +306,7 @@ def test_image_qa():
     assert (
         "taxi" in state.messages()[-1]["content"]
         or "car" in state.messages()[-1]["content"]
-    )
+    ), f"{state.messages()[-1]['content']}"
 
 
 def test_stream():

sglang/test/test_utils.py

@@ -6,9 +6,9 @@ from functools import partial
 import numpy as np
 import requests
 
-from sglang.backend.openai import OpenAI
-from sglang.backend.runtime_endpoint import RuntimeEndpoint
 from sglang.global_config import global_config
+from sglang.lang.backend.openai import OpenAI
+from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
 from sglang.utils import get_exception_traceback