sglang 0.5.3rc2__py3-none-any.whl → 0.5.4.post1__py3-none-any.whl

sglang/check_env.py

@@ -47,7 +47,7 @@ PACKAGE_LIST = [
     "tiktoken",
     "anthropic",
     "litellm",
-    "decord",
+    "decord2",
 ]
 
 

sglang/compile_deep_gemm.py

@@ -19,6 +19,7 @@ import requests
 
 from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST
 from sglang.srt.entrypoints.http_server import launch_server
+from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import GenerateReqInput
 from sglang.srt.managers.tokenizer_manager import TokenizerManager
 from sglang.srt.server_args import ServerArgs
@@ -28,9 +29,9 @@ from sglang.srt.warmup import warmup
 multiprocessing.set_start_method("spawn", force=True)
 
 # Reduce warning
-os.environ["SGL_IN_DEEPGEMM_PRECOMPILE_STAGE"] = "1"
+envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.set(True)
 # Force enable deep gemm
-os.environ["SGL_ENABLE_JIT_DEEPGEMM"] = "1"
+envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(True)
 # Force enable mha chunked kv for DeepSeek V3 to avoid missing kv_b_proj DeepGEMM case
 os.environ["SGL_CHUNKED_PREFIX_CACHE_THRESHOLD"] = "0"
 
@@ -141,6 +142,9 @@ def refine_server_args(server_args: ServerArgs, compile_args: CompileArgs):
     server_args.enable_torch_compile = False
     print(f"Disable CUDA Graph and Torch Compile to save time...")
 
+    server_args.load_format = "dummy"
+    print(f"Set load format to dummy to save time...")
+
     # Set watchdog timeout to compile_args.timeout because compilation will take a long time
     server_args.watchdog_timeout = compile_args.timeout
     server_args.warmups = "compile-deep-gemm"

sglang/global_config.py

@@ -1,14 +1,11 @@
 """Global configurations"""
 
-import os
+# FIXME: deprecate this file and move all usage to sglang.srt.environ or sglang.__init__.py
 
 
 class GlobalConfig:
     """
     Store some global constants.
-
-    See also python/sglang/srt/managers/schedule_batch.py::global_server_args_dict, which stores
-    many global runtime arguments as well.
     """
 
     def __init__(self):
@@ -20,27 +17,6 @@ class GlobalConfig:
         # Default backend of the language
         self.default_backend = None
 
-        # Runtime constants: New generation token ratio estimation
-        self.default_init_new_token_ratio = float(
-            os.environ.get("SGLANG_INIT_NEW_TOKEN_RATIO", 0.7)
-        )
-        self.default_min_new_token_ratio_factor = float(
-            os.environ.get("SGLANG_MIN_NEW_TOKEN_RATIO_FACTOR", 0.14)
-        )
-        self.default_new_token_ratio_decay_steps = float(
-            os.environ.get("SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS", 600)
-        )
-        self.torch_empty_cache_interval = float(
-            os.environ.get(
-                "SGLANG_EMPTY_CACHE_INTERVAL", -1
-            )  # in seconds. Set if you observe high memory accumulation over a long serving period.
-        )
-        # Runtime constants: others
-        self.retract_decode_steps = 20
-        self.flashinfer_workspace_size = int(
-            os.environ.get("FLASHINFER_WORKSPACE_SIZE", 384 * 1024 * 1024)
-        )
-
         # Output tokenization configs
         self.skip_special_tokens_in_output = True
         self.spaces_between_special_tokens_in_out = True

sglang/lang/api.py

@@ -79,6 +79,7 @@ def gen(
     n: Optional[int] = None,
     stop: Optional[Union[str, List[str]]] = None,
     stop_token_ids: Optional[List[int]] = None,
+    stop_regex: Optional[Union[str, List[str]]] = None,
     temperature: Optional[float] = None,
     top_p: Optional[float] = None,
     top_k: Optional[int] = None,
@@ -120,6 +121,7 @@ def gen(
         n,
         stop,
         stop_token_ids,
+        stop_regex,
         temperature,
         top_p,
         top_k,
@@ -143,6 +145,7 @@ def gen_int(
     n: Optional[int] = None,
     stop: Optional[Union[str, List[str]]] = None,
     stop_token_ids: Optional[List[int]] = None,
+    stop_regex: Optional[Union[str, List[str]]] = None,
     temperature: Optional[float] = None,
     top_p: Optional[float] = None,
     top_k: Optional[int] = None,
@@ -162,6 +165,7 @@ def gen_int(
         n,
         stop,
         stop_token_ids,
+        stop_regex,
         temperature,
         top_p,
         top_k,
@@ -184,6 +188,7 @@ def gen_string(
     n: Optional[int] = None,
     stop: Optional[Union[str, List[str]]] = None,
     stop_token_ids: Optional[List[int]] = None,
+    stop_regex: Optional[Union[str, List[str]]] = None,
     temperature: Optional[float] = None,
     top_p: Optional[float] = None,
     top_k: Optional[int] = None,
@@ -203,6 +208,7 @@ def gen_string(
         n,
         stop,
         stop_token_ids,
+        stop_regex,
         temperature,
         top_p,
         top_k,

sglang/lang/interpreter.py

@@ -792,6 +792,7 @@ class StreamExecutor:
             "n",
             "stop",
             "stop_token_ids",
+            "stop_regex",
             "temperature",
             "top_p",
             "top_k",

sglang/lang/ir.py

@@ -21,6 +21,7 @@ class SglSamplingParams:
     n: int = 1
     stop: Union[str, List[str]] = ()
     stop_token_ids: Optional[List[int]] = ()
+    stop_regex: Optional[Union[str, List[str]]] = ()
     temperature: float = 1.0
     top_p: float = 1.0
     top_k: int = -1  # -1 means disable
@@ -45,6 +46,7 @@ class SglSamplingParams:
             self.n,
             self.stop,
             self.stop_token_ids,
+            self.stop_regex,
             self.temperature,
             self.top_p,
             self.top_k,
@@ -123,6 +125,7 @@ class SglSamplingParams:
             "n": self.n,
             "stop": self.stop,
             "stop_token_ids": self.stop_token_ids,
+            "stop_regex": self.stop_regex,
             "temperature": self.temperature,
             "top_p": self.top_p,
             "top_k": self.top_k,
@@ -161,6 +164,7 @@ class SglFunction:
         n: int = 1,
         stop: Optional[Union[str, List[str]]] = None,
         stop_token_ids: Optional[List[int]] = None,
+        stop_regex: Optional[Union[str, List[str]]] = None,
         temperature: float = 1.0,
         top_p: float = 1.0,
         top_k: int = -1,
@@ -184,12 +188,15 @@ class SglFunction:
             stop = []
         if stop_token_ids is None:
             stop_token_ids = []
+        if stop_regex is None:
+            stop_regex = []
 
         default_sampling_para = SglSamplingParams(
             max_new_tokens=max_new_tokens,
             n=n,
             stop=stop,
             stop_token_ids=stop_token_ids,
+            stop_regex=stop_regex,
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
@@ -221,6 +228,7 @@ class SglFunction:
         n: int = 1,
         stop: Optional[Union[str, List[str]]] = None,
         stop_token_ids: Optional[List[int]] = None,
+        stop_regex: Optional[Union[str, List[str]]] = None,
         temperature: float = 1.0,
         top_p: float = 1.0,
         top_k: int = -1,
@@ -243,6 +251,8 @@ class SglFunction:
             stop = []
         if stop_token_ids is None:
             stop_token_ids = []
+        if stop_regex is None:
+            stop_regex = []
 
         assert isinstance(batch_kwargs, (list, tuple))
         if len(batch_kwargs) == 0:
@@ -267,6 +277,7 @@ class SglFunction:
             n=n,
             stop=stop,
             stop_token_ids=stop_token_ids,
+            stop_regex=stop_regex,
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
@@ -451,6 +462,7 @@ class SglGen(SglExpr):
         n: Optional[int] = None,
         stop: Optional[Union[str, List[str]]] = None,
         stop_token_ids: Optional[List[int]] = None,
+        stop_regex: Optional[Union[str, List[str]]] = None,
         temperature: Optional[float] = None,
         top_p: Optional[float] = None,
         top_k: Optional[int] = None,
@@ -474,6 +486,7 @@ class SglGen(SglExpr):
             min_new_tokens=min_new_tokens,
             n=n,
             stop=stop,
+            stop_regex=stop_regex,
             stop_token_ids=stop_token_ids,
             temperature=temperature,
             top_p=top_p,

sglang/launch_server.py

@@ -1,30 +1,25 @@
 """Launch the inference server."""
 
+import asyncio
 import os
 import sys
 
-from sglang.srt.entrypoints.http_server import launch_server
 from sglang.srt.server_args import prepare_server_args
 from sglang.srt.utils import kill_process_tree
 
-MOVE_ENVS_WARN = """
-########################################################################
-# For contributors and developers:                                    #
-# Please move environment variable definitions to sglang.srt.environ  #
-# using the following pattern:                                        #
-#     SGLANG_XXX = EnvBool(False)                                     #
-#                                                                     #
-########################################################################
-"""
-
 if __name__ == "__main__":
     server_args = prepare_server_args(sys.argv[1:])
 
-    from sglang.srt.server_args import print_deprecated_warning
+    try:
+        if server_args.grpc_mode:
+            # Handle gRPC server
+            from sglang.srt.entrypoints.grpc_server import serve_grpc
 
-    print_deprecated_warning(MOVE_ENVS_WARN)
+            asyncio.run(serve_grpc(server_args))
+        else:
+            # Handle HTTP server
+            from sglang.srt.entrypoints.http_server import launch_server
 
-    try:
-        launch_server(server_args)
+            launch_server(server_args)
     finally:
         kill_process_tree(os.getpid(), include_parent=False)

sglang/profiler.py

@@ -25,6 +25,7 @@ def _run_profile(
     output_dir: Optional[str] = None,
     profile_name: Optional[str] = None,
     profile_by_stage: bool = False,
+    merge_profiles: bool = False,
 ) -> str:
     if output_dir is None:
         output_dir = PROFILER_DIR
@@ -60,6 +61,7 @@ def _run_profile(
         "num_steps": str(num_steps),
         "activities": activities,
         "profile_by_stage": profile_by_stage,
+        "merge_profiles": merge_profiles,
     }
 
     response = requests.post(url=url + "/start_profile", json=json_data)
@@ -76,10 +78,17 @@ def run_profile(
     output_dir: Optional[str] = None,
     profile_name: Optional[str] = None,
     profile_by_stage: bool = False,
+    merge_profiles: bool = False,
 ):
     # step based profile will self terminate on num_steps constraints
     link = _run_profile(
-        url, num_steps, activities, output_dir, profile_name, profile_by_stage
+        url,
+        num_steps,
+        activities,
+        output_dir,
+        profile_name,
+        profile_by_stage,
+        merge_profiles,
     )
     return link
 
@@ -145,6 +154,13 @@ if __name__ == "__main__":
         default=False,
         help="Whether to use rpd profiler (https://github.com/ROCm/rocmProfileData)",
     )
+    parser.add_argument(
+        "--merge-profiles",
+        action=argparse.BooleanOptionalAction,
+        type=bool,
+        default=False,
+        help="Whether to merge profiles from all ranks into a single trace file",
+    )
 
     args = parser.parse_args()
     activities = []
@@ -163,4 +179,5 @@ if __name__ == "__main__":
         args.output_dir,
         args.profile_name,
         args.profile_by_stage,
+        args.merge_profiles,
     )

sglang/srt/_custom_ops.py

@@ -15,7 +15,7 @@ if not is_hpu():
     # ROCm does not use vllm custom allreduce
     if use_vllm_custom_allreduce and not is_hip():
         try:
-            import vllm._C
+            import vllm._C  # noqa: F401
         except ImportError as e:
             logger.warning("Failed to import from vllm._C with %r", e)
     else:

sglang/srt/batch_invariant_ops/batch_invariant_ops.py

@@ -9,6 +9,22 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
+from sglang.srt.utils.common import calc_diff, get_bool_env_var
+
+if ENABLE_JIT_DEEPGEMM:
+    import deep_gemm
+
+_ENABLE_MM_DEEPGEMM = get_bool_env_var(
+    "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM", "1"
+)
+_ENABLE_MM_COMPARISON_TEST = get_bool_env_var(
+    "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_COMPARISON_TEST"
+)
+
+if not _ENABLE_MM_DEEPGEMM:
+    print("Disable DeepGEMM in batch invariant ops. Performance may be suboptimal.")
+
 __all__ = [
     "set_batch_invariant_mode",
     "is_batch_invariant_mode_enabled",
@@ -77,8 +93,6 @@ def matmul_kernel_persistent(
     k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
     num_tiles = num_pid_m * num_pid_n
 
-    tile_id_c = start_pid - NUM_SMS
-
     offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K)
     num_pid_in_group = GROUP_SIZE_M * num_pid_n
 
@@ -120,10 +134,6 @@ def matmul_kernel_persistent(
             )
             accumulator = tl.dot(a, b, accumulator)
 
-        tile_id_c += NUM_SMS
-        pid_m, pid_n = _compute_pid(
-            tile_id_c, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS
-        )
         offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
         offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
         if C_LARGE:
@@ -137,12 +147,16 @@ def matmul_kernel_persistent(
             accumulator += bias
         if c_ptr.dtype.element_ty == tl.float8e4nv:
             c = accumulator.to(tl.float8e4nv)
+        elif c_ptr.dtype.element_ty == tl.bfloat16:
+            c = accumulator.to(tl.bfloat16)
+        elif c_ptr.dtype.element_ty == tl.float32:
+            c = accumulator.to(tl.float32)
         else:
             c = accumulator.to(tl.float16)
         tl.store(c_ptrs, c, mask=c_mask)
 
 
-def matmul_persistent(
+def _matmul_persistent_triton(
     a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None
 ):
     # Check constraints.
@@ -219,6 +233,54 @@ def matmul_persistent(
     return c
 
 
+def _matmul_persistent_deepgemm(
+    a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None
+):
+    M, K = a.shape
+    K, N = b.shape
+    dtype = a.dtype
+    out = torch.empty((M, N), device=a.device, dtype=dtype)
+
+    deep_gemm.bf16_gemm_nn(a, b, out)
+
+    # TODO can this be put in DeepGEMM's `c`?
+    if bias is not None:
+        out += bias
+
+    return out
+
+
+def matmul_persistent(
+    a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None
+):
+    if (
+        _ENABLE_MM_DEEPGEMM
+        and ENABLE_JIT_DEEPGEMM
+        and (a.dtype == torch.bfloat16)
+        and (b.dtype == torch.bfloat16)
+        and a.is_contiguous()
+        and b.transpose(0, 1).is_contiguous()
+    ):
+        if _ENABLE_MM_COMPARISON_TEST:
+            out_triton = _matmul_persistent_triton(a=a, b=b, bias=bias)
+            out_deepgemm = _matmul_persistent_deepgemm(a=a, b=b, bias=bias)
+            diff = calc_diff(out_triton, out_deepgemm)
+            assert diff < 0.0001, f"{diff=} {out_triton=} {out_deepgemm=}"
+            # can be enabled for debugging
+            # print(
+            #     f"{diff=} "
+            #     f"{(out_triton - out_deepgemm).abs().mean()=} "
+            #     f"{(out_triton - out_deepgemm).abs().sum()=} "
+            #     f"{torch.sum(out_triton != out_deepgemm)=} "
+            # )
+            # print(f"{a=} {b=} {bias=} {out_triton=} {out_deepgemm=}")
+            return out_deepgemm
+
+        return _matmul_persistent_deepgemm(a=a, b=b, bias=bias)
+
+    return _matmul_persistent_triton(a=a, b=b, bias=bias)
+
+
 @triton.jit
 def _log_softmax_kernel(
     input_ptr,
@@ -497,16 +559,39 @@ def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None =
         return torch.sum(input, dim=dim, keepdim=keepdim, dtype=torch.float32) / n_elems
 
 
+def bmm_batch_invariant(a, b, *, out=None):
+    # Batched matrix multiply: (B, M, K) x (B, K, N) -> (B, M, N)
+    # Process each batch separately with our persistent kernel
+    if a.ndim == 3 and b.ndim == 3:
+        results = []
+        for i in range(a.shape[0]):
+            results.append(matmul_persistent(a[i], b[i]))
+        result = torch.stack(results, dim=0)
+
+        if out is not None:
+            out.copy_(result)
+            return out
+        return result
+    else:
+        raise ValueError(
+            f"bmm_batch_invariant expects 3D tensors, "
+            f"got shapes {a.shape} and {b.shape}"
+        )
+
+
 _batch_invariant_MODE = False
 _batch_invariant_LIB = None
+_original_torch_bmm = None
 
 
 def is_batch_invariant_mode_enabled():
     return _batch_invariant_MODE
 
 
-def enable_batch_invariant_mode():
-    global _batch_invariant_MODE, _batch_invariant_LIB
+def enable_batch_invariant_mode(
+    enable_bmm: bool = True,
+):
+    global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
     if _batch_invariant_MODE:
         return
 
@@ -519,11 +604,21 @@ def enable_batch_invariant_mode():
     )
     _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA")
 
+    if enable_bmm:
+        _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA")
+
+        # Also monkeypatch torch.bmm directly as a fallback
+        _original_torch_bmm = torch.bmm
+        torch.bmm = bmm_batch_invariant
+
 
 def disable_batch_invariant_mode():
-    global _batch_invariant_MODE, _batch_invariant_LIB
+    global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
     if _batch_invariant_LIB is not None:
         _batch_invariant_LIB._destroy()
+    if _original_torch_bmm is not None:
+        torch.bmm = _original_torch_bmm
+        _original_torch_bmm = None
     _batch_invariant_MODE = False
     _batch_invariant_LIB = None
 

sglang/srt/checkpoint_engine/checkpoint_engine_worker.py

@@ -0,0 +1,142 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""
+Checkpoint-engine integration for SGLang.
+This module provides weight update functionality via IPC for checkpoint-engine compatibility.
+"""
+import logging
+from typing import Callable, Dict, Optional
+
+import torch
+import zmq
+
+try:
+    from checkpoint_engine.worker import update_weights_from_ipc
+except ImportError:
+    raise ImportError(
+        "checkpoint-engine is not installed. "
+        "Please install it with: pip install sglang[checkpoint-engine]"
+    )
+
+logger = logging.getLogger(__name__)
+
+
+class SGLangCheckpointEngineWorkerExtension:
+    """
+    Worker extension for SGLang to support checkpoint-engine IPC weight updates.
+    This class provides the interface needed for checkpoint-engine integration.
+    """
+
+    def __init__(self):
+        self._zmq_ctx: Optional[zmq.Context] = None
+
+    def get_device_uuid(self) -> str:
+        """Get the UUID of current device."""
+        # We need to implement this to get the device UUID
+        # This will be overridden when integrated into SGLang's worker
+        raise NotImplementedError(
+            "This method should be overridden by SGLang integration"
+        )
+
+    def get_device_id(self) -> int:
+        """Get the device ID."""
+        raise NotImplementedError(
+            "This method should be overridden by SGLang integration"
+        )
+
+    def get_model_loader(self) -> Callable:
+        """Get the model weight loader function."""
+        raise NotImplementedError(
+            "This method should be overridden by SGLang integration"
+        )
+
+    def get_post_hook(self) -> Optional[Callable]:
+        """Get the post-processing hook after weight loading."""
+        return None
+
+    def update_weights_from_ipc(self, zmq_handles: Dict[str, str]):
+        """
+        Update weights from IPC communication.
+        Args:
+            zmq_handles: Dict mapping device UUID to ZMQ socket path
+        """
+        if self._zmq_ctx is None:
+            self._zmq_ctx = zmq.Context()
+        device_uuid = self.get_device_uuid()
+        device_id = self.get_device_id()
+        if device_uuid not in zmq_handles:
+            raise ValueError(
+                f"Device UUID {device_uuid} not found in zmq_handles: {list(zmq_handles.keys())}"
+            )
+        update_weights_from_ipc(
+            self._zmq_ctx,
+            zmq_handles[device_uuid],
+            device_id=device_id,
+            run=self.get_model_loader(),
+            post_hook=self.get_post_hook(),
+        )
+
+
+class SGLangCheckpointEngineWorkerExtensionImpl(SGLangCheckpointEngineWorkerExtension):
+    """
+    Implementation of SGLangCheckpointEngineWorkerExtension that integrates with SGLang's model runner.
+    This class provides the concrete implementation for checkpoint-engine IPC weight updates.
+    """
+
+    def __init__(self, model_runner):
+        super().__init__()
+        self.model_runner = model_runner
+
+    def get_device_uuid(self) -> str:
+        """Get the UUID of current device."""
+        # Get device UUID for current device
+        device_id = torch.cuda.current_device()
+        try:
+            return f"GPU-{torch.cuda.get_device_properties(device_id).uuid!s}"
+        except AssertionError as e:
+            raise ValueError(f"Failed to get GPU UUID for device {device_id}") from e
+
+    def get_device_id(self) -> int:
+        """Get the device ID."""
+        return torch.cuda.current_device()
+
+    def get_model_loader(self) -> Callable:
+        """Get the model weight loader function."""
+        return self.model_runner.model.load_weights
+
+    def get_post_hook(self) -> Optional[Callable]:
+        """Get the post-processing hook after weight loading."""
+
+        def post_hook():
+            # Perform post-processing after weight loading similar to DefaultModelLoader
+            try:
+                from sglang.srt.model_loader.loader import device_loading_context
+
+                # Process quantization methods after loading weights
+                for _, module in self.model_runner.model.named_modules():
+                    quant_method = getattr(module, "quant_method", None)
+                    if quant_method is not None:
+                        # Move parameters to device if needed for quantization processing
+                        target_device = torch.device(
+                            "cuda", torch.cuda.current_device()
+                        )
+                        with device_loading_context(module, target_device):
+                            quant_method.process_weights_after_loading(module)
+                # Call model-specific post-loading hook if available
+                if hasattr(self.model_runner.model, "post_load_weights"):
+                    self.model_runner.model.post_load_weights()
+            except Exception as e:
+                logger.warning(f"Post-hook processing failed: {e}")
+
+        return post_hook