sglang 0.5.3rc0__py3-none-any.whl → 0.5.3rc2__py3-none-any.whl

sglang/srt/tracing/trace.py

@@ -15,7 +15,6 @@
 
 from __future__ import annotations
 
-import ctypes
 import logging
 import os
 import random
@@ -23,7 +22,10 @@ import threading
 import time
 import uuid
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Req
 
 logger = logging.getLogger(__name__)
 opentelemetry_imported = False
@@ -407,9 +409,11 @@ def trace_slice_start(
     ts: Optional[int] = None,
     anonymous: bool = False,
 ):
+    if not tracing_enabled:
+        return
 
     rid = str(rid)
-    if not tracing_enabled or rid not in reqs_context:
+    if rid not in reqs_context:
         return
 
     pid = threading.get_native_id()
@@ -458,8 +462,11 @@ def trace_slice_end(
     auto_next_anon: bool = False,
     thread_finish_flag: bool = False,
 ):
+    if not tracing_enabled:
+        return
+
     rid = str(rid)
-    if not tracing_enabled or rid not in reqs_context:
+    if rid not in reqs_context:
         return
 
     pid = threading.get_native_id()
@@ -512,10 +519,13 @@ 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):
-    if not tracing_enabled or rid not in reqs_context:
+    if not tracing_enabled:
         return
 
     rid = str(rid)
+    if rid not in reqs_context:
+        return
+
     pid = threading.get_native_id()
     if pid not in reqs_context[rid].threads_context:
         return
@@ -534,10 +544,13 @@ def trace_event(name: str, rid: str, ts: Optional[int] = None):
 
 # Add attrs to the current slice on the same thread with the same rid.
 def trace_slice_add_attr(rid: str, attrs: Dict[str, Any]):
-    if not tracing_enabled or rid not in reqs_context:
+    if not tracing_enabled:
         return
 
     rid = str(rid)
+    if rid not in reqs_context:
+        return
+
     pid = threading.get_native_id()
     if pid not in reqs_context[rid].threads_context:
         return
@@ -550,3 +563,16 @@ def trace_slice_add_attr(rid: str, attrs: Dict[str, Any]):
 
     slice_info = thread_context.cur_slice_stack[-1]
     slice_info.span.set_attributes(attrs)
+
+
+def trace_slice_batch(
+    name: str,
+    reqs: List[Req],
+):
+    for req in reqs:
+        trace_slice(
+            name,
+            req.rid,
+            auto_next_anon=not req.finished(),
+            thread_finish_flag=req.finished(),
+        )

sglang/srt/two_batch_overlap.py

@@ -30,8 +30,9 @@ from sglang.srt.model_executor.forward_batch_info import (
 )
 from sglang.srt.operations import execute_operations, execute_overlapped_operations
 from sglang.srt.operations_strategy import OperationsStrategy
-from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
-from sglang.srt.utils import BumpAllocator, get_bool_env_var, is_hip
+from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
+from sglang.srt.speculative.spec_info import SpecInput
+from sglang.srt.utils import BumpAllocator, empty_context, get_bool_env_var, is_hip
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.token_dispatcher import DispatchOutput
@@ -48,7 +49,7 @@ logger = logging.getLogger(__name__)
 
 def get_token_num_per_seq(
     forward_mode: ForwardMode,
-    spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]] = None,
+    spec_info: Optional[SpecInput] = None,
 ):
     if forward_mode.is_target_verify():
         return spec_info.draft_token_num
@@ -273,7 +274,7 @@ def compute_split_token_index(
 def compute_split_indices_for_cuda_graph_replay(
     forward_mode: ForwardMode,
     cuda_graph_num_tokens: int,
-    spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    spec_info: Optional[SpecInput],
 ):
     forward_mode_for_tbo_split = (
         forward_mode if forward_mode != ForwardMode.IDLE else ForwardMode.DECODE
@@ -333,7 +334,7 @@ class TboCudaGraphRunnerPlugin:
         forward_mode: ForwardMode,
         bs: int,
         num_token_non_padded: int,
-        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        spec_info: Optional[SpecInput],
     ):
         token_num_per_seq = get_token_num_per_seq(
             forward_mode=forward_mode, spec_info=spec_info
@@ -704,6 +705,8 @@ class TboForwardBatchPreparer:
                 extend_num_tokens=extend_num_tokens,
                 attn_backend=output_attn_backend,
                 num_token_non_padded=out_num_token_non_padded,
+                # TODO: handle it when we need TBO + DeepSeek V3.2
+                num_token_non_padded_cpu=None,
                 tbo_split_seq_index=None,
                 tbo_parent_token_range=(start_token_index, end_token_index),
                 tbo_children=None,

sglang/srt/utils/__init__.py

@@ -0,0 +1,2 @@
+# Temporarily do this to avoid changing all imports in the repo
+from .common import *