sglang 0.5.2rc1__py3-none-any.whl → 0.5.3rc0__py3-none-any.whl

sglang/srt/managers/scheduler_metrics_mixin.py

@@ -1,15 +1,24 @@
+from __future__ import annotations
+
 import logging
 import time
 from collections import defaultdict
-from typing import List, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+import torch
 
 from sglang.srt.disaggregation.kv_events import EventPublisherFactory, KVEventBatch
 from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.io_struct import TokenizedGenerateReqInput
 from sglang.srt.managers.schedule_policy import PrefillAdder
 from sglang.srt.managers.scheduler import Req, ScheduleBatch
+from sglang.srt.managers.utils import DPBalanceMeta
 from sglang.srt.metrics.collector import SchedulerMetricsCollector, SchedulerStats
 from sglang.srt.utils import get_bool_env_var
 
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Scheduler
+
 logger = logging.getLogger(__name__)
 
 RECORD_STEP_TIME = get_bool_env_var("SGLANG_RECORD_STEP_TIME")
@@ -28,7 +37,9 @@ class KvMetrics:
 
 
 class SchedulerMetricsMixin:
-    def init_metrics(self, tp_rank: int, pp_rank: int, dp_rank: Optional[int]):
+    def init_metrics(
+        self: Scheduler, tp_rank: int, pp_rank: int, dp_rank: Optional[int]
+    ):
         self.last_gen_throughput: float = 0.0
         self.last_input_throughput: float = 0.0
         self.step_time_dict = defaultdict(list)  # Dict[batch size -> step time]
@@ -50,14 +61,24 @@ class SchedulerMetricsMixin:
                 labels["dp_rank"] = dp_rank
             self.metrics_collector = SchedulerMetricsCollector(labels=labels)
 
-    def init_kv_events(self, kv_events_config: Optional[str]):
+    def init_dp_balance(self: Scheduler, dp_balance_meta: Optional[DPBalanceMeta]):
+        self.balance_meta = dp_balance_meta
+        if (
+            self.server_args.enable_dp_attention
+            and self.server_args.load_balance_method == "minimum_tokens"
+        ):
+            assert dp_balance_meta is not None
+
+        self.recv_dp_balance_id_this_term = []
+
+    def init_kv_events(self: Scheduler, kv_events_config: Optional[str]):
         if self.enable_kv_cache_events:
             self.kv_event_publisher = EventPublisherFactory.create(
                 kv_events_config, self.attn_dp_rank
             )
 
     def log_prefill_stats(
-        self,
+        self: Scheduler,
         adder: PrefillAdder,
         can_run_list: List[Req],
         running_bs: int,
@@ -138,7 +159,7 @@ class SchedulerMetricsMixin:
         self._publish_kv_events()
 
     def log_decode_stats(
-        self, can_run_cuda_graph: bool, running_batch: ScheduleBatch = None
+        self: Scheduler, can_run_cuda_graph: bool, running_batch: ScheduleBatch = None
     ):
         batch = running_batch or self.running_batch
 
@@ -193,7 +214,7 @@ class SchedulerMetricsMixin:
             msg += f"#retracted-req: {len(self.disagg_decode_prealloc_queue.retracted_queue)}, "
 
         msg += (
-            f"cuda graph: {can_run_cuda_graph}, "
+            f"{'cpu graph' if self.device == 'cpu' else 'cuda graph'}: {can_run_cuda_graph}, "
             f"gen throughput (token/s): {self.last_gen_throughput:.2f}, "
             f"#queue-req: {len(self.waiting_queue)}, "
         )
@@ -209,7 +230,7 @@ class SchedulerMetricsMixin:
             self.stats.num_grammar_queue_reqs = len(self.grammar_queue)
             self.stats.spec_accept_length = spec_accept_length
             self.stats.total_retracted_reqs = self.total_retracted_reqs
-            self.metrics_collector.log_stats(self.stats)
+            self.stats.avg_request_queue_latency = 0.0
             if self.disaggregation_mode == DisaggregationMode.DECODE:
                 self.stats.num_decode_prealloc_queue_reqs = len(
                     self.disagg_decode_prealloc_queue.queue
@@ -217,10 +238,11 @@ class SchedulerMetricsMixin:
                 self.stats.num_decode_transfer_queue_reqs = len(
                     self.disagg_decode_transfer_queue.queue
                 )
+            self.metrics_collector.log_stats(self.stats)
             self._emit_kv_metrics()
         self._publish_kv_events()
 
-    def _emit_kv_metrics(self):
+    def _emit_kv_metrics(self: Scheduler):
         kv_metrics = KvMetrics()
         kv_metrics.request_active_slots = self.stats.num_running_reqs
         kv_metrics.request_total_slots = self.max_running_requests
@@ -236,9 +258,94 @@ class SchedulerMetricsMixin:
         if not self.send_metrics_from_scheduler.closed:
             self.send_metrics_from_scheduler.send_pyobj(kv_metrics)
 
-    def _publish_kv_events(self):
+    def _publish_kv_events(self: Scheduler):
         if self.enable_kv_cache_events:
             events = self.tree_cache.take_events()
             if events:
                 batch = KVEventBatch(ts=time.time(), events=events)
                 self.kv_event_publisher.publish(batch)
+
+    def maybe_update_dp_balance_data(
+        self: Scheduler, recv_req: TokenizedGenerateReqInput
+    ):
+        if (
+            self.server_args.enable_dp_attention
+            and self.server_args.load_balance_method == "minimum_tokens"
+        ):
+            self.recv_dp_balance_id_this_term.append(recv_req.dp_balance_id)
+
+    def maybe_handle_dp_balance_data(self: Scheduler):
+        if (
+            self.server_args.load_balance_method == "minimum_tokens"
+            and self.forward_ct % 40 == 0
+        ):
+            holding_tokens = self.get_load().num_tokens
+
+            new_recv_dp_balance_id_list, holding_token_list = (
+                self.gather_dp_balance_info(holding_tokens)
+            )
+
+            self.recv_dp_balance_id_this_term.clear()
+            if self.tp_rank == 0:  # only first worker write info
+                self.write_shared_dp_balance_info(
+                    new_recv_dp_balance_id_list, holding_token_list
+                )
+
+    def gather_dp_balance_info(
+        self: Scheduler, holding_tokens_list
+    ) -> Union[None, List[List[int]]]:
+        """gather recv_dp_balance_id_this_term and holding tokens per worker for dp balance"""
+        recv_list = self.recv_dp_balance_id_this_term
+        assert len(recv_list) <= 511, (
+            "The number of requests received this round is too large. "
+            "Please increase gather_tensor_size and onfly_info_size."
+        )
+        # The maximum size of the tensor used for gathering data from all workers.
+        gather_tensor_size = 512
+
+        # recv_tensor: | holding_tokens | len(recv_dp_balance_id) | recv_dp_balance_ids
+        recv_tensor = torch.zeros(gather_tensor_size, dtype=torch.int32)
+        recv_tensor[0] = holding_tokens_list
+        recv_tensor[1] = len(recv_list)  # The first element is the length of the list.
+        recv_tensor[2 : len(recv_list) + 2] = torch.tensor(recv_list, dtype=torch.int32)
+
+        if self.tp_rank == 0:
+            gathered_list = [
+                torch.zeros(gather_tensor_size, dtype=torch.int32)
+                for _ in range(self.balance_meta.num_workers)
+            ]
+        else:
+            gathered_list = None
+
+        torch.distributed.gather(recv_tensor, gathered_list, group=self.tp_cpu_group)
+
+        gathered_id_list_per_worker = None
+        if self.tp_rank == 0:
+            gathered_id_list_per_worker = []
+            holding_tokens_list = []
+            for tensor in gathered_list:
+                holding_tokens_list.append(tensor[0].item())
+                list_length = tensor[1].item()
+                gathered_id_list_per_worker.append(tensor[2 : list_length + 2].tolist())
+
+        return gathered_id_list_per_worker, holding_tokens_list
+
+    def write_shared_dp_balance_info(self: Scheduler, new_recv_rid_lists, local_tokens):
+        meta = self.balance_meta
+
+        with meta.mutex:
+            onfly_list: List[Dict[int, int]] = meta.get_shared_onfly()
+            assert len(new_recv_rid_lists) == len(onfly_list), "num_worker not equal"
+            # 1.Check if the rid received by each worker this round is present in onfly.
+            #   If it is, remove the corresponding onfly item.
+            worker_id = 0
+            for new_recv_rids, on_fly_reqs in zip(new_recv_rid_lists, onfly_list):
+                for new_recv_rid in new_recv_rids:
+                    assert (
+                        new_recv_rid in on_fly_reqs
+                    ), f"{new_recv_rid=} not in {worker_id=} {on_fly_reqs=}, data consistency is wrong"
+                    del on_fly_reqs[new_recv_rid]
+                worker_id += 1
+            # 2. Atomically write local_tokens and onfly into shm under the mutex
+            meta.set_shared_onfly_info(onfly_list)
+            meta.set_shared_local_tokens(local_tokens)

sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -5,6 +5,8 @@ import threading
 import time
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
+import torch
+
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.io_struct import AbortReq, BatchEmbeddingOut, BatchTokenIDOut
@@ -71,6 +73,7 @@ class SchedulerOutputProcessorMixin:
 
             # Check finish conditions
             logprob_pt = 0
+
             for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
                 if req.is_retracted:
                     continue
@@ -99,6 +102,7 @@ class SchedulerOutputProcessorMixin:
                         extend_logprob_start_len = extend_logprob_start_len_per_req[i]
                         extend_input_len = extend_input_len_per_req[i]
                         num_input_logprobs = extend_input_len - extend_logprob_start_len
+
                         if req.return_logprob:
                             self.add_logprob_return_values(
                                 i,
@@ -441,27 +445,59 @@ class SchedulerOutputProcessorMixin:
         output: LogitsProcessorOutput,
     ):
         """Attach logprobs to the return values."""
-        req.output_token_logprobs_val.append(output.next_token_logprobs[i])
-        req.output_token_logprobs_idx.append(next_token_ids[i])
-
-        self.add_input_logprob_return_values(
-            i, req, output, pt, num_input_logprobs, last_prefill_chunk=True
-        )
+        if output.next_token_logprobs is not None:
+            req.output_token_logprobs_val.append(output.next_token_logprobs[i])
+            req.output_token_logprobs_idx.append(next_token_ids[i])
+
+        # Only add input logprobs if there are input tokens to process
+        # Note: For prefill-only requests with default logprob_start_len, this will be 0,
+        # meaning we only compute output logprobs (which is the intended behavior)
+        if num_input_logprobs > 0:
+            self.add_input_logprob_return_values(
+                i, req, output, pt, num_input_logprobs, last_prefill_chunk=True
+            )
+        else:
+            self._initialize_empty_logprob_containers(req)
 
         if req.top_logprobs_num > 0:
             req.output_top_logprobs_val.append(output.next_token_top_logprobs_val[i])
             req.output_top_logprobs_idx.append(output.next_token_top_logprobs_idx[i])
 
-        if req.token_ids_logprob is not None:
-            req.output_token_ids_logprobs_val.append(
-                output.next_token_token_ids_logprobs_val[i]
-            )
+        if (
+            req.token_ids_logprob is not None
+            and output.next_token_token_ids_logprobs_val is not None
+        ):
+            # Convert GPU tensor to list if needed
+            logprobs_val = output.next_token_token_ids_logprobs_val[i]
+            if isinstance(logprobs_val, torch.Tensor):
+                logprobs_val = logprobs_val.tolist()
+            req.output_token_ids_logprobs_val.append(logprobs_val)
             req.output_token_ids_logprobs_idx.append(
                 output.next_token_token_ids_logprobs_idx[i]
             )
 
         return num_input_logprobs
 
+    def _initialize_empty_logprob_containers(self, req: Req) -> None:
+        """
+        Initialize logprob fields to empty lists if unset.
+
+        This is needed for prefill-only requests where the normal initialization
+        flow might be bypassed, but downstream code expects these fields to be lists.
+        """
+        if req.input_token_logprobs_val is None:
+            req.input_token_logprobs_val = []
+        if req.input_token_logprobs_idx is None:
+            req.input_token_logprobs_idx = []
+        if req.input_top_logprobs_val is None:
+            req.input_top_logprobs_val = []
+        if req.input_top_logprobs_idx is None:
+            req.input_top_logprobs_idx = []
+        if req.input_token_ids_logprobs_val is None:
+            req.input_token_ids_logprobs_val = []
+        if req.input_token_ids_logprobs_idx is None:
+            req.input_token_ids_logprobs_idx = []
+
     def stream_output(
         self: Scheduler,
         reqs: List[Req],
@@ -700,6 +736,8 @@ class SchedulerOutputProcessorMixin:
                     output_token_ids_logprobs_val,
                     output_token_ids_logprobs_idx,
                     output_hidden_states,
+                    placeholder_tokens_idx=None,
+                    placeholder_tokens_val=None,
                 )
             )
 
@@ -719,6 +757,12 @@ class SchedulerOutputProcessorMixin:
                 cached_tokens.append(req.cached_tokens)
         self.send_to_detokenizer.send_pyobj(
             BatchEmbeddingOut(
-                rids, finished_reasons, embeddings, prompt_tokens, cached_tokens
+                rids,
+                finished_reasons,
+                embeddings,
+                prompt_tokens,
+                cached_tokens,
+                placeholder_tokens_idx=None,
+                placeholder_tokens_val=None,
             )
         )

sglang/srt/managers/scheduler_profiler_mixin.py

@@ -26,7 +26,7 @@ logger = logging.getLogger(__name__)
 
 class SchedulerProfilerMixin:
 
-    def init_profier(self):
+    def init_profiler(self):
         self.torch_profiler = None
         self.torch_profiler_output_dir: Optional[str] = None
         self.profiler_activities: Optional[List[str]] = None

sglang/srt/managers/template_manager.py

@@ -24,20 +24,20 @@ import os
 import re
 from typing import Optional
 
-from sglang.srt.code_completion_parser import (
+from sglang.srt.parser.code_completion_parser import (
     CompletionTemplate,
     FimPosition,
     completion_template_exists,
     register_completion_template,
 )
-from sglang.srt.conversation import (
+from sglang.srt.parser.conversation import (
     Conversation,
     SeparatorStyle,
     chat_template_exists,
     get_conv_template_by_model_path,
     register_conv_template,
 )
-from sglang.srt.jinja_template_utils import detect_jinja_template_content_format
+from sglang.srt.parser.jinja_template_utils import detect_jinja_template_content_format
 
 logger = logging.getLogger(__name__)