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

sglang/srt/managers/cache_controller.py

@@ -324,6 +324,22 @@ class HiCacheController:
                     group_ranks, backend="gloo"
                 )
 
+            # Select the get and set functions
+            self.page_get_func = self._generic_page_get
+            self.page_set_func = self._generic_page_set
+            self.batch_exists_func = self.storage_backend.batch_exists
+            self.is_3fs_zerocopy = (
+                self.storage_backend_type == "hf3fs"
+                and self.mem_pool_host.layout == "page_first"
+            )
+            if self.storage_backend_type == "mooncake":
+                self.page_get_func = self._mooncake_page_get
+                self.page_set_func = self._mooncake_page_set
+            elif self.is_3fs_zerocopy:
+                self.page_get_func = self._3fs_zero_copy_page_get
+                self.page_set_func = self._3fs_zero_copy_page_set
+                self.batch_exists_func = self._3fs_zero_copy_batch_exists
+
         self.load_cache_event = load_cache_event
         self.layer_done_counter = LayerDoneCounter(self.mem_pool_device.layer_num)
         self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter)
@@ -407,6 +423,7 @@ class HiCacheController:
             tp_rank=self.tp_rank,
             tp_size=self.tp_size,
             is_mla_model=is_mla_backend,
+            is_page_first_layout=self.mem_pool_host.layout == "page_first",
             model_name=model_name,
             extra_config=extra_config,
         )
@@ -616,13 +633,19 @@ class HiCacheController:
         for chunk in chunks:
             self.host_mem_release_queue.put(chunk)
 
+    def _3fs_zero_copy_batch_exists(self, batch_hashes):
+        _batch_hashes, _, factor = self.mem_pool_host.get_buffer_with_hash(batch_hashes)
+        hit_page_num = self.storage_backend.batch_exists(_batch_hashes) // factor
+        return hit_page_num
+
     def _3fs_zero_copy_page_get(self, operation, hash_values, host_indices):
-        hashes, dsts = self.mem_pool_host.get_buffer_with_hash(
+        hashes, dsts, factor = self.mem_pool_host.get_buffer_with_hash(
             hash_values, host_indices
         )
         page_data = self.storage_backend.batch_get(hashes, dsts)
         if page_data:
-            operation.increment(self.page_size * len(hashes))
+            inc = self.page_size * len(hashes) // factor
+            operation.increment(inc)
         else:
             logger.warning(
                 f"Prefetch operation {operation.request_id} failed to retrieve page {hashes}."
@@ -636,7 +659,7 @@ class HiCacheController:
         )
         get_result = self.storage_backend.batch_get(
             key_strs,
-            target_location=buffer_ptrs,
+            target_locations=buffer_ptrs,
             target_sizes=buffer_sizes,
         )
         if get_result != len(hash_values):
@@ -647,9 +670,9 @@ class HiCacheController:
             operation.increment(get_result * self.page_size)
 
     def _generic_page_get(self, operation, hash_values, host_indices):
-        dummy_page_dst = [self.mem_pool_host.get_dummy_flat_data_page()] * len(
-            hash_values
-        )
+        dummy_page_dst = [
+            self.mem_pool_host.get_dummy_flat_data_page() for _ in hash_values
+        ]
         page_data = self.storage_backend.batch_get(hash_values, dummy_page_dst)
         if page_data is None:
             return
@@ -659,26 +682,16 @@ class HiCacheController:
                     f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}."
                 )
                 break
-            if operation.increment(self.page_size):
-                self.mem_pool_host.set_from_flat_data_page(
-                    host_indices[i * self.page_size],
-                    page_data[i],
-                )
-            else:
-                break
+            # Must set the data before increasing the completed tokens.
+            # Otherwise this page may be read before being set.
+            self.mem_pool_host.set_from_flat_data_page(
+                host_indices[i * self.page_size],
+                page_data[i],
+            )
+            if not operation.increment(self.page_size):
+                break  # Operation terminated by controller
 
     def _page_transfer(self, operation):
-        # Select the get function and batch size
-        if self.storage_backend_type == "mooncake":
-            get_func = self._mooncake_page_get
-        elif (
-            self.storage_backend_type == "hf3fs"
-            and self.mem_pool_host.layout == "page_first"
-        ):
-            get_func = self._3fs_zero_copy_page_get
-        else:
-            get_func = self._generic_page_get
-
         # Transfer batch by batch
         for i in range(0, len(operation.hash_value), self.storage_batch_size):
             batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
@@ -687,7 +700,7 @@ class HiCacheController:
             ]
             prev_completed_tokens = operation.completed_tokens
             # Get one batch token, and update the completed_tokens if succeed
-            get_func(operation, batch_hashes, batch_host_indices)
+            self.page_get_func(operation, batch_hashes, batch_host_indices)
             # Check termination
             if (
                 operation.completed_tokens
@@ -744,7 +757,7 @@ class HiCacheController:
                     batch_tokens[i : i + self.page_size], last_hash
                 )
                 batch_hashes.append(last_hash)
-            hit_page_num = self.storage_backend.batch_exists(batch_hashes)
+            hit_page_num = self.batch_exists_func(batch_hashes)
             hash_value.extend(batch_hashes[:hit_page_num])
             storage_query_count += hit_page_num * self.page_size
             if hit_page_num < len(batch_hashes):
@@ -830,30 +843,20 @@ class HiCacheController:
         )
         success = self.storage_backend.batch_set(
             key_strs,
-            target_location=buffer_ptrs,
+            target_locations=buffer_ptrs,
             target_sizes=buffer_sizes,
         )
         return success
 
     # zero copy
     def _3fs_zero_copy_page_set(self, hash_values, host_indices) -> bool:
-        hashes, dsts = self.mem_pool_host.get_buffer_with_hash(
+        hashes, dsts, _ = self.mem_pool_host.get_buffer_with_hash(
             hash_values, host_indices
         )
         return self.storage_backend.batch_set(hashes, dsts)
 
     # Backup batch by batch
     def _page_backup(self, operation):
-        # Select the set function and batch size
-        if self.storage_backend_type == "mooncake":
-            backup_set_func = self._mooncake_page_set
-        elif (
-            self.storage_backend_type == "hf3fs"
-            and self.mem_pool_host.layout == "page_first"
-        ):
-            backup_set_func = self._3fs_zero_copy_page_set
-        else:
-            backup_set_func = self._generic_page_set
         # Backup batch by batch
         for i in range(0, len(operation.hash_value), self.storage_batch_size):
             batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
@@ -862,7 +865,7 @@ class HiCacheController:
             ]
             # Set one batch token, and record if success.
             # todo: allow partial success
-            success = backup_set_func(batch_hashes, batch_host_indices)
+            success = self.page_set_func(batch_hashes, batch_host_indices)
             if not success:
                 logger.warning(
                     f"Write page to storage: {len(batch_hashes)} pages failed."

sglang/srt/managers/detokenizer_manager.py

@@ -39,7 +39,6 @@ from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import (
     configure_logger,
     freeze_gc,
-    get_worker_ids_from_req_rids,
     get_zmq_socket,
     kill_itself_when_parent_died,
 )
@@ -120,39 +119,6 @@ class DetokenizerManager(MultiTokenizerMixin):
             if output is not None:
                 self.send_to_tokenizer.send_pyobj(output)
 
-    def multi_tokenizer_manager_event_loop(self):
-        """The event loop that handles requests, for multi tokenizer manager mode only"""
-        self.create_sockets_mapping()
-        while True:
-            recv_obj = self.recv_from_scheduler.recv_pyobj()
-            output = self._request_dispatcher(recv_obj)
-            if output is None:
-                continue
-            # Extract worker_id from rid
-            if isinstance(recv_obj.rids, list):
-                worker_ids = get_worker_ids_from_req_rids(recv_obj.rids)
-            else:
-                raise RuntimeError(
-                    f"for tokenizer_worker_num > 1, recv_obj.rids must be a list"
-                )
-
-            # Send data using the corresponding socket
-            for i, worker_id in enumerate(worker_ids):
-                if isinstance(recv_obj, MultiTokenizerRegisterReq):
-                    if self.register_tokenizer_ipc(recv_obj, worker_id):
-                        logger.info(
-                            f"DetokenizerManager Created ZMQ socket for worker {worker_id}"
-                        )
-                    continue
-                else:
-                    if worker_id not in self.tokenizer_mapping:
-                        logger.error(
-                            f"Tokenizer Worker ID {worker_id} not registered. Check if the server Process {worker_id} is alive"
-                        )
-                        continue
-                    new_output = self._handle_output_by_index(output, i)
-                    self.tokenizer_mapping[worker_id].send_pyobj(new_output)
-
     def trim_matched_stop(
         self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool
     ):

sglang/srt/managers/multi_tokenizer_mixin.py

@@ -23,6 +23,7 @@ import threading
 from multiprocessing import shared_memory
 from typing import Dict
 
+import setproctitle
 import zmq
 import zmq.asyncio
 
@@ -37,11 +38,7 @@ from sglang.srt.managers.io_struct import (
 )
 from sglang.srt.managers.tokenizer_manager import TokenizerManager, _Communicator
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import (
-    get_worker_ids_from_req_rids,
-    get_zmq_socket,
-    kill_process_tree,
-)
+from sglang.srt.utils import get_zmq_socket, kill_process_tree
 from sglang.utils import get_exception_traceback
 
 logger = logging.getLogger(__name__)
@@ -344,6 +341,48 @@ class MultiTokenizerMixin:
             new_output = output
         return new_output
 
+    def get_worker_ids_from_req_rids(self, rids):
+        if isinstance(rids, list):
+            worker_ids = [int(rid.split("_")[0]) for rid in rids]
+        elif isinstance(rids, str):
+            worker_ids = [int(rids.split("_")[0])]
+        else:
+            worker_ids = []
+        return worker_ids
+
+    def multi_tokenizer_manager_event_loop(self):
+        """The event loop that handles requests, for multi tokenizer manager mode only"""
+        self.create_sockets_mapping()
+        while True:
+            recv_obj = self.recv_from_scheduler.recv_pyobj()
+            output = self._request_dispatcher(recv_obj)
+            if output is None:
+                continue
+            # Extract worker_id from rid
+            if isinstance(recv_obj.rids, list):
+                worker_ids = self.get_worker_ids_from_req_rids(recv_obj.rids)
+            else:
+                raise RuntimeError(
+                    f"for tokenizer_worker_num > 1, recv_obj.rids must be a list"
+                )
+
+            # Send data using the corresponding socket
+            for i, worker_id in enumerate(worker_ids):
+                if isinstance(recv_obj, MultiTokenizerRegisterReq):
+                    if self.register_tokenizer_ipc(recv_obj, worker_id):
+                        logger.info(
+                            f"DetokenizerManager Created ZMQ socket for worker {worker_id}"
+                        )
+                    continue
+                else:
+                    if worker_id not in self.tokenizer_mapping:
+                        logger.error(
+                            f"Tokenizer Worker ID {worker_id} not registered. Check if the server Process {worker_id} is alive"
+                        )
+                        continue
+                    new_output = self._handle_output_by_index(output, i)
+                    self.tokenizer_mapping[worker_id].send_pyobj(new_output)
+
     def clear_tokenizer_mapping(self):
         if hasattr(self, "tokenizer_mapping"):
             for socket in self.tokenizer_mapping.values():
@@ -406,7 +445,7 @@ class MultiTokenizerRouter(TokenizerManager, MultiTokenizerMixin):
             worker_ids = [recv_obj.worker_id]
             recv_obj = recv_obj.obj
         else:
-            worker_ids = get_worker_ids_from_req_rids(recv_obj.rids)
+            worker_ids = self.get_worker_ids_from_req_rids(recv_obj.rids)
 
         if len(worker_ids) == 0:
             logger.error(f"Cannot find worker_id from rids {recv_obj.rids}")
@@ -438,6 +477,9 @@ class MultiTokenizerManager(TokenizerManager, MultiTokenizerMixin):
         server_args: ServerArgs,
         port_args: PortArgs,
     ):
+        setproctitle.setproctitle(
+            f"sglang::http_server/multi_tokenizer_manager:{os.getpid()}"
+        )
         # prevent init prefill bootstrapserver again
         disaggregation_mode = server_args.disaggregation_mode
         server_args.disaggregation_mode = "null"

sglang/srt/managers/schedule_policy.py

@@ -380,8 +380,9 @@ class PrefillAdder:
         self.log_input_tokens += extend_input_len
 
     def add_chunked_req(self, req: Req):
-        truncated = req.extend_input_len > self.rem_chunk_tokens
-        req.extend_input_len = min(req.extend_input_len, self.rem_chunk_tokens)
+        _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens))
+        truncated = req.extend_input_len > _rem_tokens
+        req.extend_input_len = min(req.extend_input_len, _rem_tokens)
         req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
         self.can_run_list.append(req)
         self._update_prefill_budget(

sglang/srt/managers/scheduler.py

@@ -141,7 +141,7 @@ from sglang.srt.mem_cache.lora_radix_cache import LoRARadixCache
 from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
 from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors
-from sglang.srt.reasoning_parser import ReasoningParser
+from sglang.srt.parser.reasoning_parser import ReasoningParser
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter
@@ -500,6 +500,7 @@ class Scheduler(
         # Init metrics stats
         self.init_metrics(tp_rank, pp_rank, dp_rank)
         self.init_kv_events(server_args.kv_events_config)
+        self.init_dp_balance(dp_balance_meta)
 
         # Init disaggregation
         self.disaggregation_mode = DisaggregationMode(
@@ -545,15 +546,6 @@ class Scheduler(
             ]
         )
 
-        self.balance_meta = dp_balance_meta
-        if (
-            server_args.enable_dp_attention
-            and server_args.load_balance_method == "minimum_tokens"
-        ):
-            assert dp_balance_meta is not None
-
-        self.recv_dp_balance_id_this_term = []
-
     def init_tokenizer(self):
         server_args = self.server_args
         self.is_generation = self.model_config.is_generation
@@ -1126,11 +1118,7 @@ class Scheduler(
         self,
         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)
+        self.maybe_update_dp_balance_data(recv_req)
 
         # Create a new request
         if (
@@ -1568,11 +1556,7 @@ class Scheduler(
 
         # Handle DP attention
         if need_dp_attn_preparation:
-            if (
-                self.server_args.load_balance_method == "minimum_tokens"
-                and self.forward_ct % 40 == 0
-            ):
-                self.handle_dp_balance_data(ret)
+            self.maybe_handle_dp_balance_data()
             ret = self.prepare_mlp_sync_batch(ret)
 
         return ret
@@ -1897,86 +1881,6 @@ class Scheduler(
             disable_overlap_schedule=self.server_args.disable_overlap_schedule,
         )
 
-    def handle_dp_balance_data(self, local_batch: ScheduleBatch):
-        def gather_dp_balance_info(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(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)
-
-        holding_tokens = self.get_load()
-
-        new_recv_dp_balance_id_list, holding_token_list = gather_dp_balance_info(
-            holding_tokens
-        )
-
-        self.recv_dp_balance_id_this_term.clear()
-        if self.tp_rank == 0:  # only first worker write info
-            write_shared_dp_balance_info(
-                new_recv_dp_balance_id_list, holding_token_list
-            )
-
     @staticmethod
     def prepare_mlp_sync_batch_raw(
         local_batch: ScheduleBatch,
@@ -2403,6 +2307,9 @@ class Scheduler(
             # This only works for requests that have not started anything.
             # We still need to send something back to TokenizerManager to clean up the state.
             req = self.waiting_queue.pop(i)
+            if self.enable_hicache_storage:
+                # to release prefetch events associated with the request
+                self.tree_cache.release_aborted_request(req.rid)
             self.send_to_tokenizer.send_pyobj(AbortReq(req.rid))
             # For disaggregation decode mode, the request in the waiting queue has KV cache allocated.
             if self.disaggregation_mode == DisaggregationMode.DECODE:

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
 
@@ -220,7 +241,7 @@ class SchedulerMetricsMixin:
             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 +257,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()
+
+            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/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__)
 

sglang/srt/managers/tokenizer_manager.py

@@ -329,6 +329,7 @@ class TokenizerManager:
         # Metrics
         if self.enable_metrics:
             self.metrics_collector = TokenizerMetricsCollector(
+                server_args=server_args,
                 labels={
                     "model_name": self.server_args.served_model_name,
                     # TODO: Add lora name/path in the future,