sglang 0.5.3rc2__py3-none-any.whl → 0.5.4.post1__py3-none-any.whl

sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
+from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Scheduler
+
+
+class SchedulerRuntimeCheckerMixin:
+
+    def _check_hybrid_memory(self: Scheduler):
+        (
+            full_num_used,
+            swa_num_used,
+            _,
+            _,
+            full_available_size,
+            full_evictable_size,
+            swa_available_size,
+            swa_evictable_size,
+        ) = self._get_swa_token_info()
+        memory_leak = full_num_used != 0 or swa_num_used != 0
+        token_msg = (
+            f"{self.full_tokens_per_layer=}, {full_available_size=}, {full_evictable_size=}, {self.tree_cache.full_protected_size()=}\n"
+            f"{self.swa_tokens_per_layer=}, {swa_available_size=}, {swa_evictable_size=}, {self.tree_cache.swa_protected_size()=}\n"
+        )
+        return memory_leak, token_msg
+
+    def _check_mamba_memory(self: Scheduler):
+        (
+            full_num_used,
+            mamba_num_used,
+            _,
+            _,
+            full_available_size,
+            full_evictable_size,
+            mamba_available_size,
+            mamba_evictable_size,
+        ) = self._get_mamba_token_info()
+        memory_leak = (
+            full_num_used != self.tree_cache.full_protected_size()
+            or mamba_num_used != self.tree_cache.mamba_protected_size()
+        )
+        token_msg = (
+            f"{full_available_size=}, {full_evictable_size=}, {self.token_to_kv_pool_allocator.size=}, {self.tree_cache.full_protected_size()=}\n"
+            f"{mamba_available_size=}, {mamba_evictable_size=}, {self.req_to_token_pool.mamba_pool.size=}, {self.tree_cache.mamba_protected_size()=}\n"
+        )
+        return memory_leak, token_msg
+
+    def _check_radix_cache_memory(self: Scheduler):
+        _, _, available_size, evictable_size = self._get_token_info()
+        protected_size = self.tree_cache.protected_size()
+        memory_leak = (available_size + evictable_size) != (
+            # self.max_total_num_tokens
+            # if not self.enable_hierarchical_cache
+            # else self.max_total_num_tokens - protected_size
+            self.max_total_num_tokens
+            - protected_size
+        )
+        token_msg = f"{self.max_total_num_tokens=}, {available_size=}, {evictable_size=}, {protected_size=}\n"
+        return memory_leak, token_msg
+
+    def _check_runtime_mem_leak(self: Scheduler):
+        current_batch: ScheduleBatch = self.last_batch
+
+        if current_batch is None:
+            return
+
+        _, _, available_size, evictable_size = self._get_token_info()
+        protected_size = self.tree_cache.protected_size()
+
+        extend_size = 0
+        for i, req in enumerate(current_batch.reqs):
+            seq_len = len(req.origin_input_ids) + len(req.output_ids)
+            fill_len = len(req.fill_ids) if req.fill_ids is not None else 0
+            prefix_len = (
+                len(req.prefix_indices) if req.prefix_indices is not None else 0
+            )
+
+            if current_batch.forward_mode.is_decode():
+                if req.finished():
+                    unreleased_len = 1
+                else:
+                    unreleased_len = seq_len - prefix_len
+            else:
+                unreleased_len = fill_len - prefix_len
+
+            extend_size += unreleased_len
+
+        if (
+            current_batch.forward_mode.is_extend()
+            and self.running_batch is not None
+            and not self.running_batch.is_empty()
+            and self.running_batch.forward_mode.is_decode()
+        ):
+            for i, req in enumerate(self.running_batch.reqs):
+                seq_len = len(req.origin_input_ids) + len(req.output_ids)
+                prefix_len = (
+                    len(req.prefix_indices) if req.prefix_indices is not None else 0
+                )
+
+                if req.finished():
+                    unreleased_len = 0
+                else:
+                    unreleased_len = seq_len - prefix_len - 1
+
+                extend_size += unreleased_len
+
+        total_tokens = available_size + evictable_size + protected_size + extend_size
+
+        assert (
+            total_tokens == self.max_total_num_tokens
+        ), f"Mem Leak Detected! {total_tokens=} vs {self.max_total_num_tokens=}"
+
+    def _check_req_pool(self: Scheduler):
+        if self.disaggregation_mode == DisaggregationMode.DECODE:
+            req_total_size = (
+                self.req_to_token_pool.size + self.req_to_token_pool.pre_alloc_size
+            )
+        else:
+            req_total_size = self.req_to_token_pool.size
+
+        if len(self.req_to_token_pool.free_slots) != req_total_size:
+            msg = (
+                "req_to_token_pool memory leak detected!"
+                f"available_size={len(self.req_to_token_pool.free_slots)}, "
+                f"total_size={self.req_to_token_pool.size}\n"
+            )
+            raise ValueError(msg)
+
+    def check_memory(self: Scheduler):
+        if self.is_hybrid:
+            memory_leak, token_msg = self._check_hybrid_memory()
+        elif self.is_hybrid_gdn and isinstance(self.tree_cache, MambaRadixCache):
+            memory_leak, token_msg = self._check_mamba_memory()
+        else:
+            memory_leak, token_msg = self._check_radix_cache_memory()
+
+        if memory_leak:
+            msg = "token_to_kv_pool_allocator memory leak detected! " f"{token_msg}"
+            raise ValueError(msg)
+
+        self._check_req_pool()
+
+        if (
+            self.enable_metrics
+            and self.current_scheduler_metrics_enabled()
+            and time.perf_counter() > self.metrics_collector.last_log_time + 30
+        ):
+            # During idle time, also collect metrics every 30 seconds.
+            if self.is_hybrid:
+                (
+                    full_num_used,
+                    swa_num_used,
+                    full_token_usage,
+                    swa_token_usage,
+                    _,
+                    _,
+                    _,
+                    _,
+                ) = self._get_swa_token_info()
+                num_used = max(full_num_used, swa_num_used)
+                token_usage = max(full_token_usage, swa_token_usage)
+            elif self.is_hybrid_gdn:
+                (
+                    num_used,
+                    _,
+                    token_usage,
+                    _,
+                    _,
+                    _,
+                    _,
+                    _,
+                ) = self._get_mamba_token_info()
+            else:
+                num_used, token_usage, _, _ = self._get_token_info()
+            num_running_reqs = len(self.running_batch.reqs)
+            self.stats.num_running_reqs = num_running_reqs
+            self.stats.num_used_tokens = num_used
+            self.stats.token_usage = round(token_usage, 2)
+            self.stats.gen_throughput = 0
+            self.stats.num_queue_reqs = len(self.waiting_queue)
+            self.stats.num_grammar_queue_reqs = len(self.grammar_queue)
+            if self.disaggregation_mode == DisaggregationMode.PREFILL:
+                self.stats.num_prefill_prealloc_queue_reqs = len(
+                    self.disagg_prefill_bootstrap_queue.queue
+                )
+                self.stats.num_prefill_inflight_queue_reqs = len(
+                    self.disagg_prefill_inflight_queue
+                )
+            if self.disaggregation_mode == DisaggregationMode.DECODE:
+                self.stats.num_decode_prealloc_queue_reqs = len(
+                    self.disagg_decode_prealloc_queue.queue
+                )
+                self.stats.num_decode_transfer_queue_reqs = len(
+                    self.disagg_decode_transfer_queue.queue
+                )
+            self.metrics_collector.log_stats(self.stats)
+        self._publish_kv_events()
+
+    def check_tree_cache(self: Scheduler):
+        if (self.is_hybrid and isinstance(self.tree_cache, SWARadixCache)) or (
+            self.is_hybrid_gdn and isinstance(self.tree_cache, MambaRadixCache)
+        ):
+            self.tree_cache.sanity_check()
+
+    def self_check_during_idle(self: Scheduler):
+        self.check_memory()
+        self.check_tree_cache()
+        self.new_token_ratio = self.init_new_token_ratio
+        self.maybe_sleep_on_idle()

sglang/srt/managers/scheduler_update_weights_mixin.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import logging
-from typing import Tuple
+from typing import TYPE_CHECKING, Tuple
 
 import torch
 
@@ -19,10 +21,15 @@ from sglang.srt.managers.io_struct import (
     UpdateWeightFromDiskReqOutput,
     UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromDistributedReqOutput,
+    UpdateWeightsFromIPCReqInput,
+    UpdateWeightsFromIPCReqOutput,
     UpdateWeightsFromTensorReqInput,
     UpdateWeightsFromTensorReqOutput,
 )
 
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Scheduler
+
 logger = logging.getLogger(__name__)
 
 
@@ -75,11 +82,25 @@ class SchedulerUpdateWeightsMixin:
         torch.distributed.barrier(group=self.tp_cpu_group)
         return UpdateWeightsFromTensorReqOutput(success, message)
 
+    def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput):
+        """Update the online model parameter from IPC for checkpoint-engine integration."""
+        success, message = self.tp_worker.update_weights_from_ipc(recv_req)
+        if success:
+            if recv_req.flush_cache:
+                flush_cache_success = self.flush_cache()
+                assert flush_cache_success, "Cache flush failed after updating weights"
+        else:
+            logger.error(message)
+        torch.distributed.barrier(group=self.tp_cpu_group)
+        return UpdateWeightsFromIPCReqOutput(success, message)
+
     def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
         parameter = self.tp_worker.get_weights_by_name(recv_req)
         return GetWeightsByNameReqOutput(parameter)
 
-    def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
+    def release_memory_occupation(
+        self: Scheduler, recv_req: ReleaseMemoryOccupationReqInput
+    ):
         tags = recv_req.tags
 
         if tags is None or len(tags) == 0:
@@ -94,14 +115,16 @@ class SchedulerUpdateWeightsMixin:
 
         if GPU_MEMORY_TYPE_WEIGHTS in tags:
             self.stashed_model_static_state = _export_static_state(
-                self.tp_worker.worker.model_runner.model
+                self.tp_worker.model_runner.model
             )
             torch.distributed.barrier(self.tp_cpu_group)
             self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_WEIGHTS)
 
         return ReleaseMemoryOccupationReqOutput()
 
-    def resume_memory_occupation(self, recv_req: ResumeMemoryOccupationReqInput):
+    def resume_memory_occupation(
+        self: Scheduler, recv_req: ResumeMemoryOccupationReqInput
+    ):
         tags = recv_req.tags
 
         if tags is None or len(tags) == 0:
@@ -114,7 +137,7 @@ class SchedulerUpdateWeightsMixin:
             self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_WEIGHTS)
             torch.distributed.barrier(self.tp_cpu_group)
             _import_static_state(
-                self.tp_worker.worker.model_runner.model,
+                self.tp_worker.model_runner.model,
                 self.stashed_model_static_state,
             )
             del self.stashed_model_static_state
@@ -124,24 +147,20 @@ class SchedulerUpdateWeightsMixin:
 
         return ResumeMemoryOccupationReqOutput()
 
-    def save_remote_model(self, params):
+    def save_remote_model(self: Scheduler, params):
         url = params["url"]
 
-        worker = self.tp_worker.worker
-        worker.model_runner.save_remote_model(url)
+        self.tp_worker.model_runner.save_remote_model(url)
 
         if self.draft_worker is not None:
             draft_url = params.get("draft_url", None)
             assert (
                 draft_url is not None
             ), "draft_url must be provided when draft model is enabled"
-            draft_worker = self.draft_worker.worker
-            draft_worker.model_runner.save_remote_model(draft_url)
-
-    def save_sharded_model(self, params):
-        worker = self.tp_worker.worker
+            self.draft_worker.model_runner.save_remote_model(draft_url)
 
-        worker.model_runner.save_sharded_model(
+    def save_sharded_model(self: Scheduler, params):
+        self.tp_worker.model_runner.save_sharded_model(
             path=params["path"],
             pattern=params["pattern"],
             max_size=params["max_size"],

sglang/srt/managers/tokenizer_communicator_mixin.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import asyncio
 import copy
 import logging
-import os
 import time
 import uuid
 from collections import deque
@@ -46,7 +45,6 @@ from sglang.srt.managers.io_struct import (
     LoadLoRAAdapterReqInput,
     LoadLoRAAdapterReqOutput,
     LoRAUpdateOutput,
-    MultiTokenizerWrapper,
     OpenSessionReqInput,
     ProfileReq,
     ProfileReqOutput,
@@ -65,6 +63,8 @@ from sglang.srt.managers.io_struct import (
     UnloadLoRAAdapterReqOutput,
     UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromDistributedReqOutput,
+    UpdateWeightsFromIPCReqInput,
+    UpdateWeightsFromIPCReqOutput,
     UpdateWeightsFromTensorReqInput,
     UpdateWeightsFromTensorReqOutput,
 )
@@ -83,8 +83,6 @@ logger = logging.getLogger(__name__)
 class _Communicator(Generic[T]):
     """Note: The communicator now only run up to 1 in-flight request at any time."""
 
-    enable_multi_tokenizer = False
-
     def __init__(self, sender: zmq.Socket, fan_out: int, mode="queueing"):
         self._sender = sender
         self._fan_out = fan_out
@@ -104,8 +102,6 @@ class _Communicator(Generic[T]):
             assert self._result_values is None
 
         if obj:
-            if _Communicator.enable_multi_tokenizer:
-                obj = MultiTokenizerWrapper(worker_id=os.getpid(), obj=obj)
             self._sender.send_pyobj(obj)
 
         self._result_event = asyncio.Event()
@@ -126,8 +122,6 @@ class _Communicator(Generic[T]):
             self._result_event = asyncio.Event()
 
             if obj:
-                if _Communicator.enable_multi_tokenizer:
-                    obj = MultiTokenizerWrapper(worker_id=os.getpid(), obj=obj)
                 self._sender.send_pyobj(obj)
 
         await self._result_event.wait()
@@ -146,6 +140,13 @@ class _Communicator(Generic[T]):
         if len(self._result_values) == self._fan_out:
             self._result_event.set()
 
+    @staticmethod
+    def merge_results(results):
+        all_success = all([r.success for r in results])
+        all_message = [r.message for r in results]
+        all_message = " | ".join(all_message)
+        return all_success, all_message
+
 
 class TokenizerCommunicatorMixin:
     """Mixin class for TokenizerManager to handle communication with the scheduler."""
@@ -170,6 +171,9 @@ class TokenizerCommunicatorMixin:
         self.update_weights_from_tensor_communicator = _Communicator(
             self.send_to_scheduler, server_args.dp_size
         )
+        self.update_weights_from_ipc_communicator = _Communicator(
+            self.send_to_scheduler, server_args.dp_size
+        )
         self.get_weights_by_name_communicator = _Communicator(
             self.send_to_scheduler, server_args.dp_size
         )
@@ -236,6 +240,10 @@ class TokenizerCommunicatorMixin:
                     UpdateWeightsFromTensorReqOutput,
                     self.update_weights_from_tensor_communicator.handle_recv,
                 ),
+                (
+                    UpdateWeightsFromIPCReqOutput,
+                    self.update_weights_from_ipc_communicator.handle_recv,
+                ),
                 (
                     GetWeightsByNameReqOutput,
                     self.get_weights_by_name_communicator.handle_recv,
@@ -306,6 +314,7 @@ class TokenizerCommunicatorMixin:
         with_stack: Optional[bool] = None,
         record_shapes: Optional[bool] = None,
         profile_by_stage: bool = False,
+        merge_profiles: bool = False,
     ):
         self.auto_create_handle_loop()
         env_with_stack: bool = get_bool_env_var("SGLANG_PROFILE_WITH_STACK", "true")
@@ -320,6 +329,7 @@ class TokenizerCommunicatorMixin:
             record_shapes=record_shapes,
             profile_by_stage=profile_by_stage,
             profile_id=str(time.time()),
+            merge_profiles=merge_profiles,
         )
         return await self._execute_profile(req)
 
@@ -356,10 +366,11 @@ class TokenizerCommunicatorMixin:
     ) -> Tuple[bool, str]:
         self.auto_create_handle_loop()
         assert (
-            self.server_args.dp_size == 1
-        ), "dp_size must be 1 for init parameter update group"
-        result = (await self.init_weights_update_group_communicator(obj))[0]
-        return result.success, result.message
+            self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+        ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed"
+
+        results = await self.init_weights_update_group_communicator(obj)
+        return _Communicator.merge_results(results)
 
     async def destroy_weights_update_group(
         self,
@@ -368,10 +379,11 @@ class TokenizerCommunicatorMixin:
     ) -> Tuple[bool, str]:
         self.auto_create_handle_loop()
         assert (
-            self.server_args.dp_size == 1
-        ), "dp_size must be 1 for destroy parameter update group"
-        result = (await self.destroy_weights_update_group_communicator(obj))[0]
-        return result.success, result.message
+            self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+        ), "dp_size must be 1 or dp attention must be enabled for destroy parameter update group"
+
+        results = await self.destroy_weights_update_group_communicator(obj)
+        return _Communicator.merge_results(results)
 
     async def update_weights_from_distributed(
         self: TokenizerManager,
@@ -389,8 +401,8 @@ class TokenizerCommunicatorMixin:
         # This means that weight sync
         # cannot run while requests are in progress.
         async with self.model_update_lock.writer_lock:
-            result = (await self.update_weights_from_distributed_communicator(obj))[0]
-            return result.success, result.message
+            results = await self.update_weights_from_distributed_communicator(obj)
+            return _Communicator.merge_results(results)
 
     async def init_weights_send_group_for_remote_instance(
         self,
@@ -439,6 +451,28 @@ class TokenizerCommunicatorMixin:
             result = (await self.update_weights_from_tensor_communicator(obj))[0]
             return result.success, result.message
 
+    async def update_weights_from_ipc(
+        self,
+        obj: UpdateWeightsFromIPCReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        """Update weights via IPC for checkpoint-engine integration."""
+        self.auto_create_handle_loop()
+        try:
+            # For now, we only support single data parallel instance
+            assert (
+                self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+            ), "dp_size must be 1 or dp attention must be enabled for update weights from IPC"
+            logger.info("Starting IPC weight update")
+            # This means that weight sync cannot run while requests are in progress.
+            async with self.model_update_lock.writer_lock:
+                result = (await self.update_weights_from_ipc_communicator(obj))[0]
+                return result.success, result.message
+        except Exception as e:
+            error_msg = f"IPC weight update failed: {str(e)}"
+            logger.error(error_msg)
+            return False, error_msg
+
     async def load_lora_adapter(
         self: TokenizerManager,
         obj: LoadLoRAAdapterReqInput,
@@ -606,8 +640,6 @@ class TokenizerCommunicatorMixin:
         elif obj.session_id in self.session_futures:
             return None
 
-        if self.server_args.tokenizer_worker_num > 1:
-            obj = MultiTokenizerWrapper(self.worker_id, obj)
         self.send_to_scheduler.send_pyobj(obj)
 
         self.session_futures[obj.session_id] = asyncio.Future()
@@ -627,43 +659,27 @@ class TokenizerCommunicatorMixin:
         if self.log_requests:
             if self.log_requests_level == 0:
                 max_length = 1 << 30
-                skip_names = set(
-                    [
-                        "text",
-                        "input_ids",
-                        "input_embeds",
-                        "image_data",
-                        "audio_data",
-                        "lora_path",
-                        "sampling_params",
-                    ]
-                )
-                out_skip_names = set(
-                    [
-                        "text",
-                        "output_ids",
-                        "embedding",
-                    ]
-                )
+                skip_names = {
+                    "text",
+                    "input_ids",
+                    "input_embeds",
+                    "image_data",
+                    "audio_data",
+                    "lora_path",
+                    "sampling_params",
+                }
+                out_skip_names = {"text", "output_ids", "embedding"}
             elif self.log_requests_level == 1:
                 max_length = 1 << 30
-                skip_names = set(
-                    [
-                        "text",
-                        "input_ids",
-                        "input_embeds",
-                        "image_data",
-                        "audio_data",
-                        "lora_path",
-                    ]
-                )
-                out_skip_names = set(
-                    [
-                        "text",
-                        "output_ids",
-                        "embedding",
-                    ]
-                )
+                skip_names = {
+                    "text",
+                    "input_ids",
+                    "input_embeds",
+                    "image_data",
+                    "audio_data",
+                    "lora_path",
+                }
+                out_skip_names = {"text", "output_ids", "embedding"}
             elif self.log_requests_level == 2:
                 max_length = 2048
             elif self.log_requests_level == 3: