sglang 0.5.4.post1__py3-none-any.whl → 0.5.4.post2__py3-none-any.whl

sglang/srt/tracing/trace.py

@@ -15,6 +15,8 @@
 
 from __future__ import annotations
 
+import base64
+import json
 import logging
 import os
 import random
@@ -24,6 +26,8 @@ import uuid
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
+from sglang.srt.utils import get_int_env_var
+
 if TYPE_CHECKING:
     from sglang.srt.managers.scheduler import Req
 
@@ -85,6 +89,8 @@ class SglangTraceReqContext:
     # Indicates whether this instance is a replica from the main process.
     # When True, root_span is None and only root_span_context is preserved.
     is_copy: bool = False
+    bootstrap_room_span: Optional[trace.span.Span] = None
+    bootstrap_room_span_context: Optional[context.Context] = None
     root_span: Optional[trace.span.Span] = None
     root_span_context: Optional[context.Context] = None
 
@@ -96,8 +102,7 @@ class SglangTracePropagateContext:
 
     def to_dict(self):
         carrier: dict[str, str] = {}
-        context.attach(self.root_span_context)
-        propagate.inject(carrier)
+        propagate.inject(carrier, self.root_span_context)
 
         if self.prev_span_context:
             return {
@@ -149,6 +154,7 @@ class SglangTraceCustomIdGenerator(id_generator.IdGenerator):
 
 
 # global variables
+remote_trace_contexts: Dict[str, SglangTracePropagateContext] = {}
 threads_info: Dict[int, SglangTraceThreadInfo] = {}
 reqs_context: Dict[str, SglangTraceReqContext] = {}
 
@@ -193,8 +199,17 @@ def process_tracing_init(otlp_endpoint, server_name):
             resource=resource, id_generator=SglangTraceCustomIdGenerator()
         )
 
+        schedule_delay_millis = get_int_env_var(
+            "SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS", 500
+        )
+        max_export_batch_size = get_int_env_var(
+            "SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE", 64
+        )
+
         processor = BatchSpanProcessor(
-            OTLPSpanExporter(endpoint=otlp_endpoint, insecure=True)
+            OTLPSpanExporter(endpoint=otlp_endpoint, insecure=True),
+            schedule_delay_millis=schedule_delay_millis,
+            max_export_batch_size=max_export_batch_size,
         )
         tracer_provider.add_span_processor(processor)
         trace.set_tracer_provider(tracer_provider)
@@ -266,7 +281,9 @@ def __create_thread_context(pid, req_span_context, ts: Optional[int] = None):
     return thread_context
 
 
-def trace_get_proc_propagate_context(rid) -> Optional[Dict[str, Any]]:
+def trace_get_proc_propagate_context(
+    rid, remote_propagate=False
+) -> Optional[Dict[str, Any]]:
     if not tracing_enabled:
         return None
 
@@ -283,9 +300,11 @@ def trace_get_proc_propagate_context(rid) -> Optional[Dict[str, Any]]:
     elif thread_context.last_span_context:
         prev_span_context = thread_context.last_span_context
 
-    trace_context = SglangTracePropagateContext(
-        reqs_context[rid].root_span_context, prev_span_context
-    )
+    root_span_context = reqs_context[rid].root_span_context
+    if remote_propagate:
+        root_span_context = reqs_context[rid].bootstrap_room_span_context
+
+    trace_context = SglangTracePropagateContext(root_span_context, prev_span_context)
     return trace_context.to_dict()
 
 
@@ -327,10 +346,54 @@ def trace_set_proc_propagate_context(rid, trace_context: Optional[Dict[str, Any]
     ].last_span_context = trace_context.prev_span_context
 
 
+def trace_get_remote_propagate_context(bootstrap_room_list: List[str]):
+    if not tracing_enabled:
+        return ""
+
+    reqs_trace_contexts = {}
+    for bootstrap_room in bootstrap_room_list:
+        # In the router, rid is also the bootstrap room.
+        bootstrap_room = str(bootstrap_room)
+
+        if bootstrap_room not in reqs_context:
+            continue
+
+        _context = trace_get_proc_propagate_context(
+            bootstrap_room, remote_propagate=True
+        )
+        reqs_trace_contexts[bootstrap_room] = _context
+
+    json_str = json.dumps(reqs_trace_contexts, ensure_ascii=False)
+    return base64.b64encode(json_str.encode("utf-8")).decode("utf-8")
+
+
+def trace_set_remote_propagate_context(base64_str):
+    if not tracing_enabled:
+        return
+
+    if base64_str is None or base64_str == "" or base64_str == "None":
+        return
+
+    base64_bytes = base64.b64decode(base64_str)
+    json_str = base64_bytes.decode("utf-8")
+    remote_reqs_trace_contexts = json.loads(json_str)
+
+    for bootstrap_room in remote_reqs_trace_contexts:
+        if bootstrap_room in remote_trace_contexts:
+            continue
+
+        remote_trace_contexts[bootstrap_room] = (
+            SglangTracePropagateContext.instance_from_dict(
+                remote_reqs_trace_contexts[bootstrap_room]
+            )
+        )
+
+
 def trace_req_start(
     rid: str,
     bootstrap_room: Optional[int] = None,
     ts: Optional[int] = None,
+    role: Optional[str] = "null",
 ):
     if not tracing_enabled:
         return
@@ -344,6 +407,7 @@ def trace_req_start(
         return
 
     # create req context and root span
+    bootstrap_room = 0 if bootstrap_room is None else bootstrap_room
     reqs_context[rid] = SglangTraceReqContext(
         rid=rid,
         start_time_ns=ts,
@@ -352,23 +416,42 @@ def trace_req_start(
         is_copy=False,
     )
 
+    # create bootstrap room span
+    tracer = threads_info[pid].tracer
+    if str(bootstrap_room) not in remote_trace_contexts:
+        attrs = {"bootstrap_room": str(hex(bootstrap_room))}
+        bootstrap_room_span = tracer.start_span(
+            name=f"Bootstrap Room {hex(bootstrap_room)}",
+            start_time=ts,
+            attributes=attrs,
+        )
+        reqs_context[rid].bootstrap_room_span = bootstrap_room_span
+        bootstrap_room_span_context = trace.set_span_in_context(bootstrap_room_span)
+    else:
+        bootstrap_room_span_context = remote_trace_contexts[
+            str(bootstrap_room)
+        ].root_span_context
+
     # Drop the worker_id added by MultiTokenizer
     orig_rid = rid.split("_")[-1]
-    tracer = threads_info[pid].tracer
+    role = "" if role == "null" else role
+    attrs = {"rid": orig_rid}
     root_span = tracer.start_span(
-        name=f"Req {orig_rid[:8]}",
+        name=f"{role} Req {orig_rid[:8]}",
         start_time=ts,
+        context=bootstrap_room_span_context,
+        attributes=attrs,
     )
 
     root_span.set_attributes(
         {
             "rid": rid,
-            "bootstrap_room": bootstrap_room if bootstrap_room else "None",
         }
     )
 
     reqs_context[rid].root_span = root_span
     reqs_context[rid].root_span_context = trace.set_span_in_context(root_span)
+    reqs_context[rid].bootstrap_room_span_context = bootstrap_room_span_context
 
     # create thread context and thread span
     reqs_context[rid].threads_context[pid] = __create_thread_context(
@@ -376,6 +459,10 @@ def trace_req_start(
         reqs_context[rid].root_span_context,
         ts,
     )
+    if str(bootstrap_room) in remote_trace_contexts:
+        reqs_context[rid].threads_context[pid].last_span_context = (
+            remote_trace_contexts[str(bootstrap_room)].prev_span_context
+        )
 
 
 def trace_req_finish(
@@ -399,6 +486,10 @@ def trace_req_finish(
         req_context.root_span.set_attributes(attrs)
 
     req_context.root_span.end(end_time=ts)
+    if str(req_context.bootstrap_room) in remote_trace_contexts:
+        del remote_trace_contexts[str(req_context.bootstrap_room)]
+    else:
+        req_context.bootstrap_room_span.end(end_time=ts)
 
     del reqs_context[rid]
 
@@ -518,7 +609,9 @@ trace_slice = trace_slice_end
 
 
 # Add event to the current slice on the same thread with the same rid.
-def trace_event(name: str, rid: str, ts: Optional[int] = None):
+def trace_event(
+    name: str, rid: str, ts: Optional[int] = None, attrs: Dict[str, Any] = None
+):
     if not tracing_enabled:
         return
 
@@ -539,7 +632,7 @@ def trace_event(name: str, rid: str, ts: Optional[int] = None):
     ts = ts or __get_cur_time_ns()
 
     slice_info = thread_context.cur_slice_stack[-1]
-    slice_info.span.add_event(name=name, timestamp=ts)
+    slice_info.span.add_event(name=name, timestamp=ts, attributes=attrs)
 
 
 # Add attrs to the current slice on the same thread with the same rid.
@@ -569,6 +662,9 @@ def trace_slice_batch(
     name: str,
     reqs: List[Req],
 ):
+    if not tracing_enabled:
+        return
+
     for req in reqs:
         trace_slice(
             name,
@@ -576,3 +672,16 @@ def trace_slice_batch(
             auto_next_anon=not req.finished(),
             thread_finish_flag=req.finished(),
         )
+
+
+def trace_event_batch(
+    name: str,
+    reqs: List[Req],
+    ts: Optional[int] = None,
+    attrs: Dict[str, Any] = None,
+):
+    if not tracing_enabled:
+        return
+
+    for req in reqs:
+        trace_event(name, req.rid, ts=ts, attrs=attrs)

sglang/srt/utils/common.py

@@ -188,7 +188,16 @@ is_hopper_with_cuda_12_3 = lambda: _check(9)
 def is_blackwell():
     if not is_cuda():
         return False
-    return torch.cuda.get_device_capability()[0] == 10
+    return torch.cuda.get_device_capability()[0] in [10, 12]
+
+
+@lru_cache(maxsize=1)
+def is_blackwell_supported(device=None) -> bool:
+    if not is_cuda_alike():
+        return False
+    return (torch.cuda.get_device_capability(device)[0] in [10, 12]) and (
+        torch.version.cuda >= "12.8"
+    )
 
 
 @lru_cache(maxsize=1)
@@ -1230,42 +1239,34 @@ def point_to_point_pyobj(
     dst: int = 1,
 ):
     """Send data from src to dst in group using DeviceToDevice communication."""
-
+    device = torch.get_device_module().current_device()
     if rank == src:
         if len(data) == 0:
-            tensor_size = torch.tensor(
-                [0], dtype=torch.long, device=torch.cuda.current_device()
-            )
+            tensor_size = torch.tensor([0], dtype=torch.long, device=device)
             dist.send(tensor_size, dst=dst, group=group)
         else:
             serialized_data = pickle.dumps(data)
             size = len(serialized_data)
             tensor_data = torch.ByteTensor(
                 np.frombuffer(serialized_data, dtype=np.uint8)
-            ).cuda(
-                device=torch.cuda.current_device()
+            ).to(
+                device=device
             )  # Move to GPU
-            tensor_size = torch.tensor(
-                [size], dtype=torch.long, device=torch.cuda.current_device()
-            )
+            tensor_size = torch.tensor([size], dtype=torch.long, device=device)
 
             dist.send(tensor_size, dst=dst, group=group)
             dist.send(tensor_data, dst=dst, group=group)
         return data
 
     elif rank == dst:
-        tensor_size = torch.tensor(
-            [0], dtype=torch.long, device=torch.cuda.current_device()
-        )
+        tensor_size = torch.tensor([0], dtype=torch.long, device=device)
         dist.recv(tensor_size, src=src, group=group)
         size = tensor_size.item()
 
         if size == 0:
             return []
 
-        tensor_data = torch.empty(
-            size, dtype=torch.uint8, device=torch.cuda.current_device()
-        )
+        tensor_data = torch.empty(size, dtype=torch.uint8, device=device)
         dist.recv(tensor_data, src=src, group=group)
 
         serialized_data = bytes(
@@ -2350,16 +2351,24 @@ def launch_dummy_health_check_server(host, port, enable_metrics):
     )
     server = uvicorn.Server(config=config)
 
-    try:
-        loop = asyncio.get_running_loop()
-        logger.info(
-            f"Dummy health check server scheduled on existing loop at {host}:{port}"
-        )
-        loop.create_task(server.serve())
+    # Run server in a background daemon thread with its own event loop
+    # This prevents blocking the main thread while still serving health checks
+    def run_server():
+        try:
+            asyncio.run(server.serve())
+        except Exception as e:
+            logger.error(f"Dummy health check server failed to start: {e}")
+            raise
+        finally:
+            logger.info(f"Dummy health check server stopped at {host}:{port}")
 
-    except RuntimeError:
-        logger.info(f"Starting dummy health check server at {host}:{port}")
-        server.run()
+    thread = threading.Thread(
+        target=run_server, daemon=True, name="health-check-server"
+    )
+    thread.start()
+    logger.info(
+        f"Dummy health check server started in background thread at {host}:{port}"
+    )
 
 
 def create_checksum(directory: str):
@@ -3105,12 +3114,16 @@ def apply_module_patch(target_module, target_function, wrappers):
         setattr(original_module, target_function, candidate)
 
     for key, value in sys.modules.copy().items():
-        if (
-            target_function is not None
-            and hasattr(value, target_function)
-            and id(getattr(value, target_function)) == original_function_id
-        ):
-            setattr(value, target_function, candidate)
+        try:
+            if (
+                target_function is not None
+                and hasattr(value, target_function)
+                and id(getattr(value, target_function)) == original_function_id
+            ):
+                setattr(value, target_function, candidate)
+        except ImportError as e:
+            # Ignore some modules reporting ImportError when calling hasattr
+            logger.warning(f"Ignore {value} reports ImportError with:\n{str(e)}")
 
 
 def parse_module_path(module_path, function_name, create_dummy):
@@ -3562,7 +3575,17 @@ def cached_triton_kernel(key_fn=None):
     """
 
     def decorator(fn):
-        return CachedKernel(fn, key_fn)
+        if envs.SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE.get():
+            logger.debug(
+                f"{envs.SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE.name} = True. Using custom triton kernel cache."
+            )
+            return CachedKernel(fn, key_fn)
+        else:
+            # Fallback to the native triton cache.
+            logger.debug(
+                f"{envs.SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE.name} = False. Using native triton kernel cache."
+            )
+            return fn
 
     return decorator
 

sglang/srt/utils/hf_transformers_utils.py

@@ -43,6 +43,7 @@ from sglang.srt.configs import (
     DotsVLMConfig,
     ExaoneConfig,
     FalconH1Config,
+    KimiLinearConfig,
     KimiVLConfig,
     LongcatFlashConfig,
     MultiModalityConfig,
@@ -54,6 +55,7 @@ from sglang.srt.configs import (
 from sglang.srt.configs.deepseek_ocr import DeepseekVLV2Config
 from sglang.srt.configs.internvl import InternVLChatConfig
 from sglang.srt.connector import create_remote_connector
+from sglang.srt.multimodal.customized_mm_processor_utils import _CUSTOMIZED_MM_PROCESSOR
 from sglang.srt.utils import is_remote_url, logger, lru_cache_frozenset
 
 _CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [
@@ -67,6 +69,7 @@ _CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [
     Step3VLConfig,
     LongcatFlashConfig,
     Olmo3Config,
+    KimiLinearConfig,
     Qwen3NextConfig,
     FalconH1Config,
     DotsVLMConfig,
@@ -172,6 +175,16 @@ def _load_deepseek_v32_model(
     )
 
 
+def _is_deepseek_ocr_model(config: PretrainedConfig) -> bool:
+    # TODO: Remove this workaround related when AutoConfig correctly identifies deepseek-ocr.
+    # Hugging Face's AutoConfig currently misidentifies it as deepseekvl2.
+    return (
+        getattr(config, "auto_map", None) is not None
+        and config.auto_map.get("AutoModel")
+        == "modeling_deepseekocr.DeepseekOCRForCausalLM"
+    )
+
+
 @lru_cache_frozenset(maxsize=32)
 def get_config(
     model: str,
@@ -197,14 +210,6 @@ def get_config(
         config = AutoConfig.from_pretrained(
             model, trust_remote_code=trust_remote_code, revision=revision, **kwargs
         )
-        if (
-            getattr(config, "auto_map", None) is not None
-            and config.auto_map.get("AutoModel")
-            == "modeling_deepseekocr.DeepseekOCRForCausalLM"
-        ):
-            config.model_type = "deepseek-ocr"
-            # TODO: Remove this workaround when AutoConfig correctly identifies deepseek-ocr.
-            # Hugging Face's AutoConfig currently misidentifies it as deepseekvl2.
 
     except ValueError as e:
         if not "deepseek_v32" in str(e):
@@ -241,7 +246,11 @@ def get_config(
                 setattr(config, key, val)
 
     if config.model_type in _CONFIG_REGISTRY:
-        config_class = _CONFIG_REGISTRY[config.model_type]
+        model_type = config.model_type
+        if model_type == "deepseek_vl_v2":
+            if _is_deepseek_ocr_model(config):
+                model_type = "deepseek-ocr"
+        config_class = _CONFIG_REGISTRY[model_type]
         config = config_class.from_pretrained(model, revision=revision)
         # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`.
         setattr(config, "_name_or_path", model)
@@ -445,6 +454,10 @@ def get_processor(
         **kwargs,
     )
 
+    if _is_deepseek_ocr_model(config):
+        # Temporary hack for load deepseek-ocr
+        config.model_type = "deepseek-ocr"
+
     # fix: for Qwen2-VL and Sarashina2Vision models, inject default 'size' if not provided.
     if config.model_type in {"qwen2_vl", "sarashina2_vision"}:
         if "size" not in kwargs:
@@ -462,13 +475,22 @@ def get_processor(
                 **kwargs,
             )
         else:
-            processor = AutoProcessor.from_pretrained(
-                tokenizer_name,
-                *args,
-                trust_remote_code=trust_remote_code,
-                revision=revision,
-                **kwargs,
-            )
+            if config.model_type in _CUSTOMIZED_MM_PROCESSOR:
+                processor = _CUSTOMIZED_MM_PROCESSOR[config.model_type].from_pretrained(
+                    tokenizer_name,
+                    *args,
+                    trust_remote_code=trust_remote_code,
+                    revision=revision,
+                    **kwargs,
+                )
+            else:
+                processor = AutoProcessor.from_pretrained(
+                    tokenizer_name,
+                    *args,
+                    trust_remote_code=trust_remote_code,
+                    revision=revision,
+                    **kwargs,
+                )
 
     except ValueError as e:
         error_message = str(e)

sglang/srt/utils/torch_memory_saver_adapter.py

@@ -41,6 +41,12 @@ class TorchMemorySaverAdapter(ABC):
     def region(self, tag: str, enable_cpu_backup: bool = False):
         raise NotImplementedError
 
+    def cuda_graph(self, **kwargs):
+        raise NotImplementedError
+
+    def disable(self):
+        raise NotImplementedError
+
     def pause(self, tag: str):
         raise NotImplementedError
 
@@ -61,6 +67,12 @@ class _TorchMemorySaverAdapterReal(TorchMemorySaverAdapter):
     def region(self, tag: str, enable_cpu_backup: bool = False):
         return _memory_saver.region(tag=tag, enable_cpu_backup=enable_cpu_backup)
 
+    def cuda_graph(self, **kwargs):
+        return _memory_saver.cuda_graph(**kwargs)
+
+    def disable(self):
+        return _memory_saver.disable()
+
     def pause(self, tag: str):
         return _memory_saver.pause(tag=tag)
 
@@ -81,6 +93,14 @@ class _TorchMemorySaverAdapterNoop(TorchMemorySaverAdapter):
     def region(self, tag: str, enable_cpu_backup: bool = False):
         yield
 
+    @contextmanager
+    def cuda_graph(self, **kwargs):
+        yield
+
+    @contextmanager
+    def disable(self):
+        yield
+
     def pause(self, tag: str):
         pass
 

sglang/test/kits/radix_cache_server_kit.py

@@ -0,0 +1,50 @@
+import random
+
+import requests
+
+
+def gen_radix_tree(num_nodes=400, chunk_len=256):
+    num0 = num_nodes // 2
+    num1 = num_nodes - num0
+    nodes = [{"input_ids": [37] * 117, "decode_len": 217}]
+    for _ in range(num0):
+        parent = random.choice(nodes)
+        unique_len = random.randint(0, chunk_len)
+        decode_len = random.randint(0, chunk_len)
+        token_id = random.randint(0, 32000)
+        child = {
+            "input_ids": parent["input_ids"] + [token_id] * unique_len,
+            "decode_len": decode_len,
+        }
+        nodes.append(child)
+
+    while num1 > 0:
+        num_branch = random.randint(1, min(num1, 10))
+        parent = random.choice(nodes)
+        for _ in range(num_branch):
+            unique_len = random.randint(0, chunk_len)
+            decode_len = random.randint(0, chunk_len)
+            token_id = random.randint(0, 32000)
+            child = {
+                "input_ids": parent["input_ids"] + [token_id] * unique_len,
+                "decode_len": decode_len,
+            }
+            nodes.append(child)
+
+        num1 -= num_branch
+
+    random.shuffle(nodes)
+    return nodes
+
+
+def run_radix_attention_test(base_url: str):
+    nodes = gen_radix_tree()
+    data = {
+        "input_ids": [node["input_ids"] for node in nodes],
+        "sampling_params": [
+            {"max_new_tokens": node["decode_len"], "temperature": 0} for node in nodes
+        ],
+    }
+
+    res = requests.post(base_url + "/generate", json=data)
+    assert res.status_code == 200