sglang 0.5.3rc0__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/speculative/standalone_worker.py

@@ -1,29 +1,20 @@
 import logging
-from contextlib import contextmanager
 from typing import Optional
 
 import torch
 
-from sglang.srt.distributed import GroupCoordinator, patch_tensor_parallel_group
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.server_args import ServerArgs
-from sglang.srt.speculative.eagle_worker import EAGLEWorker, load_token_map
+from sglang.srt.speculative.eagle_worker import EAGLEWorker
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_utils import draft_tp_context, load_token_map
 from sglang.srt.utils import empty_context, get_bool_env_var, is_cuda
 
 if is_cuda():
-    from sgl_kernel import segment_packbits
+    from sgl_kernel import segment_packbits  # noqa: F401
 
 logger = logging.getLogger(__name__)
-RETURN_ORIGINAL_LOGPROB = get_bool_env_var("RETURN_ORIGINAL_LOGPROB")
-
-
-@contextmanager
-def draft_tp_context(tp_group: GroupCoordinator):
-    # Draft model doesn't use dp and has its own tp group.
-    # We disable mscclpp now because it doesn't support 2 comm groups.
-    with patch_tensor_parallel_group(tp_group):
-        yield
+SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB")
 
 
 class StandaloneWorker(EAGLEWorker):
@@ -51,7 +42,6 @@ class StandaloneWorker(EAGLEWorker):
         self.speculative_algorithm = SpeculativeAlgorithm.from_string(
             server_args.speculative_algorithm
         )
-        self.padded_static_len = -1
 
         # Override the context length of the draft model to be the same as the target model.
         server_args.context_length = target_worker.model_runner.model_config.context_len

sglang/srt/tokenizer/tiktoken_tokenizer.py

@@ -133,9 +133,9 @@ class TiktokenTokenizer:
         )
         return self.encode(ret) if tokenize else ret
 
-    def __call__(self, text, **kwargs):
+    def __call__(self, text: List[str], **kwargs):
         return {
-            "input_ids": self.encode(text),
+            "input_ids": [self.encode(x) for x in text],
         }
 
     def init_xgrammar(self):

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

@@ -4,10 +4,11 @@ import copy
 import dataclasses
 import logging
 from dataclasses import replace
-from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence
 
 import torch
 
+from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.communicator import (
     CommunicateContext,
@@ -20,9 +21,11 @@ from sglang.srt.layers.moe import (
     get_tbo_token_distribution_threshold,
     is_tbo_enabled,
 )
-from sglang.srt.layers.moe.token_dispatcher import DeepEPDispatcher
-from sglang.srt.layers.quantization import deep_gemm_wrapper
-from sglang.srt.managers.schedule_batch import ScheduleBatch, global_server_args_dict
+from sglang.srt.layers.moe.token_dispatcher import (
+    DeepEPDispatcher,
+    MooncakeEPDispatcher,
+)
+from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.model_executor.forward_batch_info import (
     ForwardBatch,
     ForwardMode,
@@ -30,11 +33,13 @@ 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.server_args import get_global_server_args
+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
+    from sglang.srt.speculative.eagle_info import EagleVerifyInput
 
 _is_hip = is_hip()
 
@@ -48,7 +53,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
@@ -152,7 +157,7 @@ def _update_device_and_sum_field_from_cpu_field(
         cpu_value
         if isinstance(cpu_value, torch.Tensor)
         else torch.tensor(cpu_value, dtype=old_device_value.dtype)
-    ).to(device=global_server_args_dict["device"], non_blocking=True)
+    ).to(device=get_global_server_args().device, non_blocking=True)
     setattr(batch, device_field, new_device_value)
 
     if sum_field is not None:
@@ -273,7 +278,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 +338,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
@@ -361,7 +366,7 @@ class TboDPAttentionPreparer:
     ):
 
         deepep_mode = get_deepep_mode()
-        enable_deepep_moe = get_moe_a2a_backend().is_deepep()
+        enable_a2a_moe = not get_moe_a2a_backend().is_none()
         enable_two_batch_overlap = is_tbo_enabled()
 
         self.enable_two_batch_overlap = enable_two_batch_overlap
@@ -390,7 +395,7 @@ class TboDPAttentionPreparer:
                     local_batch.forward_mode.is_extend()
                     and not local_batch.forward_mode.is_target_verify()
                 )
-                and enable_deepep_moe
+                and enable_a2a_moe
                 and (resolved_deepep_mode.is_low_latency())
             )
         else:
@@ -581,7 +586,7 @@ class TboForwardBatchPreparer:
             sum_field=None,
         )
         _, child_b.extend_start_loc = compute_position(
-            global_server_args_dict["attention_backend"],
+            get_global_server_args().attention_backend,
             child_b.extend_prefix_lens,
             child_b.extend_seq_lens,
             child_b.extend_num_tokens,
@@ -666,6 +671,7 @@ class TboForwardBatchPreparer:
             "can_run_dp_cuda_graph",
             "dp_padding_mode",
             "global_forward_mode",
+            "is_prefill_only",
             "spec_algorithm",
             "capture_hidden_mode",
             "padded_static_len",
@@ -685,7 +691,7 @@ class TboForwardBatchPreparer:
 
         # TODO improve, e.g. unify w/ `init_raw`
         if (
-            global_server_args_dict["moe_dense_tp_size"] == 1
+            get_global_server_args().moe_dense_tp_size == 1
             and batch.global_dp_buffer_len is not None
         ):
             sum_len = end_token_index - start_token_index
@@ -704,6 +710,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,
@@ -751,7 +759,7 @@ class TboForwardBatchPreparer:
         value_a = min(tbo_split_token_index, num_token_non_padded)
         value_b = max(0, num_token_non_padded - tbo_split_token_index)
         return torch.tensor([value_a, value_b], dtype=torch.int32).to(
-            device=global_server_args_dict["device"], non_blocking=True
+            device=get_global_server_args().device, non_blocking=True
         )
 
     @classmethod
@@ -963,9 +971,14 @@ def _model_forward_tbo_merge_outputs(output_a, output_b):
 class MaybeTboDeepEPDispatcher:
     def __init__(self, **kwargs):
         num_inner_dispatchers = 2 if is_tbo_enabled() else 1
-        self._inners = [
-            DeepEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
-        ]
+        if get_moe_a2a_backend().is_deepep():
+            self._inners = [
+                DeepEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
+            ]
+        elif get_moe_a2a_backend().is_mooncake():
+            self._inners = [
+                MooncakeEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
+            ]
 
     def _execute(self, name, tbo_subbatch_index: Optional[int] = None, **kwargs):
         return getattr(self._inners[tbo_subbatch_index or 0], name)(**kwargs)
@@ -987,3 +1000,7 @@ class MaybeTboDeepEPDispatcher:
 
     def combine_b(self, **kwargs):
         return self._execute("combine_b", **kwargs)
+
+    def set_quant_config(self, quant_config: dict):
+        for inner in self._inners:
+            inner.set_quant_config(quant_config)

sglang/srt/utils/__init__.py

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

sglang/srt/{bench_utils.py → utils/bench_utils.py}

@@ -1,4 +1,5 @@
 import os
+import re
 import sys
 from contextlib import nullcontext
 
@@ -108,7 +109,8 @@ def bench_kineto(
     if not with_multiple_kernels:
         for name in kernel_names:
             assert (
-                sum([name in line for line in prof_lines]) == 1
+                sum([int(re.search(name, line) is not None) for line in prof_lines])
+                == 1
             ), f"Errors of the kernel {name} in the profiling table (table: {prof_lines})"
 
     # Save chrome traces
@@ -122,7 +124,7 @@ def bench_kineto(
         total_time = 0
         total_num = 0
         for line in prof_lines:
-            if name in line:
+            if re.search(name, line) is not None:
                 time_str = line.split()[-2]
                 num_str = line.split()[-1]
                 for unit, scale in units.items():