sglang 0.4.9.post3__py3-none-any.whl → 0.4.9.post5__py3-none-any.whl

sglang/srt/utils.py

@@ -15,6 +15,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import builtins
 import ctypes
 import dataclasses
@@ -85,6 +86,8 @@ from torch.profiler import ProfilerActivity, profile, record_function
 from torch.utils._contextlib import _DecoratorContextManager
 from triton.runtime.cache import FileCacheManager
 
+from sglang.srt.metrics.func_timer import enable_func_timer
+
 logger = logging.getLogger(__name__)
 
 show_time_cost = False
@@ -744,9 +747,13 @@ def load_image(
         image = Image.open(BytesIO(image_file))
     elif image_file.startswith("http://") or image_file.startswith("https://"):
         timeout = int(os.getenv("REQUEST_TIMEOUT", "3"))
-        response = requests.get(image_file, stream=True, timeout=timeout).raw
-        image = Image.open(response)
-        response.close()
+        response = requests.get(image_file, stream=True, timeout=timeout)
+        try:
+            response.raise_for_status()
+            image = Image.open(response.raw)
+            image.load()  # Force loading to avoid issues after closing the stream
+        finally:
+            response.close()
     elif image_file.lower().endswith(("png", "jpg", "jpeg", "webp", "gif")):
         image = Image.open(image_file)
     elif image_file.startswith("data:"):
@@ -933,71 +940,6 @@ def monkey_patch_vllm_gguf_config():
     setattr(GGUFConfig, "get_quant_method", get_quant_method_with_embedding_replaced)
 
 
-def maybe_set_triton_cache_manager() -> None:
-    """Set environment variable to tell Triton to use a
-    custom cache manager"""
-    cache_manger = os.environ.get("TRITON_CACHE_MANAGER", None)
-    if cache_manger is None:
-        manager = "sglang.srt.utils:CustomCacheManager"
-        logger.debug("Setting Triton cache manager to: %s", manager)
-        os.environ["TRITON_CACHE_MANAGER"] = manager
-
-
-class CustomCacheManager(FileCacheManager):
-    # Adapted from: https://github.com/tdoublep/vllm/blob/3307522289fdfefe323b6c00d0db696651989a2f/vllm/triton_utils/custom_cache_manager.py
-    def __init__(self, key, override=False, dump=False):
-        from sglang.srt.distributed.parallel_state import get_tp_group
-
-        self.key = key
-        self.lock_path = None
-
-        try:
-            module_path = "triton.runtime.cache"
-            cache_module = importlib.import_module(module_path)
-
-            default_cache_dir = getattr(cache_module, "default_cache_dir", None)
-            default_dump_dir = getattr(cache_module, "default_dump_dir", None)
-            default_override_dir = getattr(cache_module, "default_override_dir", None)
-        except (ModuleNotFoundError, AttributeError) as e:
-            default_cache_dir = None
-            default_dump_dir = None
-            default_override_dir = None
-
-        if dump:
-            self.cache_dir = (
-                default_dump_dir()
-                if default_dump_dir is not None
-                else os.path.join(Path.home(), ".triton", "dump")
-            )
-            self.cache_dir = os.path.join(self.cache_dir, self.key)
-            self.lock_path = os.path.join(self.cache_dir, "lock")
-            os.makedirs(self.cache_dir, exist_ok=True)
-        elif override:
-            self.cache_dir = (
-                default_override_dir()
-                if default_override_dir is not None
-                else os.path.join(Path.home(), ".triton", "override")
-            )
-            self.cache_dir = os.path.join(self.cache_dir, self.key)
-        else:
-            # create cache directory if it doesn't exist
-            self.cache_dir = os.getenv("TRITON_CACHE_DIR", "").strip() or (
-                default_cache_dir()
-                if default_cache_dir is not None
-                else os.path.join(Path.home(), ".triton", "cache")
-            )
-            if self.cache_dir:
-                try:
-                    self.cache_dir = f"{self.cache_dir}_{get_tp_group().local_rank}"
-                except:
-                    self.cache_dir = f"{self.cache_dir}_{os.getpid()}"
-                self.cache_dir = os.path.join(self.cache_dir, self.key)
-                self.lock_path = os.path.join(self.cache_dir, "lock")
-                os.makedirs(self.cache_dir, exist_ok=True)
-            else:
-                raise RuntimeError("Could not create or locate cache dir")
-
-
 def set_ulimit(target_soft_limit=65535):
     # number of open files
     resource_type = resource.RLIMIT_NOFILE
@@ -2061,6 +2003,16 @@ def is_valid_ipv6_address(address: str) -> bool:
         return False
 
 
+def maybe_wrap_ipv6_address(address: str) -> str:
+    if is_valid_ipv6_address(address):
+        return f"[{address}]"
+    return address
+
+
+def format_tcp_address(ip: str, port: int) -> str:
+    return f"tcp://{maybe_wrap_ipv6_address(ip)}:{port}"
+
+
 def configure_ipv6(dist_init_addr):
     addr = dist_init_addr
     end = addr.find("]")
@@ -2100,7 +2052,7 @@ def rank0_log(msg: str):
         logger.info(msg)
 
 
-def launch_dummy_health_check_server(host, port):
+def launch_dummy_health_check_server(host, port, enable_metrics):
     import asyncio
 
     import uvicorn
@@ -2118,6 +2070,11 @@ def launch_dummy_health_check_server(host, port):
         """Check the health of the http server."""
         return Response(status_code=200)
 
+    # Add prometheus middleware
+    if enable_metrics:
+        add_prometheus_middleware(app)
+        enable_func_timer()
+
     config = uvicorn.Config(
         app,
         host=host,
@@ -2386,6 +2343,7 @@ def is_fa3_default_architecture(hf_config):
         "Gemma3ForConditionalGeneration",
         "Qwen3ForCausalLM",
         "Qwen3MoeForCausalLM",
+        "Glm4MoeForCausalLM",
     }
     return architectures[0] in default_archs
 
@@ -2906,3 +2864,89 @@ SUPPORTED_LORA_TARGET_MODULES = [
 ]
 
 LORA_TARGET_ALL_MODULES = "all"
+
+
+class ConcurrentCounter:
+    """
+    An asynchronous counter for managing concurrent tasks that need
+    coordinated increments, decrements, and waiting until the count reaches zero.
+
+    This class is useful for scenarios like tracking the number of in-flight tasks
+    and waiting for them to complete.
+    """
+
+    def __init__(self, initial: int = 0):
+        """
+        Initialize the counter with an optional initial value.
+
+        Args:
+            initial (int): The initial value of the counter. Default is 0.
+        """
+        self._count = initial
+        self._condition = asyncio.Condition()
+
+    def value(self) -> int:
+        """
+        Return the current value of the counter.
+
+        Note:
+            This method is not synchronized. It may return a stale value
+            if other coroutines are concurrently modifying the counter.
+
+        Returns:
+            int: The current counter value.
+        """
+        return self._count
+
+    def __repr__(self) -> str:
+        """Return an informative string representation of the counter."""
+        return f"<ConcurrentCounter value={self.value()}>"
+
+    async def increment(self, n: int = 1, notify_all: bool = True):
+        """
+        Atomically increment the counter by a given amount and notify all waiters.
+
+        Args:
+            n (int): The amount to increment the counter by. Default is 1.
+            notify_all (bool): Whether to notify all waiters after incrementing. Default is True.
+        """
+        async with self._condition:
+            self._count += n
+            if notify_all:
+                self._condition.notify_all()
+
+    async def decrement(self, n: int = 1, notify_all: bool = True):
+        """
+        Atomically decrement the counter by a given amount and notify all waiters.
+
+        Args:
+            n (int): The amount to decrement the counter by. Default is 1.
+            notify_all (bool): Whether to notify all waiters after decrementing. Default is True.
+        """
+        async with self._condition:
+            self._count -= n
+            if notify_all:
+                self._condition.notify_all()
+
+    async def wait_for(self, condition: Callable[[int], bool]):
+        """
+        Asynchronously wait until the counter satisfies a given condition.
+
+        This suspends the calling coroutine without blocking the thread, allowing
+        other tasks to run while waiting. When the condition is met, the coroutine resumes.
+
+        Args:
+            condition (Callable[[int], bool]): A function that takes the current counter value
+                and returns True when the condition is satisfied.
+        """
+        async with self._condition:
+            await self._condition.wait_for(lambda: condition(self._count))
+
+    async def wait_for_zero(self):
+        """
+        Asynchronously wait until the counter reaches zero.
+
+        This suspends the calling coroutine without blocking the thread, allowing
+        other tasks to run while waiting. When the counter becomes zero, the coroutine resumes.
+        """
+        self.wait_for(lambda count: count == 0)

sglang/srt/weight_sync/utils.py

@@ -0,0 +1,119 @@
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+
+from sglang.srt.entrypoints.engine import Engine
+from sglang.srt.managers.tokenizer_manager import UpdateWeightsFromTensorReqInput
+from sglang.srt.model_executor.model_runner import LocalSerializedTensor
+from sglang.srt.utils import MultiprocessingSerializer
+
+
+async def update_weights(
+    engine: Engine,
+    params_batch: list[tuple[str, torch.Tensor]],
+    device_mesh_key: str,
+    device_mesh: DeviceMesh,
+    load_format: Optional[str] = None,
+):
+    """
+    Update weights for the inference engine.
+    This function is designed to be stateless, so that the caller process could keep the stateful engine.
+    Example Use Case:
+        - Multiple Producer Process will call this function in a SPMD style
+
+    Args:
+        engine: The inference engine created by the caller process.
+        params_batch: A list of (name, tensor) tuples. We batched the tensors to avoid the overhead of cpu call.
+        device_mesh_key: The key of the device mesh. Typically "tp" or "infer_tp"
+        device_mesh: The device mesh.
+        load_format: The format of the weights.
+    """
+    infer_tp_size = device_mesh[device_mesh_key].mesh.size()[0]
+    infer_tp_rank = device_mesh[device_mesh_key].get_local_rank()
+    from sglang.srt.patch_torch import monkey_patch_torch_reductions
+
+    monkey_patch_torch_reductions()
+
+    # [
+    #   (name0, ipc_tensor0_tp0),
+    #   (name1, ipc_tensor1_tp0),
+    # ]
+    named_tensors_batch = [
+        (
+            name,
+            MultiprocessingSerializer.serialize(
+                _preprocess_tensor_for_update_weights(tensor)
+            ),
+        )
+        for name, tensor in params_batch
+    ]
+
+    if infer_tp_rank == 0:
+        gathered_serialized_batches = [None for _ in range(infer_tp_size)]
+    else:
+        gathered_serialized_batches = None
+
+    # [
+    #   [ (name0, ipc_tensor0_tp0), (name1, ipc_tensor1_tp0) ],
+    #   [ (name0, ipc_tensor0_tp1), (name1, ipc_tensor1_tp1) ],
+    # ]
+    dist.gather_object(
+        obj=named_tensors_batch,
+        object_gather_list=gathered_serialized_batches,
+        dst=device_mesh[device_mesh_key].mesh.tolist()[0],
+        group=device_mesh[device_mesh_key].get_group(),
+    )
+
+    if infer_tp_rank == 0:
+        # Use zip(*) to "transpose" the data structure.
+        # After transpose, the data structure is like:
+        # [
+        #   ( (name0, ipc_tensor0_tp0), (name0, ipc_tensor0_tp1) ),
+        #   ( (name1, ipc_tensor1_tp0), (name1, ipc_tensor1_tp1) ),
+        # ]
+        logical_tensors = zip(*gathered_serialized_batches, strict=True)
+
+        named_tensors = [
+            # [
+            #   (name0, LocalSerializedTensor(values=[ipc_tensor0_tp0, ipc_tensor0_tp1])),
+            #   (name1, LocalSerializedTensor(values=[ipc_tensor1_tp0, ipc_tensor1_tp1])),
+            # ]
+            (
+                tensor_group[0][0],
+                LocalSerializedTensor(
+                    values=[rank_part[1] for rank_part in tensor_group]
+                ),
+            )
+            for tensor_group in logical_tensors
+        ]
+
+        update_weights_request = UpdateWeightsFromTensorReqInput(
+            serialized_named_tensors=[
+                MultiprocessingSerializer.serialize(named_tensors)
+                for _ in range(infer_tp_size)
+            ],
+            load_format=load_format,
+        )
+
+        return await engine.update_weights_from_tensor(update_weights_request)
+
+
+def _preprocess_tensor_for_update_weights(tensor: torch.Tensor):
+    """
+    Preprocess the tensor for update weights.
+    Example Use Case:
+        - FSDP: we gather tensor by calling full_tensor in _preprocess_tensor_for_update_weights
+        - Megatron: we do nothing here, assuming it is gathered when feed into this func
+
+    Args:
+        tensor: The tensor to be preprocessed.
+
+    Returns:
+        The full tensor if it is a DTensor, otherwise the original tensor.
+    """
+    if isinstance(tensor, DTensor):
+        return tensor.full_tensor()
+    return tensor

sglang/test/runners.py

@@ -491,6 +491,8 @@ class SRTRunner:
         lora_paths: List[str] = None,
         max_loras_per_batch: int = 4,
         attention_backend: Optional[str] = None,
+        prefill_attention_backend: Optional[str] = None,
+        decode_attention_backend: Optional[str] = None,
         lora_backend: str = "triton",
         disable_cuda_graph: bool = False,
         disable_radix_cache: bool = False,
@@ -540,6 +542,8 @@ class SRTRunner:
             max_loras_per_batch=max_loras_per_batch,
             lora_backend=lora_backend,
             attention_backend=attention_backend,
+            prefill_attention_backend=prefill_attention_backend,
+            decode_attention_backend=decode_attention_backend,
             disable_cuda_graph=disable_cuda_graph,
             disable_radix_cache=disable_radix_cache,
             chunked_prefill_size=chunked_prefill_size,

sglang/test/test_activation.py

@@ -3,9 +3,12 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.activation import GeluAndMul
+from sglang.srt.layers.activation import GeluAndMul, QuickGELU
+from sglang.srt.utils import is_hip
 from sglang.test.test_utils import CustomTestCase
 
+_is_hip = is_hip()
+
 
 class TestGeluAndMul(CustomTestCase):
     DTYPES = [torch.half, torch.bfloat16]
@@ -52,5 +55,51 @@ class TestGeluAndMul(CustomTestCase):
                 self._run_gelu_and_mul_test(*params)
 
 
+class TestQuickGELU(CustomTestCase):
+    DTYPES = [torch.half, torch.bfloat16]
+    NUM_TOKENS = [7, 83, 2048]  # batch = sequence length
+    DIMS = [512, 4096, 5120, 13824]  # all multiples of 16 bytes
+    SEEDS = [0]
+
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is not available")
+        torch.set_default_device("cuda")
+
+    def _run_gelu_quick_test(self, n_tok: int, dim: int, dtype: torch.dtype, seed: int):
+        torch.manual_seed(seed)
+
+        layer = QuickGELU().to(dtype=dtype)
+
+        x = torch.randn(n_tok, dim, dtype=dtype, device="cuda")
+
+        with torch.inference_mode():
+            ref = layer.forward_native(x)  # x * sigmoid(1.702 * x), fp32 math
+            if _is_hip:
+                out = layer.forward_hip(x)  # 128-bit vectorised kernel from sgl-kernel
+            else:
+                out = layer.forward_cuda(x)
+
+        tol = 1e-2 if dtype is torch.bfloat16 else 1e-3
+        self.assertTrue(
+            torch.allclose(out, ref, atol=tol, rtol=tol),
+            msg=f"Mismatch @ B={n_tok}, D={dim}, dtype={dtype}",
+        )
+        print(f"Match @ B={n_tok}, D={dim}, dtype={dtype}")
+
+    def test_quick_gelu(self):
+        for params in itertools.product(
+            self.NUM_TOKENS, self.DIMS, self.DTYPES, self.SEEDS
+        ):
+            with self.subTest(
+                num_tokens=params[0],
+                dim=params[1],
+                dtype=params[2],
+                seed=params[3],
+            ):
+                self._run_gelu_quick_test(*params)
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=2)

sglang/test/test_utils.py

@@ -1,6 +1,7 @@
 """Common utilities for testing and benchmarking"""
 
 import argparse
+import asyncio
 import copy
 import json
 import logging
@@ -14,8 +15,9 @@ import unittest
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from functools import partial
+from pathlib import Path
 from types import SimpleNamespace
-from typing import Callable, List, Optional, Tuple
+from typing import Awaitable, Callable, List, Optional, Tuple
 
 import numpy as np
 import requests
@@ -26,6 +28,7 @@ from sglang.bench_serving import run_benchmark
 from sglang.global_config import global_config
 from sglang.lang.backend.openai import OpenAI
 from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
+from sglang.lang.interpreter import ProgramState
 from sglang.srt.utils import (
     get_bool_env_var,
     get_device,
@@ -347,6 +350,7 @@ def add_common_sglang_args_and_parse(parser: argparse.ArgumentParser):
         help="Device type (auto/cuda/rocm/cpu). Auto will detect available platforms",
     )
     parser.add_argument("--result-file", type=str, default="result.jsonl")
+    parser.add_argument("--raw-result-file", type=str)
     args = parser.parse_args()
 
     return args
@@ -714,6 +718,7 @@ def get_benchmark_args(
     seed: int = 0,
     device="auto",
     pd_separated: bool = False,
+    lora_name=None,
 ):
     return SimpleNamespace(
         backend="sglang",
@@ -741,7 +746,7 @@ def get_benchmark_args(
         extra_request_body=None,
         apply_chat_template=False,
         profile=None,
-        lora_name=None,
+        lora_name=lora_name,
         prompt_suffix="",
         device=device,
         pd_separated=pd_separated,
@@ -764,6 +769,8 @@ def run_bench_serving(
     need_warmup=False,
     seed: int = 0,
     device="auto",
+    background_task: Optional[Callable[[str, asyncio.Event], Awaitable[None]]] = None,
+    lora_name: Optional[str] = None,
 ):
     if device == "auto":
         device = auto_config_device()
@@ -791,14 +798,35 @@ def run_bench_serving(
         disable_ignore_eos=disable_ignore_eos,
         seed=seed,
         device=device,
+        lora_name=lora_name,
     )
 
-    try:
+    async def _run():
         if need_warmup:
             warmup_args = copy.deepcopy(args)
             warmup_args.num_prompts = 16
-            run_benchmark(warmup_args)
-        res = run_benchmark(args)
+            await asyncio.to_thread(run_benchmark, warmup_args)
+
+        start_event = asyncio.Event()
+        stop_event = asyncio.Event()
+        task_handle = (
+            asyncio.create_task(background_task(base_url, start_event, stop_event))
+            if background_task
+            else None
+        )
+
+        try:
+            start_event.set()
+            result = await asyncio.to_thread(run_benchmark, args)
+        finally:
+            if task_handle:
+                stop_event.set()
+                await task_handle
+
+        return result
+
+    try:
+        res = asyncio.run(_run())
     finally:
         kill_process_tree(process.pid)
 
@@ -1284,3 +1312,35 @@ class CustomTestCase(unittest.TestCase):
             lambda: super(CustomTestCase, self)._callTestMethod(method),
             max_retry=max_retry,
         )
+
+
+def dump_bench_raw_result(
+    path: str,
+    states,
+    preds,
+    labels,
+):
+    if not path:
+        return
+
+    rows = []
+    for i in range(len(states)):
+        state = states[i]
+        output = state["answer"]
+        prompt = _ensure_remove_suffix(state.text(), output)
+        rows.append(
+            dict(
+                prompt_id=i,
+                prompt=prompt,
+                output=output,
+                correct=bool(preds[i] == labels[i]),
+            )
+        )
+
+    print(f"BenchRawResultDumper save results to {path}")
+    Path(path).write_text("\n".join(json.dumps(row) for row in rows))
+
+
+def _ensure_remove_suffix(text: str, suffix: str):
+    assert text.endswith(suffix)
+    return text.removesuffix(suffix)

sglang/utils.py

@@ -14,6 +14,7 @@ import traceback
 import urllib.request
 import weakref
 from concurrent.futures import ThreadPoolExecutor
+from functools import wraps
 from io import BytesIO
 from json import dumps
 from typing import Any, Callable, List, Optional, Tuple, Type, Union
@@ -28,6 +29,24 @@ from tqdm import tqdm
 logger = logging.getLogger(__name__)
 
 
+def execute_once(func):
+    has_run = None
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        nonlocal has_run
+        if not has_run:
+            func(*args, **kwargs)
+            has_run = True
+
+    return wrapper
+
+
+@execute_once
+def info_once(message: str):
+    logger.info(message)
+
+
 def convert_json_schema_to_str(json_schema: Union[dict, str, Type[BaseModel]]) -> str:
     """Convert a JSON schema to a string.
     Parameters

sglang/version.py

@@ -1 +1 @@
-__version__ = "0.4.9.post3"
+__version__ = "0.4.9.post5"

{sglang-0.4.9.post3.dist-info → sglang-0.4.9.post5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sglang
-Version: 0.4.9.post3
+Version: 0.4.9.post5
 Summary: SGLang is yet another fast serving framework for large language models and vision language models.
 License:                                  Apache License
                                    Version 2.0, January 2004
@@ -246,20 +246,20 @@ Requires-Dist: sentencepiece; extra == "runtime-common"
 Requires-Dist: soundfile==0.13.1; extra == "runtime-common"
 Requires-Dist: scipy; extra == "runtime-common"
 Requires-Dist: torchao==0.9.0; extra == "runtime-common"
-Requires-Dist: transformers==4.53.2; extra == "runtime-common"
+Requires-Dist: transformers==4.54.0; extra == "runtime-common"
 Requires-Dist: timm==1.0.16; extra == "runtime-common"
 Requires-Dist: uvicorn; extra == "runtime-common"
 Requires-Dist: uvloop; extra == "runtime-common"
 Requires-Dist: xgrammar==0.1.21; extra == "runtime-common"
 Provides-Extra: srt
 Requires-Dist: sglang[runtime_common]; extra == "srt"
-Requires-Dist: sgl-kernel==0.2.6.post1; extra == "srt"
+Requires-Dist: sgl-kernel==0.2.7; extra == "srt"
 Requires-Dist: torch==2.7.1; extra == "srt"
 Requires-Dist: torchaudio==2.7.1; extra == "srt"
 Requires-Dist: torchvision==0.22.1; extra == "srt"
 Requires-Dist: cuda-python; extra == "srt"
 Requires-Dist: einops; extra == "srt"
-Requires-Dist: flashinfer_python==0.2.7.post1; extra == "srt"
+Requires-Dist: flashinfer_python==0.2.9rc2; extra == "srt"
 Provides-Extra: blackwell
 Requires-Dist: sglang[runtime_common]; extra == "blackwell"
 Requires-Dist: sgl-kernel; extra == "blackwell"
@@ -268,11 +268,11 @@ Requires-Dist: torchaudio==2.7.1; extra == "blackwell"
 Requires-Dist: torchvision==0.22.1; extra == "blackwell"
 Requires-Dist: cuda-python; extra == "blackwell"
 Requires-Dist: einops; extra == "blackwell"
-Requires-Dist: flashinfer_python==0.2.7.post1; extra == "blackwell"
+Requires-Dist: flashinfer_python==0.2.9rc2; extra == "blackwell"
 Provides-Extra: srt-hip
 Requires-Dist: sglang[runtime_common]; extra == "srt-hip"
 Requires-Dist: torch; extra == "srt-hip"
-Requires-Dist: petit_kernel; extra == "srt-hip"
+Requires-Dist: petit_kernel==0.0.2; extra == "srt-hip"
 Provides-Extra: srt-xpu
 Requires-Dist: sglang[runtime_common]; extra == "srt-xpu"
 Provides-Extra: srt-hpu