Carve out SchedulerWeightUpdaterManager for weight-update state (#25615)

python/sglang/srt/managers/scheduler.py

@@ -173,6 +173,9 @@ from sglang.srt.managers.scheduler_components.profiler_manager import (
 from sglang.srt.managers.scheduler_components.request_receiver import (
     SchedulerRequestReceiver,
 )
+from sglang.srt.managers.scheduler_components.weight_updater import (
+    SchedulerWeightUpdaterManager,
+)
 from sglang.srt.managers.scheduler_input_blocker import SchedulerInputBlocker
 from sglang.srt.managers.scheduler_output_processor_mixin import (
     SchedulerOutputProcessorMixin,
@@ -545,6 +548,15 @@ class Scheduler(
         # Init prefill kv split size when deterministic inference is enabled with various attention backends
         self.init_deterministic_inference_config()
 
+        self.weight_updater = SchedulerWeightUpdaterManager(
+            tp_worker=self.tp_worker,
+            draft_worker=self.draft_worker,
+            tp_cpu_group=self.tp_cpu_group,
+            memory_saver_adapter=self.memory_saver_adapter,
+            flush_cache=self.flush_cache,
+            is_fully_idle=self.is_fully_idle,
+        )
+
         # Init request dispatcher
         self.init_request_dispatcher()
 
@@ -594,7 +606,7 @@ class Scheduler(
             req_to_token_pool=self.req_to_token_pool,
             token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
             tree_cache=self.tree_cache,
-            offload_tags=self.offload_tags,
+            offload_tags=self.weight_updater.offload_tags,
             ps=self.ps,
             server_args=self.server_args,
             model_config=self.model_config,
@@ -1039,7 +1051,6 @@ class Scheduler(
         self.memory_saver_adapter = TorchMemorySaverAdapter.create(
             enable=self.server_args.enable_memory_saver
         )
-        self.offload_tags = set()
 
         # Init recv skipper and input blocker
         self.recv_skipper = SchedulerRecvSkipper.maybe_create(self.server_args)
@@ -1299,9 +1310,22 @@ class Scheduler(
                 (AbortReq, self.abort_request),
                 (OpenSessionReqInput, self.open_session),
                 (CloseSessionReqInput, self.close_session),
-                (UpdateWeightFromDiskReqInput, self.update_weights_from_disk),
+                (
-                (InitWeightsUpdateGroupReqInput, self.init_weights_update_group),
+                    UpdateWeightFromDiskReqInput,
-                (DestroyWeightsUpdateGroupReqInput, self.destroy_weights_update_group),
+                    lambda req: self.update_weights_from_disk(self.weight_updater, req),
+                ),
+                (
+                    InitWeightsUpdateGroupReqInput,
+                    lambda req: self.init_weights_update_group(
+                        self.weight_updater, req
+                    ),
+                ),
+                (
+                    DestroyWeightsUpdateGroupReqInput,
+                    lambda req: self.destroy_weights_update_group(
+                        self.weight_updater, req
+                    ),
+                ),
                 (
                     InitWeightsSendGroupForRemoteInstanceReqInput,
                     self.init_weights_send_group_for_remote_instance,
@@ -1312,14 +1336,38 @@ class Scheduler(
                 ),
                 (
                     UpdateWeightsFromDistributedReqInput,
-                    self.update_weights_from_distributed,
+                    lambda req: self.update_weights_from_distributed(
+                        self.weight_updater, req
+                    ),
+                ),
+                (
+                    UpdateWeightsFromTensorReqInput,
+                    lambda req: self.update_weights_from_tensor(
+                        self.weight_updater, req
+                    ),
+                ),
+                (
+                    UpdateWeightsFromIPCReqInput,
+                    lambda req: self.update_weights_from_ipc(self.weight_updater, req),
+                ),
+                (
+                    GetWeightsByNameReqInput,
+                    lambda req: self.get_weights_by_name(self.weight_updater, req),
+                ),
+                (
+                    ReleaseMemoryOccupationReqInput,
+                    lambda req: self.release_memory_occupation(
+                        self.weight_updater, req
+                    ),
+                ),
+                (
+                    ResumeMemoryOccupationReqInput,
+                    lambda req: self.resume_memory_occupation(self.weight_updater, req),
+                ),
+                (
+                    CheckWeightsReqInput,
+                    lambda req: self.check_weights(self.weight_updater, req),
                 ),
-                (UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor),
-                (UpdateWeightsFromIPCReqInput, self.update_weights_from_ipc),
-                (GetWeightsByNameReqInput, self.get_weights_by_name),
-                (ReleaseMemoryOccupationReqInput, self.release_memory_occupation),
-                (ResumeMemoryOccupationReqInput, self.resume_memory_occupation),
-                (CheckWeightsReqInput, self.check_weights),
                 (SlowDownReqInput, self.slow_down),
                 (
                     ProfileReq,
@@ -3240,6 +3288,12 @@ class Scheduler(
             server_args=vars(get_global_server_args()),
         )
 
+    def save_remote_model(self, **kwargs):
+        SchedulerUpdateWeightsMixin.save_remote_model(self.weight_updater, kwargs)
+
+    def save_sharded_model(self, **kwargs):
+        SchedulerUpdateWeightsMixin.save_sharded_model(self.weight_updater, kwargs)
+
     def handle_rpc_request(self, recv_req: RpcReqInput):
         # Handle RPC requests
         logger.info(

python/sglang/srt/managers/scheduler_components/weight_updater.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+
+@dataclass(kw_only=True, slots=True)
+class SchedulerWeightUpdaterManager:
+    tp_worker: Any
+    draft_worker: Any
+    tp_cpu_group: Any
+    memory_saver_adapter: Any
+    flush_cache: Callable[..., bool]
+    is_fully_idle: Callable[..., bool]
+    offload_tags: set = field(default_factory=set)
+    stashed_model_static_state: Any = None

python/sglang/srt/managers/scheduler_update_weights_mixin.py

@@ -36,21 +36,27 @@ from sglang.srt.managers.io_struct import (
 )
 
 if TYPE_CHECKING:
-    from sglang.srt.managers.scheduler import Scheduler
+    from sglang.srt.managers.scheduler_components.weight_updater import (
+        SchedulerWeightUpdaterManager,
+    )
 
 logger = logging.getLogger(__name__)
 
 
 class SchedulerUpdateWeightsMixin:
-    def flush_cache_after_weight_update(self: Scheduler, recv_req) -> None:
+    @staticmethod
+    def flush_cache_after_weight_update(
+        self: "SchedulerWeightUpdaterManager", recv_req
+    ) -> None:
         if recv_req.flush_cache:
             flush_cache_success = self.flush_cache(
                 empty_cache=recv_req.torch_empty_cache
             )
             assert flush_cache_success, "Cache flush failed after updating weights"
 
+    @staticmethod
     def update_weights_from_disk(
-        self: Scheduler, recv_req: UpdateWeightFromDiskReqInput
+        self: "SchedulerWeightUpdaterManager", recv_req: UpdateWeightFromDiskReqInput
     ):
         """In-place update of the weights from disk."""
         success, message = self.tp_worker.update_weights_from_disk(recv_req)
@@ -58,39 +64,44 @@ class SchedulerUpdateWeightsMixin:
         if success and self.draft_worker is not None:
             success, message = self.draft_worker.update_weights_from_disk(recv_req)
         if tp_success:
-            self.flush_cache_after_weight_update(recv_req)
+            SchedulerUpdateWeightsMixin.flush_cache_after_weight_update(self, recv_req)
         if not success:
             logger.error(message)
         return UpdateWeightFromDiskReqOutput(success, message, 0)
 
+    @staticmethod
     def init_weights_update_group(
-        self: Scheduler, recv_req: InitWeightsUpdateGroupReqInput
+        self: "SchedulerWeightUpdaterManager", recv_req: InitWeightsUpdateGroupReqInput
     ):
         """Initialize the online model parameter update group."""
         success, message = self.tp_worker.init_weights_update_group(recv_req)
         return InitWeightsUpdateGroupReqOutput(success, message)
 
+    @staticmethod
     def destroy_weights_update_group(
-        self: Scheduler, recv_req: DestroyWeightsUpdateGroupReqInput
+        self: "SchedulerWeightUpdaterManager",
+        recv_req: DestroyWeightsUpdateGroupReqInput,
     ):
         """Destroy the online model parameter update group."""
         success, message = self.tp_worker.destroy_weights_update_group(recv_req)
         return DestroyWeightsUpdateGroupReqOutput(success, message)
 
+    @staticmethod
     def update_weights_from_distributed(
-        self,
+        self: "SchedulerWeightUpdaterManager",
         recv_req: UpdateWeightsFromDistributedReqInput,
     ) -> Tuple[bool, str]:
         """Update the online model parameter."""
         success, message = self.tp_worker.update_weights_from_distributed(recv_req)
         if success:
-            self.flush_cache_after_weight_update(recv_req)
+            SchedulerUpdateWeightsMixin.flush_cache_after_weight_update(self, recv_req)
         else:
             logger.error(message)
         return UpdateWeightsFromDistributedReqOutput(success, message)
 
+    @staticmethod
     def update_weights_from_tensor(
-        self: Scheduler, recv_req: UpdateWeightsFromTensorReqInput
+        self: "SchedulerWeightUpdaterManager", recv_req: UpdateWeightsFromTensorReqInput
     ):
         """Update the online model parameter from tensors."""
         if recv_req.disable_draft_model:
@@ -99,14 +110,15 @@ class SchedulerUpdateWeightsMixin:
             worker = self.draft_worker or self.tp_worker
         success, message = worker.update_weights_from_tensor(recv_req)
         if success:
-            self.flush_cache_after_weight_update(recv_req)
+            SchedulerUpdateWeightsMixin.flush_cache_after_weight_update(self, recv_req)
         else:
             logger.error(message)
         torch.distributed.barrier(group=self.tp_cpu_group)
         return UpdateWeightsFromTensorReqOutput(success, message)
 
+    @staticmethod
     def update_weights_from_ipc(
-        self: Scheduler, recv_req: UpdateWeightsFromIPCReqInput
+        self: "SchedulerWeightUpdaterManager", 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)
@@ -114,18 +126,22 @@ class SchedulerUpdateWeightsMixin:
         if success and self.draft_worker is not None:
             success, message = self.draft_worker.update_weights_from_ipc(recv_req)
         if tp_success:
-            self.flush_cache_after_weight_update(recv_req)
+            SchedulerUpdateWeightsMixin.flush_cache_after_weight_update(self, recv_req)
         if not success:
             logger.error(message)
         torch.distributed.barrier(group=self.tp_cpu_group)
         return UpdateWeightsFromIPCReqOutput(success, message)
 
-    def get_weights_by_name(self: Scheduler, recv_req: GetWeightsByNameReqInput):
+    @staticmethod
+    def get_weights_by_name(
+        self: "SchedulerWeightUpdaterManager", recv_req: GetWeightsByNameReqInput
+    ):
         parameter = self.tp_worker.get_weights_by_name(recv_req)
         return GetWeightsByNameReqOutput(parameter)
 
+    @staticmethod
     def release_memory_occupation(
-        self: Scheduler, recv_req: ReleaseMemoryOccupationReqInput
+        self: "SchedulerWeightUpdaterManager", recv_req: ReleaseMemoryOccupationReqInput
     ):
         assert (
             self.is_fully_idle()
@@ -157,8 +173,9 @@ class SchedulerUpdateWeightsMixin:
 
         return ReleaseMemoryOccupationReqOutput()
 
+    @staticmethod
     def resume_memory_occupation(
-        self: Scheduler, recv_req: ResumeMemoryOccupationReqInput
+        self: "SchedulerWeightUpdaterManager", recv_req: ResumeMemoryOccupationReqInput
     ):
         tags = recv_req.tags
 
@@ -185,7 +202,10 @@ class SchedulerUpdateWeightsMixin:
 
         return ResumeMemoryOccupationReqOutput()
 
-    def check_weights(self: Scheduler, recv_req: CheckWeightsReqInput):
+    @staticmethod
+    def check_weights(
+        self: "SchedulerWeightUpdaterManager", recv_req: CheckWeightsReqInput
+    ):
         try:
             payload = self.tp_worker.model_runner.check_weights(action=recv_req.action)
             return CheckWeightsReqOutput(
@@ -196,7 +216,8 @@ class SchedulerUpdateWeightsMixin:
             traceback.print_exc()
             return CheckWeightsReqOutput(success=False, message=f"{e}")
 
-    def save_remote_model(self: Scheduler, params):
+    @staticmethod
+    def save_remote_model(self: "SchedulerWeightUpdaterManager", params):
         url = params["url"]
 
         self.tp_worker.model_runner.save_remote_model(url)
@@ -208,7 +229,8 @@ class SchedulerUpdateWeightsMixin:
             ), "draft_url must be provided when draft model is enabled"
             self.draft_worker.model_runner.save_remote_model(draft_url)
 
-    def save_sharded_model(self: Scheduler, params):
+    @staticmethod
+    def save_sharded_model(self: "SchedulerWeightUpdaterManager", params):
         self.tp_worker.model_runner.save_sharded_model(
             path=params["path"],
             pattern=params["pattern"],