diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index d03c541bbd..0f9b80d6e2 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -557,6 +557,10 @@ jobs: "num_gpus": 0, "test_file": "test_megatron_argument_validation.py" }, + { + "num_gpus": 0, + "test_file": "test_mooncake_weight_transport.py" + }, { "num_gpus": 0, "test_file": "utils/test_megatron_server_arguments.py" @@ -1121,4 +1125,4 @@ jobs: else python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" fi - ' \ No newline at end of file + ' diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index b6f8a4ffc3..11781a965d 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -64,6 +64,7 @@ 'cpu': True, 'tests': [ {'test_file': 'test_megatron_argument_validation.py', 'num_gpus': 0}, + {'test_file': 'test_mooncake_weight_transport.py', 'num_gpus': 0}, {'test_file': 'utils/test_megatron_server_arguments.py', 'num_gpus': 0}, {'test_file': 'test_dp_schedule.py', 'num_gpus': 0}, {'test_file': 'test_cp_utils.py', 'num_gpus': 0}, diff --git a/docker/patch/latest/sglang.patch b/docker/patch/latest/sglang.patch index 191c20ad4a..949ecf7f53 100644 --- a/docker/patch/latest/sglang.patch +++ b/docker/patch/latest/sglang.patch @@ -682,7 +682,18 @@ diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/en index d7368383d89..2c881d95bd5 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py -@@ -127,6 +127,7 @@ from sglang.srt.managers.io_struct import ( +@@ -118,24 +118,28 @@ from sglang.srt.managers.io_struct import ( ++ DestroyMooncakeWeightReceiverReqInput, + DestroyWeightsUpdateGroupReqInput, + DumperControlReqInput, + EmbeddingReqInput, + GenerateReqInput, + GetWeightsByNameReqInput, ++ InitMooncakeWeightReceiverReqInput, + InitWeightsSendGroupForRemoteInstanceReqInput, + InitWeightsUpdateGroupReqInput, + LoadLoRAAdapterFromTensorsReqInput, + LoadLoRAAdapterReqInput, OpenSessionReqInput, ParseFunctionCallReq, PauseGenerationReqInput, @@ -690,6 +701,16 @@ index d7368383d89..2c881d95bd5 100644 ProfileReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, + SendWeightsToRemoteInstanceReqInput, + SeparateReasoningReqInput, + SetInternalStateReq, + SlowDownReqInput, + UnloadLoRAAdapterReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, ++ UpdateWeightsFromMooncakeReqInput, + UpdateWeightsFromTensorReqInput, @@ -618,10 +619,8 @@ async def model_info(): @app.get("/weight_version") async def weight_version(): @@ -725,10 +746,63 @@ index d7368383d89..2c881d95bd5 100644 server_args = _global_state.tokenizer_manager.server_args -@@ -1249,6 +1257,20 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re +@@ -1249,6 +1257,73 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) ++@app.post("/init_mooncake_weight_receiver") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def init_mooncake_weight_receiver( ++ obj: InitMooncakeWeightReceiverReqInput, request: Request ++): ++ """Initialize a Mooncake Transfer Engine receiver for trainer-side weight writes.""" ++ success, message, descriptor = ( ++ await _global_state.tokenizer_manager.init_mooncake_weight_receiver( ++ obj, request ++ ) ++ ) ++ content = {"success": success, "message": message} ++ if descriptor: ++ content.update(descriptor) ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ ++@app.post("/update_weights_from_mooncake") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def update_weights_from_mooncake( ++ obj: UpdateWeightsFromMooncakeReqInput, request: Request ++): ++ """Load weights that were already written into the local Mooncake receiver buffer.""" ++ success, message = ( ++ await _global_state.tokenizer_manager.update_weights_from_mooncake( ++ obj, request ++ ) ++ ) ++ content = {"success": success, "message": message} ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ ++@app.post("/destroy_mooncake_weight_receiver") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def destroy_mooncake_weight_receiver( ++ obj: DestroyMooncakeWeightReceiverReqInput, request: Request ++): ++ """Release the Mooncake receiver buffers and Transfer Engine instance.""" ++ success, message = ( ++ await _global_state.tokenizer_manager.destroy_mooncake_weight_receiver( ++ obj, request ++ ) ++ ) ++ content = {"success": success, "message": message} ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ +@app.post("/post_process_weights") +@auth_level(AuthLevel.ADMIN_OPTIONAL) +async def post_process_weights(req: PostProcessWeightsReqInput, request: Request): @@ -1082,10 +1156,64 @@ index 987ec512122..e098565729b 100644 # Whether to call torch.cuda.empty_cache() during flush torch_empty_cache: bool = False -@@ -1673,6 +1711,18 @@ class ResumeMemoryOccupationReqOutput(BaseReq): +@@ -1673,6 +1711,72 @@ class ResumeMemoryOccupationReqOutput(BaseReq): pass ++@dataclass ++class InitMooncakeWeightReceiverReqInput(BaseReq): ++ metadata_server: str = "P2PHANDSHAKE" ++ protocol: str = "tcp" ++ device_name: str = "" ++ buffer_size: int = 0 ++ buffer_count: int = 1 ++ rpc_port_base: Optional[int] = None ++ group_name: str = "slime-mooncake-weight" ++ server_name: Optional[str] = None ++ ++ ++@dataclass ++class InitMooncakeWeightReceiverReqOutput(BaseReq): ++ success: bool ++ message: str ++ target_name: Optional[str] = None ++ server_name: Optional[str] = None ++ rpc_port: Optional[int] = None ++ buffer_device: Optional[str] = None ++ buffers: Optional[List[Dict[str, Any]]] = None ++ ++ ++@dataclass ++class UpdateWeightsFromMooncakeReqInput(BaseReq): ++ weight_version: Optional[str] ++ target_name: str ++ buffer_slot: int ++ remote_addr: int ++ nbytes: int ++ tensors: List[Dict[str, Any]] ++ load_format: Optional[str] = None ++ delta: Optional[Dict[str, Any]] = None ++ flush_cache: bool = True ++ abort_all_requests: bool = False ++ ++ ++@dataclass ++class UpdateWeightsFromMooncakeReqOutput(BaseReq): ++ success: bool ++ message: str ++ ++ ++@dataclass ++class DestroyMooncakeWeightReceiverReqInput(BaseReq): ++ pass ++ ++ ++@dataclass ++class DestroyMooncakeWeightReceiverReqOutput(BaseReq): ++ success: bool ++ message: str ++ ++ +@dataclass +class PostProcessWeightsReqInput(BaseReq): + restore_weights_before_load: bool = False @@ -1162,7 +1290,34 @@ diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/manager index 8e32640fc6a..98966842506 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py -@@ -124,6 +124,7 @@ from sglang.srt.managers.io_struct import ( +@@ -99,49 +99,53 @@ from sglang.srt.managers.io_struct import ( + ConfigureLoggingReq, + ContinueGenerationReqInput, ++ DestroyMooncakeWeightReceiverReqInput, + DestroyWeightsUpdateGroupReqInput, + DetachHiCacheStorageReqInput, + DetachHiCacheStorageReqOutput, + DumperControlReqInput, + DumperControlReqOutput, + ExpertDistributionReq, + ExpertDistributionReqOutput, + ExpertDistributionReqType, + FlushCacheReqInput, + FreezeGCReq, + GetInternalStateReq, + GetInternalStateReqOutput, + GetLoadsReqInput, + GetWeightsByNameReqInput, + HealthCheckOutput, ++ InitMooncakeWeightReceiverReqInput, + InitWeightsSendGroupForRemoteInstanceReqInput, + InitWeightsSendGroupForRemoteInstanceReqOutput, + InitWeightsUpdateGroupReqInput, + ListExternalCorporaReqInput, + ListExternalCorporaReqOutput, + LoadLoRAAdapterFromTensorsReqInput, + LoadLoRAAdapterFromTensorsReqOutput, + LoadLoRAAdapterReqInput, LoadLoRAAdapterReqOutput, OpenSessionReqInput, PauseGenerationReqInput, @@ -1170,11 +1325,42 @@ index 8e32640fc6a..98966842506 100644 ProfileReq, ReleaseMemoryOccupationReqInput, RemoveExternalCorpusReqInput, -@@ -1318,6 +1319,10 @@ class Scheduler( + RemoveExternalCorpusReqOutput, + ResumeMemoryOccupationReqInput, + RpcReqInput, + RpcReqOutput, + SendWeightsToRemoteInstanceReqInput, + SendWeightsToRemoteInstanceReqOutput, + SetInternalStateReq, + SetInternalStateReqOutput, + SlowDownReqInput, + SlowDownReqOutput, + TokenizedEmbeddingReqInput, + TokenizedGenerateReqInput, + UnloadLoRAAdapterReqInput, + UnloadLoRAAdapterReqOutput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, ++ UpdateWeightsFromMooncakeReqInput, + UpdateWeightsFromTensorReqInput, +@@ -1318,6 +1319,22 @@ class Scheduler( UpdateWeightsFromIPCReqInput, self.weight_updater.update_weights_from_ipc, ), + ( ++ InitMooncakeWeightReceiverReqInput, ++ self.weight_updater.init_mooncake_weight_receiver, ++ ), ++ ( ++ UpdateWeightsFromMooncakeReqInput, ++ self.weight_updater.update_weights_from_mooncake, ++ ), ++ ( ++ DestroyMooncakeWeightReceiverReqInput, ++ self.weight_updater.destroy_mooncake_weight_receiver, ++ ), ++ ( + PostProcessWeightsReqInput, + self.weight_updater.post_process_weights, + ), @@ -1334,7 +1520,7 @@ diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b index 77bf823b081..9ab3abe5618 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py -@@ -16,6 +16,7 @@ from sglang.srt.constants import ( +@@ -16,17 +16,26 @@ from sglang.srt.constants import ( GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS, ) @@ -1342,8 +1528,14 @@ index 77bf823b081..9ab3abe5618 100644 from sglang.srt.managers.io_struct import ( CheckWeightsReqInput, CheckWeightsReqOutput, -@@ -25,6 +26,8 @@ from sglang.srt.managers.io_struct import ( ++ DestroyMooncakeWeightReceiverReqInput, ++ DestroyMooncakeWeightReceiverReqOutput, + DestroyWeightsUpdateGroupReqInput, + DestroyWeightsUpdateGroupReqOutput, + GetWeightsByNameReqInput, GetWeightsByNameReqOutput, ++ InitMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqOutput, InitWeightsUpdateGroupReqInput, InitWeightsUpdateGroupReqOutput, + PostProcessWeightsReqInput, @@ -1351,6 +1543,10 @@ index 77bf823b081..9ab3abe5618 100644 ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, ResumeMemoryOccupationReqInput, + ResumeMemoryOccupationReqOutput, ++ UpdateWeightsFromMooncakeReqInput, ++ UpdateWeightsFromMooncakeReqOutput, + UpdateWeightFromDiskReqInput, @@ -78,6 +81,10 @@ class SchedulerWeightUpdaterManager: flush_cache: Callable[..., bool] is_fully_idle: Callable[..., bool] @@ -1362,10 +1558,41 @@ index 77bf823b081..9ab3abe5618 100644 offload_tags: set = field(default_factory=set) stashed_model_static_state: Any = None -@@ -175,6 +182,19 @@ class SchedulerWeightUpdaterManager: +@@ -175,6 +182,50 @@ class SchedulerWeightUpdaterManager: parameter = self.tp_worker.get_weights_by_name(recv_req) return GetWeightsByNameReqOutput(parameter) ++ def init_mooncake_weight_receiver( ++ self, recv_req: InitMooncakeWeightReceiverReqInput ++ ): ++ success, message, descriptor = self.tp_worker.init_mooncake_weight_receiver( ++ recv_req ++ ) ++ if not success: ++ logger.error(message) ++ return InitMooncakeWeightReceiverReqOutput( ++ success=success, message=message, **(descriptor or {}) ++ ) ++ ++ def update_weights_from_mooncake(self, recv_req: UpdateWeightsFromMooncakeReqInput): ++ success, message = self.tp_worker.update_weights_from_mooncake(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 UpdateWeightsFromMooncakeReqOutput(success, message) ++ ++ def destroy_mooncake_weight_receiver( ++ self, recv_req: DestroyMooncakeWeightReceiverReqInput ++ ): ++ success, message = self.tp_worker.destroy_mooncake_weight_receiver(recv_req) ++ if not success: ++ logger.error(message) ++ return DestroyMooncakeWeightReceiverReqOutput(success, message) ++ + def post_process_weights(self, recv_req: PostProcessWeightsReqInput): + success, message = self.tp_worker.post_process_weights(recv_req) + if ( @@ -1424,7 +1651,36 @@ diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sgla index c9939a1fc93..ee25e5e70e0 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py -@@ -48,6 +48,8 @@ from sglang.srt.managers.io_struct import ( +@@ -23,50 +23,58 @@ from sglang.srt.managers.io_struct import ( ++ DestroyMooncakeWeightReceiverReqInput, ++ DestroyMooncakeWeightReceiverReqOutput, + DestroyWeightsUpdateGroupReqInput, + DestroyWeightsUpdateGroupReqOutput, + DetachHiCacheStorageReqInput, + DetachHiCacheStorageReqOutput, + DumperControlReqInput, + DumperControlReqOutput, + ExpertDistributionReq, + ExpertDistributionReqOutput, + ExpertDistributionReqType, + FlushCacheReqInput, + FlushCacheReqOutput, + GetInternalStateReq, + GetInternalStateReqOutput, + GetLoadsReqOutput, + GetWeightsByNameReqInput, + GetWeightsByNameReqOutput, ++ InitMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqOutput, + InitWeightsSendGroupForRemoteInstanceReqInput, + InitWeightsSendGroupForRemoteInstanceReqOutput, + InitWeightsUpdateGroupReqInput, + InitWeightsUpdateGroupReqOutput, + ListExternalCorporaReqInput, + ListExternalCorporaReqOutput, + LoadLoRAAdapterFromTensorsReqInput, + LoadLoRAAdapterFromTensorsReqOutput, + LoadLoRAAdapterReqInput, LoadLoRAAdapterReqOutput, LoRAUpdateOutput, OpenSessionReqInput, @@ -1433,18 +1689,92 @@ index c9939a1fc93..ee25e5e70e0 100644 ProfileReq, ProfileReqOutput, ProfileReqType, -@@ -96,6 +98,7 @@ _COMMUNICATOR_SPECS = [ + ReleaseMemoryOccupationReqInput, + ReleaseMemoryOccupationReqOutput, + RemoveExternalCorpusReqInput, + RemoveExternalCorpusReqOutput, + ResumeMemoryOccupationReqInput, + ResumeMemoryOccupationReqOutput, + SendWeightsToRemoteInstanceReqInput, + SendWeightsToRemoteInstanceReqOutput, + SetInternalStateReq, + SetInternalStateReqOutput, + SlowDownReqInput, + SlowDownReqOutput, + UnloadLoRAAdapterReqInput, + UnloadLoRAAdapterReqOutput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromDistributedReqOutput, + UpdateWeightsFromIPCReqInput, + UpdateWeightsFromIPCReqOutput, ++ UpdateWeightsFromMooncakeReqInput, ++ UpdateWeightsFromMooncakeReqOutput, + UpdateWeightsFromTensorReqInput, +@@ -96,6 +98,10 @@ _COMMUNICATOR_SPECS = [ ("send_weights_to_remote_instance", SendWeightsToRemoteInstanceReqOutput), ("update_weights_from_tensor", UpdateWeightsFromTensorReqOutput), ("update_weights_from_ipc", UpdateWeightsFromIPCReqOutput), ++ ("init_mooncake_weight_receiver", InitMooncakeWeightReceiverReqOutput), ++ ("update_weights_from_mooncake", UpdateWeightsFromMooncakeReqOutput), ++ ("destroy_mooncake_weight_receiver", DestroyMooncakeWeightReceiverReqOutput), + ("post_process_weights", PostProcessWeightsReqOutput), ("get_weights_by_name", GetWeightsByNameReqOutput), ("release_memory_occupation", ReleaseMemoryOccupationReqOutput), ("resume_memory_occupation", ResumeMemoryOccupationReqOutput), -@@ -754,6 +757,16 @@ class TokenizerControlMixin: +@@ -754,6 +757,66 @@ class TokenizerControlMixin: self.auto_create_handle_loop() await self.resume_memory_occupation_communicator(obj) ++ async def init_mooncake_weight_receiver( ++ self: TokenizerManager, ++ obj: InitMooncakeWeightReceiverReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str, dict]: ++ self.auto_create_handle_loop() ++ unsupported = "Mooncake weight receiver currently supports only dp_size=1." ++ if self.server_args.dp_size != 1: ++ return False, unsupported, {} ++ result = (await self.init_mooncake_weight_receiver_communicator(obj))[0] ++ descriptor = { ++ "target_name": result.target_name, ++ "server_name": result.server_name, ++ "rpc_port": result.rpc_port, ++ "buffer_device": result.buffer_device, ++ "buffers": result.buffers or [], ++ } ++ return result.success, result.message, descriptor ++ ++ async def update_weights_from_mooncake( ++ self: TokenizerManager, ++ obj: UpdateWeightsFromMooncakeReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str]: ++ self.auto_create_handle_loop() ++ unsupported = "Mooncake weight receiver currently supports only dp_size=1." ++ if self.server_args.dp_size != 1: ++ return False, unsupported ++ if obj.abort_all_requests: ++ self.abort_request(abort_all=True) ++ async with self.model_update_lock.writer_lock: ++ results = await self.update_weights_from_mooncake_communicator(obj) ++ success, message = FanOutCommunicator.merge_results(results) ++ if success and obj.weight_version is not None: ++ self._update_weight_version_if_provided(obj.weight_version) ++ message += f" Weight version updated to {obj.weight_version}." ++ return success, message ++ ++ async def destroy_mooncake_weight_receiver( ++ self: TokenizerManager, ++ obj: DestroyMooncakeWeightReceiverReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str]: ++ self.auto_create_handle_loop() ++ unsupported = "Mooncake weight receiver currently supports only dp_size=1." ++ if self.server_args.dp_size != 1: ++ return False, unsupported ++ results = await self.destroy_mooncake_weight_receiver_communicator(obj) ++ return FanOutCommunicator.merge_results(results) ++ + async def post_process_weights( + self: TokenizerManager, + obj: PostProcessWeightsReqInput, @@ -1535,7 +1865,12 @@ diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/manager index bd9184408ed..71bbe8f400f 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py -@@ -29,6 +29,7 @@ from sglang.srt.managers.io_struct import ( +@@ -26,12 +26,16 @@ from sglang.srt.managers.io_struct import ( ++ DestroyMooncakeWeightReceiverReqInput, + DestroyWeightsUpdateGroupReqInput, + GetWeightsByNameReqInput, ++ InitMooncakeWeightReceiverReqInput, + InitWeightsSendGroupForRemoteInstanceReqInput, InitWeightsUpdateGroupReqInput, LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, @@ -1543,6 +1878,10 @@ index bd9184408ed..71bbe8f400f 100644 SendWeightsToRemoteInstanceReqInput, UnloadLoRAAdapterReqInput, UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, ++ UpdateWeightsFromMooncakeReqInput, + UpdateWeightsFromTensorReqInput, @@ -98,6 +99,7 @@ class BaseTpWorker(ABC): recv_req.model_path, recv_req.load_format, @@ -1551,7 +1890,7 @@ index bd9184408ed..71bbe8f400f 100644 ) return success, message -@@ -152,6 +154,14 @@ class BaseTpWorker(ABC): +@@ -152,6 +154,47 @@ class BaseTpWorker(ABC): recv_req.shapes, recv_req.group_name, recv_req.load_format, @@ -1559,6 +1898,39 @@ index bd9184408ed..71bbe8f400f 100644 + ) + return success, message + ++ def init_mooncake_weight_receiver( ++ self, recv_req: InitMooncakeWeightReceiverReqInput ++ ): ++ success, message, descriptor = self.model_runner.init_mooncake_weight_receiver( ++ metadata_server=recv_req.metadata_server, ++ protocol=recv_req.protocol, ++ device_name=recv_req.device_name, ++ buffer_size=recv_req.buffer_size, ++ buffer_count=recv_req.buffer_count, ++ rpc_port_base=recv_req.rpc_port_base, ++ group_name=recv_req.group_name, ++ server_name=recv_req.server_name, ++ ) ++ return success, message, descriptor ++ ++ def update_weights_from_mooncake(self, recv_req: UpdateWeightsFromMooncakeReqInput): ++ success, message = self.model_runner.update_weights_from_mooncake( ++ target_name=recv_req.target_name, ++ buffer_slot=recv_req.buffer_slot, ++ remote_addr=recv_req.remote_addr, ++ nbytes=recv_req.nbytes, ++ tensors=recv_req.tensors, ++ load_format=recv_req.load_format, ++ delta=recv_req.delta, ++ ) ++ return success, message ++ ++ def destroy_mooncake_weight_receiver( ++ self, recv_req: DestroyMooncakeWeightReceiverReqInput ++ ): ++ success, message = self.model_runner.destroy_mooncake_weight_receiver() ++ return success, message ++ + def post_process_weights(self, recv_req: PostProcessWeightsReqInput): + success, message = self.model_runner.post_process_weights( + restore_weights_before_load=recv_req.restore_weights_before_load, @@ -1910,10 +2282,207 @@ index 3b30eb0e1f7..d715bc6893d 100644 no_copy_to_cpu=no_copy_to_cpu, ) -@@ -3718,6 +3895,39 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -3718,6 +3895,236 @@ class ModelRunner(ModelRunnerKVCacheMixin): logger.error(f"IPC weight update failed: {e}") return False, str(e) ++ def _default_mooncake_receiver_server_name( ++ self, metadata_server: str, group_name: str ++ ) -> str: ++ if metadata_server != "P2PHANDSHAKE": ++ return f"{group_name}-tp{self.tp_rank}" ++ try: ++ with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: ++ sock.connect(("8.8.8.8", 80)) ++ return sock.getsockname()[0] ++ except OSError: ++ hostname = socket.gethostname() ++ try: ++ return socket.gethostbyname(hostname) ++ except socket.gaierror: ++ return hostname ++ ++ @staticmethod ++ def _format_mooncake_target_name(server_name: str, rpc_port: int) -> str: ++ if ":" in server_name and not server_name.startswith("["): ++ return f"[{server_name}]:{rpc_port}" ++ return f"{server_name}:{rpc_port}" ++ ++ def init_mooncake_weight_receiver( ++ self, ++ metadata_server: str, ++ protocol: str, ++ device_name: str, ++ buffer_size: int, ++ buffer_count: int, ++ rpc_port_base: Optional[int] = None, ++ group_name: str = "slime-mooncake-weight", ++ server_name: Optional[str] = None, ++ ) -> tuple[bool, str, Optional[Dict[str, Any]]]: ++ if rpc_port_base is not None: ++ return ( ++ False, ++ "Mooncake Python TransferEngine does not expose fixed rpc_port; " ++ "unset --mooncake-rpc-port-base for now.", ++ None, ++ ) ++ if buffer_size <= 0: ++ return False, "Mooncake receiver buffer_size must be positive.", None ++ if buffer_count <= 0: ++ return False, "Mooncake receiver buffer_count must be positive.", None ++ ++ base_server_name = server_name or self._default_mooncake_receiver_server_name( ++ metadata_server, group_name ++ ) ++ candidate_devices = [torch.device(self.device)] ++ ++ last_error = None ++ for candidate_device in candidate_devices: ++ try: ++ cleanup_ok, cleanup_message = self.destroy_mooncake_weight_receiver() ++ if not cleanup_ok: ++ return False, cleanup_message, None ++ from mooncake.engine import TransferEngine ++ ++ engine = TransferEngine() ++ ret = engine.initialize( ++ base_server_name, ++ metadata_server, ++ protocol, ++ device_name or "", ++ ) ++ if ret != 0: ++ free_engine = getattr(engine, "freeEngine", None) ++ if callable(free_engine): ++ free_engine() ++ last_error = f"Mooncake TransferEngine.initialize failed: {ret}" ++ continue ++ ++ buffers = [] ++ buffer_descs = [] ++ self._mooncake_weight_receiver = { ++ "engine": engine, ++ "buffers": buffers, ++ "descriptor": {}, ++ } ++ for slot in range(buffer_count): ++ buf = torch.empty( ++ (buffer_size,), dtype=torch.uint8, device=candidate_device ++ ) ++ if buf.is_cuda: ++ torch.cuda.synchronize(buf.device) ++ ret = engine.register_memory(int(buf.data_ptr()), int(buf.numel())) ++ if ret != 0: ++ raise RuntimeError(f"Mooncake register_memory failed: {ret}") ++ buffers.append(buf) ++ buffer_descs.append( ++ { ++ "slot": slot, ++ "addr": int(buf.data_ptr()), ++ "capacity": int(buf.numel()), ++ "device": buf.device.type, ++ } ++ ) ++ ++ rpc_port = int(engine.get_rpc_port()) ++ target_name = self._format_mooncake_target_name(base_server_name, rpc_port) ++ buffer_device = buffer_descs[0]["device"] if buffer_descs else candidate_device.type ++ descriptor = { ++ "target_name": target_name, ++ "server_name": base_server_name, ++ "rpc_port": rpc_port, ++ "buffer_device": buffer_device, ++ "buffers": buffer_descs, ++ } ++ self._mooncake_weight_receiver["descriptor"] = descriptor ++ return ( ++ True, ++ f"Succeeded to initialize Mooncake weight receiver on {buffer_device} buffer.", ++ descriptor, ++ ) ++ except Exception as e: ++ last_error = str(e) ++ logger.warning( ++ "Failed to initialize Mooncake weight receiver on %s: %s", ++ candidate_device, ++ e, ++ ) ++ cleanup_ok, cleanup_message = self.destroy_mooncake_weight_receiver() ++ if not cleanup_ok: ++ return False, cleanup_message, None ++ ++ return False, f"Failed to initialize Mooncake weight receiver: {last_error}", None ++ ++ def update_weights_from_mooncake( ++ self, ++ target_name: str, ++ buffer_slot: int, ++ remote_addr: int, ++ nbytes: int, ++ tensors: List[Dict[str, Any]], ++ load_format: Optional[str] = None, ++ delta: Optional[Dict[str, Any]] = None, ++ ) -> tuple[bool, str]: ++ if delta is not None: ++ return False, "Mooncake weight receiver does not support delta manifests yet." ++ state = getattr(self, "_mooncake_weight_receiver", None) ++ if state is None: ++ return False, "Mooncake weight receiver is not initialized." ++ ++ descriptor = state["descriptor"] ++ if target_name != descriptor["target_name"]: ++ return False, f"Mooncake target mismatch: {target_name} != {descriptor['target_name']}" ++ buffers = state["buffers"] ++ if not (0 <= buffer_slot < len(buffers)): ++ return False, f"Invalid Mooncake buffer_slot={buffer_slot}." ++ buf = buffers[buffer_slot] ++ if int(remote_addr) != int(buf.data_ptr()): ++ return False, "Mooncake manifest remote_addr does not match receiver buffer." ++ if nbytes < 0 or nbytes > int(buf.numel()): ++ return False, "Mooncake manifest nbytes exceeds receiver buffer capacity." ++ ++ try: ++ if getattr(buf, "is_cuda", False) and str(self.device).startswith("cuda"): ++ torch.cuda.synchronize(self.device) ++ named_tensors = [] ++ for spec in tensors: ++ offset = int(spec["offset"]) ++ size = int(spec["nbytes"]) ++ if offset < 0 or size < 0 or offset + size > nbytes: ++ return False, f"Invalid Mooncake tensor slice: {spec!r}" ++ dtype = _resolve_torch_dtype(spec["dtype"]) ++ tensor = ( ++ buf.narrow(0, offset, size) ++ .view(dtype) ++ .view(tuple(spec["shape"])) ++ .clone() ++ ) ++ named_tensors.append((spec["name"], tensor)) ++ return self.update_weights_from_tensor(named_tensors, load_format) ++ except Exception as e: ++ logger.error("Mooncake weight update failed: %s", e) ++ return False, f"Mooncake weight update failed: {e}" ++ ++ def destroy_mooncake_weight_receiver(self) -> tuple[bool, str]: ++ state = getattr(self, "_mooncake_weight_receiver", None) ++ if state is None: ++ return True, "Mooncake weight receiver is not initialized." ++ errors = [] ++ engine = state["engine"] ++ for buf in state["buffers"]: ++ ret = engine.unregister_memory(int(buf.data_ptr())) ++ if ret != 0: ++ errors.append(f"unregister_memory({int(buf.data_ptr())})={ret}") ++ free_engine = getattr(engine, "freeEngine", None) ++ if callable(free_engine): ++ ret = free_engine() ++ if ret not in (None, 0): ++ errors.append(f"freeEngine={ret}") ++ self._mooncake_weight_receiver = None ++ if errors: ++ return False, "; ".join(errors) ++ return True, "Succeeded to destroy Mooncake weight receiver." ++ + def post_process_weights( + self, + restore_weights_before_load: bool = False, @@ -1950,7 +2519,7 @@ index 3b30eb0e1f7..d715bc6893d 100644 def prealloc_symmetric_memory_pool(self): # PyTorch mempools never de-fragment memory in OOM scenarios, so we need to pre-allocate a large chunk of memory to limit fragmentation. if ( -@@ -3767,6 +3977,123 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -3767,6 +4148,123 @@ class ModelRunner(ModelRunnerKVCacheMixin): return output diff --git a/docker/patch/v0.5.12.post1/sglang.patch b/docker/patch/v0.5.12.post1/sglang.patch index 861f252460..5d07087ef2 100644 --- a/docker/patch/v0.5.12.post1/sglang.patch +++ b/docker/patch/v0.5.12.post1/sglang.patch @@ -759,11 +759,14 @@ diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/en index 80081fc..cec2c57 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py -@@ -127,6 +127,7 @@ from sglang.srt.managers.io_struct import ( +@@ -127,6 +127,10 @@ from sglang.srt.managers.io_struct import ( OpenSessionReqInput, ParseFunctionCallReq, PauseGenerationReqInput, ++ DestroyMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqInput, + PostProcessWeightsReqInput, ++ UpdateWeightsFromMooncakeReqInput, ProfileReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, @@ -803,10 +806,63 @@ index 80081fc..cec2c57 100644 # server_args.model_config is not serializable but should be excluded by asdict. return { -@@ -1222,6 +1231,23 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re +@@ -1222,6 +1231,76 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) ++@app.post("/init_mooncake_weight_receiver") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def init_mooncake_weight_receiver( ++ obj: InitMooncakeWeightReceiverReqInput, request: Request ++): ++ """Initialize a Mooncake Transfer Engine receiver for trainer-side weight writes.""" ++ success, message, descriptor = ( ++ await _global_state.tokenizer_manager.init_mooncake_weight_receiver( ++ obj, request ++ ) ++ ) ++ content = {"success": success, "message": message} ++ if descriptor: ++ content.update(descriptor) ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ ++@app.post("/update_weights_from_mooncake") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def update_weights_from_mooncake( ++ obj: UpdateWeightsFromMooncakeReqInput, request: Request ++): ++ """Load weights that were already written into the local Mooncake receiver buffer.""" ++ success, message = ( ++ await _global_state.tokenizer_manager.update_weights_from_mooncake( ++ obj, request ++ ) ++ ) ++ content = {"success": success, "message": message} ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ ++@app.post("/destroy_mooncake_weight_receiver") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def destroy_mooncake_weight_receiver( ++ obj: DestroyMooncakeWeightReceiverReqInput, request: Request ++): ++ """Release the Mooncake receiver buffers and Transfer Engine instance.""" ++ success, message = ( ++ await _global_state.tokenizer_manager.destroy_mooncake_weight_receiver( ++ obj, request ++ ) ++ ) ++ content = {"success": success, "message": message} ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ +@app.post("/post_process_weights") +@auth_level(AuthLevel.ADMIN_OPTIONAL) +async def post_process_weights(req: PostProcessWeightsReqInput, request: Request): @@ -1150,10 +1206,64 @@ index 293335f..51e5669 100644 # Whether to call torch.cuda.empty_cache() during flush torch_empty_cache: bool = False -@@ -1638,6 +1677,18 @@ class ResumeMemoryOccupationReqOutput(BaseReq): +@@ -1638,6 +1677,72 @@ class ResumeMemoryOccupationReqOutput(BaseReq): pass ++@dataclass ++class InitMooncakeWeightReceiverReqInput(BaseReq): ++ metadata_server: str = "P2PHANDSHAKE" ++ protocol: str = "tcp" ++ device_name: str = "" ++ buffer_size: int = 0 ++ buffer_count: int = 1 ++ rpc_port_base: Optional[int] = None ++ group_name: str = "slime-mooncake-weight" ++ server_name: Optional[str] = None ++ ++ ++@dataclass ++class InitMooncakeWeightReceiverReqOutput(BaseReq): ++ success: bool ++ message: str ++ target_name: Optional[str] = None ++ server_name: Optional[str] = None ++ rpc_port: Optional[int] = None ++ buffer_device: Optional[str] = None ++ buffers: Optional[List[Dict[str, Any]]] = None ++ ++ ++@dataclass ++class UpdateWeightsFromMooncakeReqInput(BaseReq): ++ weight_version: Optional[str] ++ target_name: str ++ buffer_slot: int ++ remote_addr: int ++ nbytes: int ++ tensors: List[Dict[str, Any]] ++ load_format: Optional[str] = None ++ delta: Optional[Dict[str, Any]] = None ++ flush_cache: bool = True ++ abort_all_requests: bool = False ++ ++ ++@dataclass ++class UpdateWeightsFromMooncakeReqOutput(BaseReq): ++ success: bool ++ message: str ++ ++ ++@dataclass ++class DestroyMooncakeWeightReceiverReqInput(BaseReq): ++ pass ++ ++ ++@dataclass ++class DestroyMooncakeWeightReceiverReqOutput(BaseReq): ++ success: bool ++ message: str ++ ++ +@dataclass +class PostProcessWeightsReqInput(BaseReq): + restore_weights_before_load: bool = False @@ -1189,18 +1299,24 @@ diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/manager index 143054c..40e05e2 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py -@@ -124,6 +124,7 @@ from sglang.srt.managers.io_struct import ( +@@ -124,6 +124,10 @@ from sglang.srt.managers.io_struct import ( LoadLoRAAdapterReqOutput, OpenSessionReqInput, PauseGenerationReqInput, ++ DestroyMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqInput, + PostProcessWeightsReqInput, ++ UpdateWeightsFromMooncakeReqInput, ProfileReq, ReleaseMemoryOccupationReqInput, RemoveExternalCorpusReqInput, -@@ -1452,6 +1453,7 @@ class Scheduler( +@@ -1452,6 +1453,10 @@ class Scheduler( ), (UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor), (UpdateWeightsFromIPCReqInput, self.update_weights_from_ipc), ++ (InitMooncakeWeightReceiverReqInput, self.init_mooncake_weight_receiver), ++ (UpdateWeightsFromMooncakeReqInput, self.update_weights_from_mooncake), ++ (DestroyMooncakeWeightReceiverReqInput, self.destroy_mooncake_weight_receiver), + (PostProcessWeightsReqInput, self.post_process_weights), (GetWeightsByNameReqInput, self.get_weights_by_name), (ReleaseMemoryOccupationReqInput, self.release_memory_occupation), @@ -1253,7 +1369,7 @@ diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/pyth index f2daf64..534d1b2 100644 --- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py +++ b/python/sglang/srt/managers/scheduler_update_weights_mixin.py -@@ -12,6 +12,7 @@ from sglang.srt.constants import ( +@@ -12,17 +12,26 @@ from sglang.srt.constants import ( GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS, ) @@ -1261,8 +1377,14 @@ index f2daf64..534d1b2 100644 from sglang.srt.managers.io_struct import ( CheckWeightsReqInput, CheckWeightsReqOutput, -@@ -21,6 +22,8 @@ from sglang.srt.managers.io_struct import ( ++ DestroyMooncakeWeightReceiverReqInput, ++ DestroyMooncakeWeightReceiverReqOutput, + DestroyWeightsUpdateGroupReqInput, + DestroyWeightsUpdateGroupReqOutput, + GetWeightsByNameReqInput, GetWeightsByNameReqOutput, ++ InitMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqOutput, InitWeightsUpdateGroupReqInput, InitWeightsUpdateGroupReqOutput, + PostProcessWeightsReqInput, @@ -1270,10 +1392,47 @@ index f2daf64..534d1b2 100644 ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, ResumeMemoryOccupationReqInput, -@@ -120,6 +123,11 @@ class SchedulerUpdateWeightsMixin: + ResumeMemoryOccupationReqOutput, ++ UpdateWeightsFromMooncakeReqInput, ++ UpdateWeightsFromMooncakeReqOutput, + UpdateWeightFromDiskReqInput, +@@ -120,6 +123,44 @@ class SchedulerUpdateWeightsMixin: torch.distributed.barrier(group=self.tp_cpu_group) return UpdateWeightsFromIPCReqOutput(success, message) ++ def init_mooncake_weight_receiver( ++ self: Scheduler, recv_req: InitMooncakeWeightReceiverReqInput ++ ): ++ success, message, descriptor = self.tp_worker.init_mooncake_weight_receiver( ++ recv_req ++ ) ++ if not success: ++ logger.error(message) ++ return InitMooncakeWeightReceiverReqOutput( ++ success=success, message=message, **(descriptor or {}) ++ ) ++ ++ def update_weights_from_mooncake( ++ self: Scheduler, recv_req: UpdateWeightsFromMooncakeReqInput ++ ): ++ success, message = self.tp_worker.update_weights_from_mooncake(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 UpdateWeightsFromMooncakeReqOutput(success, message) ++ ++ def destroy_mooncake_weight_receiver( ++ self: Scheduler, recv_req: DestroyMooncakeWeightReceiverReqInput ++ ): ++ success, message = self.tp_worker.destroy_mooncake_weight_receiver(recv_req) ++ if not success: ++ logger.error(message) ++ return DestroyMooncakeWeightReceiverReqOutput(success, message) ++ + def post_process_weights(self, recv_req: PostProcessWeightsReqInput): + """Optional post-processing for updated weights (e.g., Marlin conversion).""" + success, message = self.tp_worker.post_process_weights(recv_req) @@ -1318,27 +1477,86 @@ diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sgla index 05382e0..7eea220 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py -@@ -76,6 +76,8 @@ from sglang.srt.managers.io_struct import ( +@@ -76,6 +76,14 @@ from sglang.srt.managers.io_struct import ( UpdateWeightsFromDistributedReqOutput, UpdateWeightsFromIPCReqInput, UpdateWeightsFromIPCReqOutput, ++ DestroyMooncakeWeightReceiverReqInput, ++ DestroyMooncakeWeightReceiverReqOutput, ++ InitMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqOutput, + PostProcessWeightsReqInput, + PostProcessWeightsReqOutput, ++ UpdateWeightsFromMooncakeReqInput, ++ UpdateWeightsFromMooncakeReqOutput, UpdateWeightsFromTensorReqInput, UpdateWeightsFromTensorReqOutput, ) -@@ -102,6 +104,7 @@ _COMMUNICATOR_SPECS = [ +@@ -102,6 +104,10 @@ _COMMUNICATOR_SPECS = [ ("send_weights_to_remote_instance", SendWeightsToRemoteInstanceReqOutput), ("update_weights_from_tensor", UpdateWeightsFromTensorReqOutput), ("update_weights_from_ipc", UpdateWeightsFromIPCReqOutput), ++ ("init_mooncake_weight_receiver", InitMooncakeWeightReceiverReqOutput), ++ ("update_weights_from_mooncake", UpdateWeightsFromMooncakeReqOutput), ++ ("destroy_mooncake_weight_receiver", DestroyMooncakeWeightReceiverReqOutput), + ("post_process_weights", PostProcessWeightsReqOutput), ("get_weights_by_name", GetWeightsByNameReqOutput), ("release_memory_occupation", ReleaseMemoryOccupationReqOutput), ("resume_memory_occupation", ResumeMemoryOccupationReqOutput), -@@ -531,6 +534,17 @@ class TokenizerControlMixin: +@@ -531,6 +534,67 @@ class TokenizerControlMixin: return success, message ++ async def init_mooncake_weight_receiver( ++ self: TokenizerManager, ++ obj: InitMooncakeWeightReceiverReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str, dict]: ++ self.auto_create_handle_loop() ++ unsupported = "Mooncake weight receiver currently supports only dp_size=1." ++ if self.server_args.dp_size != 1: ++ return False, unsupported, {} ++ result = (await self.init_mooncake_weight_receiver_communicator(obj))[0] ++ descriptor = { ++ "target_name": result.target_name, ++ "server_name": result.server_name, ++ "rpc_port": result.rpc_port, ++ "buffer_device": result.buffer_device, ++ "buffers": result.buffers or [], ++ } ++ return result.success, result.message, descriptor ++ ++ async def update_weights_from_mooncake( ++ self: TokenizerManager, ++ obj: UpdateWeightsFromMooncakeReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str]: ++ self.auto_create_handle_loop() ++ unsupported = "Mooncake weight receiver currently supports only dp_size=1." ++ if self.server_args.dp_size != 1: ++ return False, unsupported ++ if obj.abort_all_requests: ++ self.abort_request(abort_all=True) ++ async with self.model_update_lock.writer_lock: ++ results = await self.update_weights_from_mooncake_communicator(obj) ++ success, message = FanOutCommunicator.merge_results(results) ++ if success and obj.weight_version is not None: ++ self._update_weight_version_if_provided(obj.weight_version) ++ message += f" Weight version updated to {obj.weight_version}." ++ return success, message ++ ++ async def destroy_mooncake_weight_receiver( ++ self: TokenizerManager, ++ obj: DestroyMooncakeWeightReceiverReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str]: ++ self.auto_create_handle_loop() ++ unsupported = "Mooncake weight receiver currently supports only dp_size=1." ++ if self.server_args.dp_size != 1: ++ return False, unsupported ++ results = await self.destroy_mooncake_weight_receiver_communicator(obj) ++ return FanOutCommunicator.merge_results(results) ++ + async def post_process_weights( + self: TokenizerManager, + obj: PostProcessWeightsReqInput, @@ -1430,11 +1648,14 @@ diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/manager index 60e105d..6a37a82 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py -@@ -29,6 +29,7 @@ from sglang.srt.managers.io_struct import ( +@@ -29,6 +29,10 @@ from sglang.srt.managers.io_struct import ( InitWeightsUpdateGroupReqInput, LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, ++ DestroyMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqInput, + PostProcessWeightsReqInput, ++ UpdateWeightsFromMooncakeReqInput, SendWeightsToRemoteInstanceReqInput, UnloadLoRAAdapterReqInput, UpdateWeightFromDiskReqInput, @@ -1454,10 +1675,43 @@ index 60e105d..6a37a82 100644 ) return success, message -@@ -171,6 +174,11 @@ class BaseTpWorker(ABC): +@@ -171,6 +174,44 @@ class BaseTpWorker(ABC): success, message = self.model_runner.update_weights_from_ipc(recv_req) return success, message ++ def init_mooncake_weight_receiver( ++ self, recv_req: InitMooncakeWeightReceiverReqInput ++ ): ++ success, message, descriptor = self.model_runner.init_mooncake_weight_receiver( ++ metadata_server=recv_req.metadata_server, ++ protocol=recv_req.protocol, ++ device_name=recv_req.device_name, ++ buffer_size=recv_req.buffer_size, ++ buffer_count=recv_req.buffer_count, ++ rpc_port_base=recv_req.rpc_port_base, ++ group_name=recv_req.group_name, ++ server_name=recv_req.server_name, ++ ) ++ return success, message, descriptor ++ ++ def update_weights_from_mooncake(self, recv_req: UpdateWeightsFromMooncakeReqInput): ++ success, message = self.model_runner.update_weights_from_mooncake( ++ target_name=recv_req.target_name, ++ buffer_slot=recv_req.buffer_slot, ++ remote_addr=recv_req.remote_addr, ++ nbytes=recv_req.nbytes, ++ tensors=recv_req.tensors, ++ load_format=recv_req.load_format, ++ delta=recv_req.delta, ++ ) ++ return success, message ++ ++ def destroy_mooncake_weight_receiver( ++ self, recv_req: DestroyMooncakeWeightReceiverReqInput ++ ): ++ success, message = self.model_runner.destroy_mooncake_weight_receiver() ++ return success, message ++ + def post_process_weights(self, recv_req: PostProcessWeightsReqInput): + """Perform optional post-processing on the updated model weights (e.g., Marlin conversion).""" + success, message = self.model_runner.post_process_weights(recv_req) @@ -1643,7 +1897,7 @@ index eff3a56..5b1eec9 100644 from dataclasses import dataclass from pathlib import Path -from typing import Callable, List, Optional, Tuple, Union -+from typing import Callable, Dict, List, Optional, Tuple, Union ++from typing import Any, Callable, Dict, List, Optional, Tuple, Union import torch import torch.distributed as dist @@ -1910,10 +2164,217 @@ index eff3a56..5b1eec9 100644 no_copy_to_cpu=no_copy_to_cpu, ) -@@ -3641,6 +3828,161 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -3641,6 +3828,368 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) return output ++ def _default_mooncake_receiver_server_name( ++ self, metadata_server: str, group_name: str ++ ) -> str: ++ if metadata_server != "P2PHANDSHAKE": ++ return f"{group_name}-tp{self.tp_rank}" ++ try: ++ with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: ++ sock.connect(("8.8.8.8", 80)) ++ return sock.getsockname()[0] ++ except OSError: ++ hostname = socket.gethostname() ++ try: ++ return socket.gethostbyname(hostname) ++ except socket.gaierror: ++ return hostname ++ ++ @staticmethod ++ def _format_mooncake_target_name(server_name: str, rpc_port: int) -> str: ++ if ":" in server_name and not server_name.startswith("["): ++ return f"[{server_name}]:{rpc_port}" ++ return f"{server_name}:{rpc_port}" ++ ++ @staticmethod ++ def _resolve_mooncake_torch_dtype(dtype: Any) -> torch.dtype: ++ if isinstance(dtype, torch.dtype): ++ return dtype ++ dtype_name = str(dtype).replace("torch.", "") ++ resolved = getattr(torch, dtype_name, None) ++ if not isinstance(resolved, torch.dtype): ++ raise ValueError(f"Unsupported Mooncake tensor dtype: {dtype}") ++ return resolved ++ ++ def init_mooncake_weight_receiver( ++ self, ++ metadata_server: str, ++ protocol: str, ++ device_name: str, ++ buffer_size: int, ++ buffer_count: int, ++ rpc_port_base: Optional[int] = None, ++ group_name: str = "slime-mooncake-weight", ++ server_name: Optional[str] = None, ++ ) -> tuple[bool, str, Optional[Dict[str, Any]]]: ++ if rpc_port_base is not None: ++ return ( ++ False, ++ "Mooncake Python TransferEngine does not expose fixed rpc_port; " ++ "unset --mooncake-rpc-port-base for now.", ++ None, ++ ) ++ if buffer_size <= 0: ++ return False, "Mooncake receiver buffer_size must be positive.", None ++ if buffer_count <= 0: ++ return False, "Mooncake receiver buffer_count must be positive.", None ++ ++ base_server_name = server_name or self._default_mooncake_receiver_server_name( ++ metadata_server, group_name ++ ) ++ candidate_devices = [torch.device(self.device)] ++ ++ last_error = None ++ for candidate_device in candidate_devices: ++ try: ++ cleanup_ok, cleanup_message = self.destroy_mooncake_weight_receiver() ++ if not cleanup_ok: ++ return False, cleanup_message, None ++ from mooncake.engine import TransferEngine ++ ++ engine = TransferEngine() ++ ret = engine.initialize( ++ base_server_name, ++ metadata_server, ++ protocol, ++ device_name or "", ++ ) ++ if ret != 0: ++ free_engine = getattr(engine, "freeEngine", None) ++ if callable(free_engine): ++ free_engine() ++ last_error = f"Mooncake TransferEngine.initialize failed: {ret}" ++ continue ++ ++ buffers = [] ++ buffer_descs = [] ++ self._mooncake_weight_receiver = { ++ "engine": engine, ++ "buffers": buffers, ++ "descriptor": {}, ++ } ++ for slot in range(buffer_count): ++ buf = torch.empty( ++ (buffer_size,), dtype=torch.uint8, device=candidate_device ++ ) ++ if buf.is_cuda: ++ torch.cuda.synchronize(buf.device) ++ ret = engine.register_memory(int(buf.data_ptr()), int(buf.numel())) ++ if ret != 0: ++ raise RuntimeError(f"Mooncake register_memory failed: {ret}") ++ buffers.append(buf) ++ buffer_descs.append( ++ { ++ "slot": slot, ++ "addr": int(buf.data_ptr()), ++ "capacity": int(buf.numel()), ++ "device": buf.device.type, ++ } ++ ) ++ ++ rpc_port = int(engine.get_rpc_port()) ++ target_name = self._format_mooncake_target_name(base_server_name, rpc_port) ++ buffer_device = buffer_descs[0]["device"] if buffer_descs else candidate_device.type ++ descriptor = { ++ "target_name": target_name, ++ "server_name": base_server_name, ++ "rpc_port": rpc_port, ++ "buffer_device": buffer_device, ++ "buffers": buffer_descs, ++ } ++ self._mooncake_weight_receiver["descriptor"] = descriptor ++ return ( ++ True, ++ f"Succeeded to initialize Mooncake weight receiver on {buffer_device} buffer.", ++ descriptor, ++ ) ++ except Exception as e: ++ last_error = str(e) ++ logger.warning( ++ "Failed to initialize Mooncake weight receiver on %s: %s", ++ candidate_device, ++ e, ++ ) ++ cleanup_ok, cleanup_message = self.destroy_mooncake_weight_receiver() ++ if not cleanup_ok: ++ return False, cleanup_message, None ++ ++ return False, f"Failed to initialize Mooncake weight receiver: {last_error}", None ++ ++ def update_weights_from_mooncake( ++ self, ++ target_name: str, ++ buffer_slot: int, ++ remote_addr: int, ++ nbytes: int, ++ tensors: List[Dict[str, Any]], ++ load_format: Optional[str] = None, ++ delta: Optional[Dict[str, Any]] = None, ++ ) -> tuple[bool, str]: ++ if delta is not None: ++ return False, "Mooncake weight receiver does not support delta manifests yet." ++ state = getattr(self, "_mooncake_weight_receiver", None) ++ if state is None: ++ return False, "Mooncake weight receiver is not initialized." ++ ++ descriptor = state["descriptor"] ++ if target_name != descriptor["target_name"]: ++ return False, f"Mooncake target mismatch: {target_name} != {descriptor['target_name']}" ++ buffers = state["buffers"] ++ if not (0 <= buffer_slot < len(buffers)): ++ return False, f"Invalid Mooncake buffer_slot={buffer_slot}." ++ buf = buffers[buffer_slot] ++ if int(remote_addr) != int(buf.data_ptr()): ++ return False, "Mooncake manifest remote_addr does not match receiver buffer." ++ if nbytes < 0 or nbytes > int(buf.numel()): ++ return False, "Mooncake manifest nbytes exceeds receiver buffer capacity." ++ ++ try: ++ if getattr(buf, "is_cuda", False) and str(self.device).startswith("cuda"): ++ torch.cuda.synchronize(self.device) ++ named_tensors = [] ++ for spec in tensors: ++ offset = int(spec["offset"]) ++ size = int(spec["nbytes"]) ++ if offset < 0 or size < 0 or offset + size > nbytes: ++ return False, f"Invalid Mooncake tensor slice: {spec!r}" ++ dtype = self._resolve_mooncake_torch_dtype(spec["dtype"]) ++ tensor = ( ++ buf.narrow(0, offset, size) ++ .view(dtype) ++ .view(tuple(spec["shape"])) ++ .clone() ++ ) ++ named_tensors.append((spec["name"], tensor)) ++ return self.update_weights_from_tensor(named_tensors, load_format) ++ except Exception as e: ++ logger.error("Mooncake weight update failed: %s", e) ++ return False, f"Mooncake weight update failed: {e}" ++ ++ def destroy_mooncake_weight_receiver(self) -> tuple[bool, str]: ++ state = getattr(self, "_mooncake_weight_receiver", None) ++ if state is None: ++ return True, "Mooncake weight receiver is not initialized." ++ errors = [] ++ engine = state["engine"] ++ for buf in state["buffers"]: ++ ret = engine.unregister_memory(int(buf.data_ptr())) ++ if ret != 0: ++ errors.append(f"unregister_memory({int(buf.data_ptr())})={ret}") ++ free_engine = getattr(engine, "freeEngine", None) ++ if callable(free_engine): ++ ret = free_engine() ++ if ret not in (None, 0): ++ errors.append(f"freeEngine={ret}") ++ self._mooncake_weight_receiver = None ++ if errors: ++ return False, "; ".join(errors) ++ return True, "Succeeded to destroy Mooncake weight receiver." ++ + def post_process_weights(self, recv_req): + """ + Execute post-processing logic for model weights, such as Marlin quantization format conversion. diff --git a/docker/patch/v0.5.9/sglang.patch b/docker/patch/v0.5.9/sglang.patch index e9145702e1..d826e0a5af 100644 --- a/docker/patch/v0.5.9/sglang.patch +++ b/docker/patch/v0.5.9/sglang.patch @@ -1,5 +1,5 @@ diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py -index 6fbd1db823..f80ec11bb4 100644 +index 6fbd1db82..f80ec11bb 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -274,6 +274,7 @@ class ModelConfig: @@ -11,7 +11,7 @@ index 6fbd1db823..f80ec11bb4 100644 ]: self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN" diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py -index da4629e525..c03f98231a 100644 +index da4629e52..c03f98231 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -17,6 +17,7 @@ class KVArgs: @@ -23,7 +23,7 @@ index da4629e525..c03f98231a 100644 aux_data_lens: List[int] aux_item_lens: List[int] diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py -index 67fe82ad67..ed5fa7b0e3 100644 +index 67fe82ad6..ed5fa7b0e 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -333,6 +333,10 @@ class CommonKVReceiver(BaseKVReceiver): @@ -60,7 +60,7 @@ index 67fe82ad67..ed5fa7b0e3 100644 "prefill_pp_size": self.pp_size, "prefill_page_size": self.page_size, diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py -index 1d8baf0028..1ebb959298 100644 +index 1d8baf002..1ebb95929 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -21,6 +21,7 @@ Life cycle of a request in the decode server @@ -246,7 +246,7 @@ index 1d8baf0028..1ebb959298 100644 if not hasattr(self, "polling_count"): diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py -index d0d4efd958..b3a207063e 100644 +index d0d4efd95..b3a207063 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -30,7 +30,7 @@ from sglang.srt.disaggregation.common.utils import ( @@ -434,7 +434,7 @@ index d0d4efd958..b3a207063e 100644 def _register_kv_args(self): for bootstrap_info in self.bootstrap_infos: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py -index fbc8016351..7de53ba292 100644 +index fbc801635..7de53ba29 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -20,6 +20,7 @@ Life cycle of a request in the prefill server @@ -582,7 +582,7 @@ index fbc8016351..7de53ba292 100644 release_kv_cache(req, self.tree_cache) # unlock the tree req.finished_reason = FINISH_LENGTH(length=0) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py -index 6d58f415a7..84723c342c 100644 +index 6d58f415a..84723c342 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -21,6 +21,17 @@ if TYPE_CHECKING: @@ -774,7 +774,7 @@ index 6d58f415a7..84723c342c 100644 ######################### diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py -index 0ed5a1b44b..67e33c650d 100644 +index 0ed5a1b44..67e33c650 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -52,6 +52,7 @@ from sglang.srt.managers.io_struct import ( @@ -811,18 +811,21 @@ index 0ed5a1b44b..67e33c650d 100644 """Get weights by parameter name.""" obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py -index 1d6816c010..402b42e05b 100644 +index 1d6816c01..db6e0b92b 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py -@@ -115,6 +115,7 @@ from sglang.srt.managers.io_struct import ( +@@ -115,6 +115,10 @@ from sglang.srt.managers.io_struct import ( OpenSessionReqInput, ParseFunctionCallReq, PauseGenerationReqInput, ++ DestroyMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqInput, + PostProcessWeightsReqInput, ++ UpdateWeightsFromMooncakeReqInput, ProfileReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, -@@ -574,10 +575,8 @@ async def model_info(): +@@ -574,10 +578,8 @@ async def model_info(): @app.get("/weight_version") async def weight_version(): """Get the current weight version.""" @@ -835,7 +838,7 @@ index 1d6816c010..402b42e05b 100644 @app.get("/get_server_info") -@@ -594,9 +593,19 @@ async def get_server_info(): +@@ -594,9 +596,19 @@ async def get_server_info(): async def server_info(): """Get the server information.""" # Returns internal states per DP. @@ -858,10 +861,63 @@ index 1d6816c010..402b42e05b 100644 # This field is not serializable. if hasattr(_global_state.tokenizer_manager.server_args, "model_config"): -@@ -1084,6 +1093,23 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re +@@ -1084,6 +1096,76 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) ++@app.post("/init_mooncake_weight_receiver") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def init_mooncake_weight_receiver( ++ obj: InitMooncakeWeightReceiverReqInput, request: Request ++): ++ """Initialize a Mooncake Transfer Engine receiver for trainer-side weight writes.""" ++ success, message, descriptor = ( ++ await _global_state.tokenizer_manager.init_mooncake_weight_receiver( ++ obj, request ++ ) ++ ) ++ content = {"success": success, "message": message} ++ if descriptor: ++ content.update(descriptor) ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ ++@app.post("/update_weights_from_mooncake") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def update_weights_from_mooncake( ++ obj: UpdateWeightsFromMooncakeReqInput, request: Request ++): ++ """Load weights that were already written into the local Mooncake receiver buffer.""" ++ success, message = ( ++ await _global_state.tokenizer_manager.update_weights_from_mooncake( ++ obj, request ++ ) ++ ) ++ content = {"success": success, "message": message} ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ ++@app.post("/destroy_mooncake_weight_receiver") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def destroy_mooncake_weight_receiver( ++ obj: DestroyMooncakeWeightReceiverReqInput, request: Request ++): ++ """Release the Mooncake receiver buffers and Transfer Engine instance.""" ++ success, message = ( ++ await _global_state.tokenizer_manager.destroy_mooncake_weight_receiver( ++ obj, request ++ ) ++ ) ++ content = {"success": success, "message": message} ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ +@app.post("/post_process_weights") +@auth_level(AuthLevel.ADMIN_OPTIONAL) +async def post_process_weights(req: PostProcessWeightsReqInput, request: Request): @@ -883,7 +939,7 @@ index 1d6816c010..402b42e05b 100644 @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Request): diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py -index 8293796a2e..bff34e4221 100644 +index 8293796a2..bff34e422 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -244,6 +244,7 @@ class Envs: @@ -895,7 +951,7 @@ index 8293796a2e..bff34e4221 100644 # Scheduler: others: SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period. diff --git a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py -index 1cdf65b91c..4783cd18fb 100644 +index 1cdf65b91..4783cd18f 100644 --- a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py +++ b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py @@ -630,7 +630,6 @@ def _get_k_and_s_triton( @@ -915,7 +971,7 @@ index 1cdf65b91c..4783cd18fb 100644 buf_numel_per_page: tl.constexpr, index_head_dim: tl.constexpr, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py -index ca54a931b7..3540f77bae 100644 +index ca54a931b..3540f77ba 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1,6 +1,7 @@ @@ -1000,7 +1056,7 @@ index ca54a931b7..3540f77bae 100644 if enable_dual_stream: current_stream = torch.cuda.current_stream() diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py -index de8a07ab30..5c9f4813a6 100644 +index de8a07ab3..5c9f4813a 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -697,6 +697,7 @@ class FusedMoE(torch.nn.Module): @@ -1020,7 +1076,7 @@ index de8a07ab30..5c9f4813a6 100644 ) diff --git a/python/sglang/srt/layers/moe/routed_experts_capturer.py b/python/sglang/srt/layers/moe/routed_experts_capturer.py -index 00bd687555..12d5577af2 100644 +index 00bd68755..12d5577af 100644 --- a/python/sglang/srt/layers/moe/routed_experts_capturer.py +++ b/python/sglang/srt/layers/moe/routed_experts_capturer.py @@ -8,10 +8,15 @@ import torch @@ -1081,7 +1137,7 @@ index 00bd687555..12d5577af2 100644 def get_routed_experts( diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py -index 4cbfed6f90..88b4527443 100644 +index 4cbfed6f9..88b452744 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -499,7 +499,7 @@ class CompressedTensorsConfig(QuantizationConfig): @@ -1104,7 +1160,7 @@ index 4cbfed6f90..88b4527443 100644 self, layer: torch.nn.Module, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py -index 6264f36d04..f0310e305e 100644 +index 6264f36d0..f0310e305 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py @@ -17,7 +17,10 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import ( @@ -1223,7 +1279,7 @@ index 6264f36d04..f0310e305e 100644 is_k_full=self.is_k_full, routed_scaling_factor=self.moe_runner_config.routed_scaling_factor, diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py -index 6522278603..7d3a5d0c4c 100644 +index 652227860..7d3a5d0c4 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -405,6 +405,17 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): @@ -1245,7 +1301,7 @@ index 6522278603..7d3a5d0c4c 100644 def handle_multimodal_decode_req(self, recv_obj: BatchMultimodalDecodeReq): diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index ff17745673..f947e71d7d 100644 +index ff1774567..66805d02f 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -101,6 +101,42 @@ class RequestTimingMetricsMixin: @@ -1291,10 +1347,64 @@ index ff17745673..f947e71d7d 100644 @dataclass class SpeculativeDecodingMetricsMixin: -@@ -1403,6 +1439,20 @@ class UpdateWeightsFromIPCReqOutput(BaseReq): +@@ -1403,6 +1439,74 @@ class UpdateWeightsFromIPCReqOutput(BaseReq): message: str ++@dataclass ++class InitMooncakeWeightReceiverReqInput(BaseReq): ++ metadata_server: str = "P2PHANDSHAKE" ++ protocol: str = "tcp" ++ device_name: str = "" ++ buffer_size: int = 0 ++ buffer_count: int = 1 ++ rpc_port_base: Optional[int] = None ++ group_name: str = "slime-mooncake-weight" ++ server_name: Optional[str] = None ++ ++ ++@dataclass ++class InitMooncakeWeightReceiverReqOutput(BaseReq): ++ success: bool ++ message: str ++ target_name: Optional[str] = None ++ server_name: Optional[str] = None ++ rpc_port: Optional[int] = None ++ buffer_device: Optional[str] = None ++ buffers: Optional[List[Dict[str, Any]]] = None ++ ++ ++@dataclass ++class UpdateWeightsFromMooncakeReqInput(BaseReq): ++ weight_version: Optional[str] ++ target_name: str ++ buffer_slot: int ++ remote_addr: int ++ nbytes: int ++ tensors: List[Dict[str, Any]] ++ load_format: Optional[str] = None ++ delta: Optional[Dict[str, Any]] = None ++ flush_cache: bool = True ++ abort_all_requests: bool = False ++ ++ ++@dataclass ++class UpdateWeightsFromMooncakeReqOutput(BaseReq): ++ success: bool ++ message: str ++ ++ ++@dataclass ++class DestroyMooncakeWeightReceiverReqInput(BaseReq): ++ pass ++ ++ ++@dataclass ++class DestroyMooncakeWeightReceiverReqOutput(BaseReq): ++ success: bool ++ message: str ++ ++ +@dataclass +class PostProcessWeightsReqInput(BaseReq): + # Whether to restore weights before loading new weights @@ -1312,7 +1422,7 @@ index ff17745673..f947e71d7d 100644 @dataclass class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq): success: bool -@@ -1802,6 +1852,10 @@ class GetLoadReqOutput(BaseReq): +@@ -1802,6 +1906,10 @@ class GetLoadReqOutput(BaseReq): num_waiting_reqs: int num_tokens: int ts_tic: float @@ -1324,7 +1434,7 @@ index ff17745673..f947e71d7d 100644 @dataclass diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py -index e1236aa0f3..daa598a1f6 100644 +index e1236aa0f..daa598a1f 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -142,6 +142,39 @@ def _handle_output_by_index(output, i): @@ -1520,7 +1630,7 @@ index e1236aa0f3..daa598a1f6 100644 class SenderWrapper: def __init__(self, port_args: PortArgs, send_to_scheduler: zmq.Socket): diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py -index c079957980..dd8ca7167d 100644 +index c07995798..dd8ca7167 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1869,7 +1869,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): @@ -1536,27 +1646,33 @@ index c079957980..dd8ca7167d 100644 break diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index a9ff0ac94b..a50dd5122b 100644 +index a9ff0ac94..80326d3b1 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py -@@ -114,6 +114,7 @@ from sglang.srt.managers.io_struct import ( +@@ -114,6 +114,10 @@ from sglang.srt.managers.io_struct import ( OpenSessionReqInput, OpenSessionReqOutput, PauseGenerationReqInput, ++ DestroyMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqInput, + PostProcessWeightsReqInput, ++ UpdateWeightsFromMooncakeReqInput, ProfileReq, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, -@@ -1063,6 +1064,7 @@ class Scheduler( +@@ -1063,6 +1067,10 @@ class Scheduler( ), (UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor), (UpdateWeightsFromIPCReqInput, self.update_weights_from_ipc), ++ (InitMooncakeWeightReceiverReqInput, self.init_mooncake_weight_receiver), ++ (UpdateWeightsFromMooncakeReqInput, self.update_weights_from_mooncake), ++ (DestroyMooncakeWeightReceiverReqInput, self.destroy_mooncake_weight_receiver), + (PostProcessWeightsReqInput, self.post_process_weights), (GetWeightsByNameReqInput, self.get_weights_by_name), (ReleaseMemoryOccupationReqInput, self.release_memory_occupation), (ResumeMemoryOccupationReqInput, self.resume_memory_occupation), diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py -index 30b2732b9f..68090b1617 100644 +index 30b2732b9..68090b161 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -609,12 +609,54 @@ class SchedulerMetricsMixin: @@ -1615,7 +1731,7 @@ index 30b2732b9f..68090b1617 100644 def get_loads(self: Scheduler, req: GetLoadsReqInput = None) -> GetLoadsReqOutput: diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py -index 482bc6ca66..fbc4864176 100644 +index 482bc6ca6..fbc486417 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -922,6 +922,18 @@ class SchedulerOutputProcessorMixin: @@ -1784,7 +1900,7 @@ index 482bc6ca66..fbc4864176 100644 embeddings=embeddings, prompt_tokens=prompt_tokens, diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py -index 7d08f12b35..afc045da20 100644 +index 7d08f12b3..afc045da2 100644 --- a/python/sglang/srt/managers/scheduler_profiler_mixin.py +++ b/python/sglang/srt/managers/scheduler_profiler_mixin.py @@ -347,7 +347,7 @@ class SchedulerProfilerMixin: @@ -1797,10 +1913,10 @@ index 7d08f12b35..afc045da20 100644 if self.profile_in_progress: # force trace flush diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py -index 293a843508..244ea4eb1b 100644 +index 293a84350..4c3ddc099 100644 --- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py +++ b/python/sglang/srt/managers/scheduler_update_weights_mixin.py -@@ -12,6 +12,7 @@ from sglang.srt.constants import ( +@@ -12,21 +12,30 @@ from sglang.srt.constants import ( GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS, ) @@ -1808,19 +1924,66 @@ index 293a843508..244ea4eb1b 100644 from sglang.srt.managers.io_struct import ( CheckWeightsReqInput, CheckWeightsReqOutput, -@@ -21,6 +22,8 @@ from sglang.srt.managers.io_struct import ( + DestroyWeightsUpdateGroupReqInput, + DestroyWeightsUpdateGroupReqOutput, ++ DestroyMooncakeWeightReceiverReqInput, ++ DestroyMooncakeWeightReceiverReqOutput, + GetWeightsByNameReqInput, GetWeightsByNameReqOutput, InitWeightsUpdateGroupReqInput, InitWeightsUpdateGroupReqOutput, ++ InitMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqOutput, + PostProcessWeightsReqInput, + PostProcessWeightsReqOutput, ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, ResumeMemoryOccupationReqInput, -@@ -114,6 +117,11 @@ class SchedulerUpdateWeightsMixin: + ResumeMemoryOccupationReqOutput, + UpdateWeightFromDiskReqInput, + UpdateWeightFromDiskReqOutput, ++ UpdateWeightsFromMooncakeReqInput, ++ UpdateWeightsFromMooncakeReqOutput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromDistributedReqOutput, + UpdateWeightsFromIPCReqInput, +@@ -114,6 +123,44 @@ class SchedulerUpdateWeightsMixin: torch.distributed.barrier(group=self.tp_cpu_group) return UpdateWeightsFromIPCReqOutput(success, message) ++ def init_mooncake_weight_receiver( ++ self: Scheduler, recv_req: InitMooncakeWeightReceiverReqInput ++ ): ++ success, message, descriptor = self.tp_worker.init_mooncake_weight_receiver( ++ recv_req ++ ) ++ if not success: ++ logger.error(message) ++ return InitMooncakeWeightReceiverReqOutput( ++ success=success, message=message, **(descriptor or {}) ++ ) ++ ++ def update_weights_from_mooncake( ++ self: Scheduler, recv_req: UpdateWeightsFromMooncakeReqInput ++ ): ++ success, message = self.tp_worker.update_weights_from_mooncake(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 UpdateWeightsFromMooncakeReqOutput(success, message) ++ ++ def destroy_mooncake_weight_receiver( ++ self: Scheduler, recv_req: DestroyMooncakeWeightReceiverReqInput ++ ): ++ success, message = self.tp_worker.destroy_mooncake_weight_receiver(recv_req) ++ if not success: ++ logger.error(message) ++ return DestroyMooncakeWeightReceiverReqOutput(success, message) ++ + def post_process_weights(self, recv_req: PostProcessWeightsReqInput): + """Optional post-processing for updated weights (e.g., Marlin conversion).""" + success, message = self.tp_worker.post_process_weights(recv_req) @@ -1829,7 +1992,7 @@ index 293a843508..244ea4eb1b 100644 def get_weights_by_name(self: Scheduler, recv_req: GetWeightsByNameReqInput): parameter = self.tp_worker.get_weights_by_name(recv_req) return GetWeightsByNameReqOutput(parameter) -@@ -137,6 +145,15 @@ class SchedulerUpdateWeightsMixin: +@@ -137,6 +184,15 @@ class SchedulerUpdateWeightsMixin: self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE) self.flush_cache() @@ -1845,7 +2008,7 @@ index 293a843508..244ea4eb1b 100644 if GPU_MEMORY_TYPE_WEIGHTS in tags: self.stashed_model_static_state = _export_static_state( self.tp_worker.model_runner.model -@@ -177,6 +194,15 @@ class SchedulerUpdateWeightsMixin: +@@ -177,6 +233,15 @@ class SchedulerUpdateWeightsMixin: if GPU_MEMORY_TYPE_KV_CACHE in tags: self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_KV_CACHE) @@ -1862,43 +2025,128 @@ index 293a843508..244ea4eb1b 100644 def check_weights(self: Scheduler, recv_req: CheckWeightsReqInput): diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py -index f2ffa9909d..6e4d1d460b 100644 +index f2ffa9909..87daac566 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py -@@ -59,6 +59,8 @@ from sglang.srt.managers.io_struct import ( +@@ -59,6 +59,14 @@ from sglang.srt.managers.io_struct import ( LoadLoRAAdapterReqOutput, LoRAUpdateOutput, OpenSessionReqInput, ++ DestroyMooncakeWeightReceiverReqInput, ++ DestroyMooncakeWeightReceiverReqOutput, ++ InitMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqOutput, + PostProcessWeightsReqInput, + PostProcessWeightsReqOutput, ++ UpdateWeightsFromMooncakeReqInput, ++ UpdateWeightsFromMooncakeReqOutput, ProfileReq, ProfileReqOutput, ProfileReqType, -@@ -187,6 +189,9 @@ class TokenizerCommunicatorMixin: +@@ -187,6 +195,18 @@ class TokenizerCommunicatorMixin: self.update_weights_from_ipc_communicator = _Communicator( self.send_to_scheduler, server_args.dp_size ) ++ self.init_mooncake_weight_receiver_communicator = _Communicator( ++ self.send_to_scheduler, server_args.dp_size ++ ) ++ self.update_weights_from_mooncake_communicator = _Communicator( ++ self.send_to_scheduler, server_args.dp_size ++ ) ++ self.destroy_mooncake_weight_receiver_communicator = _Communicator( ++ self.send_to_scheduler, server_args.dp_size ++ ) + self.post_process_weights_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 ) -@@ -272,6 +277,10 @@ class TokenizerCommunicatorMixin: +@@ -272,6 +292,22 @@ class TokenizerCommunicatorMixin: UpdateWeightsFromIPCReqOutput, self.update_weights_from_ipc_communicator.handle_recv, ), + ( ++ InitMooncakeWeightReceiverReqOutput, ++ self.init_mooncake_weight_receiver_communicator.handle_recv, ++ ), ++ ( ++ UpdateWeightsFromMooncakeReqOutput, ++ self.update_weights_from_mooncake_communicator.handle_recv, ++ ), ++ ( ++ DestroyMooncakeWeightReceiverReqOutput, ++ self.destroy_mooncake_weight_receiver_communicator.handle_recv, ++ ), ++ ( + PostProcessWeightsReqOutput, + self.post_process_weights_communicator.handle_recv, + ), ( GetWeightsByNameReqOutput, self.get_weights_by_name_communicator.handle_recv, -@@ -522,6 +531,17 @@ class TokenizerCommunicatorMixin: +@@ -522,6 +558,75 @@ class TokenizerCommunicatorMixin: return success, message ++ async def init_mooncake_weight_receiver( ++ self: TokenizerManager, ++ obj: InitMooncakeWeightReceiverReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str, dict]: ++ self.auto_create_handle_loop() ++ unsupported = "Mooncake weight receiver currently supports only dp_size=1." ++ if self.server_args.dp_size != 1: ++ return False, unsupported, {} ++ result = (await self.init_mooncake_weight_receiver_communicator(obj))[0] ++ descriptor = { ++ "target_name": result.target_name, ++ "server_name": result.server_name, ++ "rpc_port": result.rpc_port, ++ "buffer_device": result.buffer_device, ++ "buffers": result.buffers or [], ++ } ++ return result.success, result.message, descriptor ++ ++ async def update_weights_from_mooncake( ++ self: TokenizerManager, ++ obj: UpdateWeightsFromMooncakeReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str]: ++ self.auto_create_handle_loop() ++ unsupported = "Mooncake weight receiver currently supports only dp_size=1." ++ if self.server_args.dp_size != 1: ++ return False, unsupported ++ if obj.abort_all_requests: ++ self.abort_request(abort_all=True) ++ ++ async with self.is_pause_cond: ++ is_paused = self.is_pause ++ ++ lock_context = ( ++ self.model_update_lock.writer_lock if not is_paused else nullcontext() ++ ) ++ async with lock_context: ++ results = await self.update_weights_from_mooncake_communicator(obj) ++ ++ success, message = _Communicator.merge_results(results) ++ if success and obj.weight_version is not None: ++ self._update_weight_version_if_provided(obj.weight_version) ++ message += f" Weight version updated to {obj.weight_version}." ++ return success, message ++ ++ async def destroy_mooncake_weight_receiver( ++ self: TokenizerManager, ++ obj: DestroyMooncakeWeightReceiverReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str]: ++ self.auto_create_handle_loop() ++ unsupported = "Mooncake weight receiver currently supports only dp_size=1." ++ if self.server_args.dp_size != 1: ++ return False, unsupported ++ results = await self.destroy_mooncake_weight_receiver_communicator(obj) ++ return _Communicator.merge_results(results) ++ + async def post_process_weights( + self: TokenizerManager, + obj: PostProcessWeightsReqInput, @@ -1914,7 +2162,7 @@ index f2ffa9909d..6e4d1d460b 100644 self, obj: InitWeightsSendGroupForRemoteInstanceReqInput, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index 0914a5230b..9114e3e713 100644 +index 0914a5230..9114e3e71 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1327,7 +1327,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi @@ -2012,21 +2260,57 @@ index 0914a5230b..9114e3e713 100644 if state.finished: retraction_count = ( diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py -index 86b009df4e..16ebd52ae3 100644 +index 86b009df4..ab013996f 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py -@@ -29,6 +29,7 @@ from sglang.srt.managers.io_struct import ( +@@ -29,6 +29,10 @@ from sglang.srt.managers.io_struct import ( InitWeightsUpdateGroupReqInput, LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, ++ DestroyMooncakeWeightReceiverReqInput, ++ InitMooncakeWeightReceiverReqInput, + PostProcessWeightsReqInput, ++ UpdateWeightsFromMooncakeReqInput, SendWeightsToRemoteInstanceReqInput, UnloadLoRAAdapterReqInput, UpdateWeightFromDiskReqInput, -@@ -168,6 +169,11 @@ class BaseTpWorker(ABC): +@@ -168,6 +172,44 @@ class BaseTpWorker(ABC): success, message = self.model_runner.update_weights_from_ipc(recv_req) return success, message ++ def init_mooncake_weight_receiver( ++ self, recv_req: InitMooncakeWeightReceiverReqInput ++ ): ++ success, message, descriptor = self.model_runner.init_mooncake_weight_receiver( ++ metadata_server=recv_req.metadata_server, ++ protocol=recv_req.protocol, ++ device_name=recv_req.device_name, ++ buffer_size=recv_req.buffer_size, ++ buffer_count=recv_req.buffer_count, ++ rpc_port_base=recv_req.rpc_port_base, ++ group_name=recv_req.group_name, ++ server_name=recv_req.server_name, ++ ) ++ return success, message, descriptor ++ ++ def update_weights_from_mooncake(self, recv_req: UpdateWeightsFromMooncakeReqInput): ++ success, message = self.model_runner.update_weights_from_mooncake( ++ target_name=recv_req.target_name, ++ buffer_slot=recv_req.buffer_slot, ++ remote_addr=recv_req.remote_addr, ++ nbytes=recv_req.nbytes, ++ tensors=recv_req.tensors, ++ load_format=recv_req.load_format, ++ delta=recv_req.delta, ++ ) ++ return success, message ++ ++ def destroy_mooncake_weight_receiver( ++ self, recv_req: DestroyMooncakeWeightReceiverReqInput ++ ): ++ success, message = self.model_runner.destroy_mooncake_weight_receiver() ++ return success, message ++ + def post_process_weights(self, recv_req: PostProcessWeightsReqInput): + """Perform optional post-processing on the updated model weights (e.g., Marlin conversion).""" + success, message = self.model_runner.post_process_weights(recv_req) @@ -2036,7 +2320,7 @@ index 86b009df4e..16ebd52ae3 100644 parameter = self.model_runner.get_weights_by_name( recv_req.name, recv_req.truncate_size diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py -index fa08bb66a4..22c1c2a127 100644 +index fa08bb66a..22c1c2a12 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/python/sglang/srt/mem_cache/allocator.py @@ -411,7 +411,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): @@ -2058,7 +2342,7 @@ index fa08bb66a4..22c1c2a127 100644 prefix_lens, seq_lens, diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py -index d7cd472a98..9cf1185cbb 100644 +index d7cd472a9..9cf1185cb 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -750,9 +750,8 @@ class HiRadixCache(RadixCache): @@ -2098,7 +2382,7 @@ index d7cd472a98..9cf1185cbb 100644 self._inc_hit_count(new_node, chunked) total_prefix_length += prefix_len diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py -index 1d917137c6..669e5c5181 100644 +index 1d917137c..669e5c518 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1777,9 +1777,12 @@ class NSATokenToKVPool(MLATokenToKVPool): @@ -2181,7 +2465,7 @@ index 1d917137c6..669e5c5181 100644 kv_size_bytes = super().get_kv_size_bytes() for index_k_cache in self.index_k_with_scale_buffer: diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py -index 42b169728a..8e799196a4 100644 +index 42b169728..8e799196a 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -495,7 +495,17 @@ class RadixCache(BasePrefixCache): @@ -2216,7 +2500,7 @@ index 42b169728a..8e799196a4 100644 return delta diff --git a/python/sglang/srt/metrics/collector.py b/python/sglang/srt/metrics/collector.py -index 255d41ccc0..f93bedb4dc 100644 +index 255d41ccc..f93bedb4d 100644 --- a/python/sglang/srt/metrics/collector.py +++ b/python/sglang/srt/metrics/collector.py @@ -20,7 +20,10 @@ import time @@ -2399,7 +2683,7 @@ index 255d41ccc0..f93bedb4dc 100644 if self.disagg_mode == DisaggregationMode.NULL: queue_duration = self.forward_entry_time - self.wait_queue_entry_time diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index 275775a73d..e4e2fdc398 100644 +index 275775a73..c2cf40f1d 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -395,7 +395,12 @@ class ModelRunner(ModelRunnerKVCacheMixin): @@ -2451,10 +2735,217 @@ index 275775a73d..e4e2fdc398 100644 if self.eplb_manager is not None: self.eplb_manager.on_forward_pass_end() -@@ -2664,6 +2678,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -2664,6 +2678,249 @@ class ModelRunner(ModelRunnerKVCacheMixin): device=self.device, ) ++ def _default_mooncake_receiver_server_name( ++ self, metadata_server: str, group_name: str ++ ) -> str: ++ if metadata_server != "P2PHANDSHAKE": ++ return f"{group_name}-tp{self.tp_rank}" ++ try: ++ with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: ++ sock.connect(("8.8.8.8", 80)) ++ return sock.getsockname()[0] ++ except OSError: ++ hostname = socket.gethostname() ++ try: ++ return socket.gethostbyname(hostname) ++ except socket.gaierror: ++ return hostname ++ ++ @staticmethod ++ def _format_mooncake_target_name(server_name: str, rpc_port: int) -> str: ++ if ":" in server_name and not server_name.startswith("["): ++ return f"[{server_name}]:{rpc_port}" ++ return f"{server_name}:{rpc_port}" ++ ++ @staticmethod ++ def _resolve_mooncake_torch_dtype(dtype: object) -> torch.dtype: ++ if isinstance(dtype, torch.dtype): ++ return dtype ++ dtype_name = str(dtype).replace("torch.", "") ++ resolved = getattr(torch, dtype_name, None) ++ if not isinstance(resolved, torch.dtype): ++ raise ValueError(f"Unsupported Mooncake tensor dtype: {dtype}") ++ return resolved ++ ++ def init_mooncake_weight_receiver( ++ self, ++ metadata_server: str, ++ protocol: str, ++ device_name: str, ++ buffer_size: int, ++ buffer_count: int, ++ rpc_port_base: Optional[int] = None, ++ group_name: str = "slime-mooncake-weight", ++ server_name: Optional[str] = None, ++ ) -> Tuple[bool, str, Optional[dict[str, object]]]: ++ if rpc_port_base is not None: ++ return ( ++ False, ++ "Mooncake Python TransferEngine does not expose fixed rpc_port; " ++ "unset --mooncake-rpc-port-base for now.", ++ None, ++ ) ++ if buffer_size <= 0: ++ return False, "Mooncake receiver buffer_size must be positive.", None ++ if buffer_count <= 0: ++ return False, "Mooncake receiver buffer_count must be positive.", None ++ ++ base_server_name = server_name or self._default_mooncake_receiver_server_name( ++ metadata_server, group_name ++ ) ++ candidate_devices = [torch.device(self.device)] ++ ++ last_error = None ++ for candidate_device in candidate_devices: ++ try: ++ cleanup_ok, cleanup_message = self.destroy_mooncake_weight_receiver() ++ if not cleanup_ok: ++ return False, cleanup_message, None ++ from mooncake.engine import TransferEngine ++ ++ engine = TransferEngine() ++ ret = engine.initialize( ++ base_server_name, ++ metadata_server, ++ protocol, ++ device_name or "", ++ ) ++ if ret != 0: ++ free_engine = getattr(engine, "freeEngine", None) ++ if callable(free_engine): ++ free_engine() ++ last_error = f"Mooncake TransferEngine.initialize failed: {ret}" ++ continue ++ ++ buffers = [] ++ buffer_descs = [] ++ self._mooncake_weight_receiver = { ++ "engine": engine, ++ "buffers": buffers, ++ "descriptor": {}, ++ } ++ for slot in range(buffer_count): ++ buf = torch.empty( ++ (buffer_size,), dtype=torch.uint8, device=candidate_device ++ ) ++ if buf.is_cuda: ++ torch.cuda.synchronize(buf.device) ++ ret = engine.register_memory(int(buf.data_ptr()), int(buf.numel())) ++ if ret != 0: ++ raise RuntimeError(f"Mooncake register_memory failed: {ret}") ++ buffers.append(buf) ++ buffer_descs.append( ++ { ++ "slot": slot, ++ "addr": int(buf.data_ptr()), ++ "capacity": int(buf.numel()), ++ "device": buf.device.type, ++ } ++ ) ++ ++ rpc_port = int(engine.get_rpc_port()) ++ target_name = self._format_mooncake_target_name(base_server_name, rpc_port) ++ buffer_device = buffer_descs[0]["device"] if buffer_descs else candidate_device.type ++ descriptor = { ++ "target_name": target_name, ++ "server_name": base_server_name, ++ "rpc_port": rpc_port, ++ "buffer_device": buffer_device, ++ "buffers": buffer_descs, ++ } ++ self._mooncake_weight_receiver["descriptor"] = descriptor ++ return ( ++ True, ++ f"Succeeded to initialize Mooncake weight receiver on {buffer_device} buffer.", ++ descriptor, ++ ) ++ except Exception as e: ++ last_error = str(e) ++ logger.warning( ++ "Failed to initialize Mooncake weight receiver on %s: %s", ++ candidate_device, ++ e, ++ ) ++ cleanup_ok, cleanup_message = self.destroy_mooncake_weight_receiver() ++ if not cleanup_ok: ++ return False, cleanup_message, None ++ ++ return False, f"Failed to initialize Mooncake weight receiver: {last_error}", None ++ ++ def update_weights_from_mooncake( ++ self, ++ target_name: str, ++ buffer_slot: int, ++ remote_addr: int, ++ nbytes: int, ++ tensors: List[dict[str, object]], ++ load_format: Optional[str] = None, ++ delta: Optional[dict[str, object]] = None, ++ ) -> Tuple[bool, str]: ++ if delta is not None: ++ return False, "Mooncake weight receiver does not support delta manifests yet." ++ state = getattr(self, "_mooncake_weight_receiver", None) ++ if state is None: ++ return False, "Mooncake weight receiver is not initialized." ++ ++ descriptor = state["descriptor"] ++ if target_name != descriptor["target_name"]: ++ return False, f"Mooncake target mismatch: {target_name} != {descriptor['target_name']}" ++ buffers = state["buffers"] ++ if not (0 <= buffer_slot < len(buffers)): ++ return False, f"Invalid Mooncake buffer_slot={buffer_slot}." ++ buf = buffers[buffer_slot] ++ if int(remote_addr) != int(buf.data_ptr()): ++ return False, "Mooncake manifest remote_addr does not match receiver buffer." ++ if nbytes < 0 or nbytes > int(buf.numel()): ++ return False, "Mooncake manifest nbytes exceeds receiver buffer capacity." ++ ++ try: ++ if getattr(buf, "is_cuda", False) and str(self.device).startswith("cuda"): ++ torch.cuda.synchronize(torch.device(self.device)) ++ named_tensors = [] ++ for spec in tensors: ++ offset = int(spec["offset"]) ++ size = int(spec["nbytes"]) ++ if offset < 0 or size < 0 or offset + size > nbytes: ++ return False, f"Invalid Mooncake tensor slice: {spec!r}" ++ dtype = self._resolve_mooncake_torch_dtype(spec["dtype"]) ++ tensor = ( ++ buf.narrow(0, offset, size) ++ .view(dtype) ++ .view(tuple(spec["shape"])) ++ .clone() ++ ) ++ named_tensors.append((spec["name"], tensor)) ++ return self.update_weights_from_tensor(named_tensors, load_format) ++ except Exception as e: ++ logger.error("Mooncake weight update failed: %s", e) ++ return False, f"Mooncake weight update failed: {e}" ++ ++ def destroy_mooncake_weight_receiver(self) -> Tuple[bool, str]: ++ state = getattr(self, "_mooncake_weight_receiver", None) ++ if state is None: ++ return True, "Mooncake weight receiver is not initialized." ++ errors = [] ++ engine = state["engine"] ++ for buf in state["buffers"]: ++ ret = engine.unregister_memory(int(buf.data_ptr())) ++ if ret != 0: ++ errors.append(f"unregister_memory({int(buf.data_ptr())})={ret}") ++ free_engine = getattr(engine, "freeEngine", None) ++ if callable(free_engine): ++ ret = free_engine() ++ if ret not in (None, 0): ++ errors.append(f"freeEngine={ret}") ++ self._mooncake_weight_receiver = None ++ if errors: ++ return False, "; ".join(errors) ++ return True, "Succeeded to destroy Mooncake weight receiver." ++ + def post_process_weights(self, recv_req): + """ + Execute post-processing logic for model weights, such as Marlin quantization format conversion. @@ -2495,7 +2986,7 @@ index 275775a73d..e4e2fdc398 100644 def _model_load_weights_direct(model, named_tensors: List[Tuple[str, torch.Tensor]]): params_dict = dict(model.named_parameters()) diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py -index cc673a9cac..06c430d2c4 100644 +index cc673a9ca..06c430d2c 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py @@ -1,4 +1,5 @@ @@ -2514,7 +3005,7 @@ index cc673a9cac..06c430d2c4 100644 return AttnForwardMethod.MHA_ONE_SHOT return AttnForwardMethod.MLA diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py -index cb13a7c676..d9669ce086 100644 +index cb13a7c67..d9669ce08 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -29,6 +29,7 @@ from sglang.srt.layers.attention.nsa.utils import ( @@ -2545,7 +3036,7 @@ index cb13a7c676..d9669ce086 100644 if not forward_batch.forward_mode.is_idle(): if residual is not None: diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py -index 1583dd7880..a35c00f96c 100644 +index 1583dd788..a35c00f96 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1085,6 +1085,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): @@ -2766,7 +3257,7 @@ index 1583dd7880..a35c00f96c 100644 if normal_end_layer != self.end_layer: hidden_states, residual = model_forward_maybe_tbo( diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py -index db8c1c7ce7..53ffadf6d0 100644 +index db8c1c7ce..53ffadf6d 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -678,8 +678,13 @@ class Glm4MoeDecoderLayer(nn.Module): @@ -2794,7 +3285,7 @@ index db8c1c7ce7..53ffadf6d0 100644 hidden_states, residual = self.layer_communicator.prepare_attn( diff --git a/python/sglang/srt/models/glm4_moe_nextn.py b/python/sglang/srt/models/glm4_moe_nextn.py -index 1f6e753646..546cce4ab5 100644 +index 1f6e75364..546cce4ab 100644 --- a/python/sglang/srt/models/glm4_moe_nextn.py +++ b/python/sglang/srt/models/glm4_moe_nextn.py @@ -103,7 +103,7 @@ class Glm4MoeModelNextN(nn.Module): @@ -2807,7 +3298,7 @@ index 1f6e753646..546cce4ab5 100644 ) diff --git a/python/sglang/srt/models/glm4v_moe.py b/python/sglang/srt/models/glm4v_moe.py -index 324de18b49..fc72faa031 100644 +index 324de18b4..fc72faa03 100644 --- a/python/sglang/srt/models/glm4v_moe.py +++ b/python/sglang/srt/models/glm4v_moe.py @@ -52,11 +52,31 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): @@ -2912,7 +3403,7 @@ index 324de18b49..fc72faa031 100644 continue diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py -index f01225487b..1dad8bb8e5 100644 +index f01225487..1dad8bb8e 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -372,6 +372,7 @@ class Qwen3_5LinearDecoderLayer(nn.Module): @@ -2988,7 +3479,7 @@ index f01225487b..1dad8bb8e5 100644 return hidden_states, residual diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py -index d641826e33..3abc39ef32 100644 +index d641826e3..3abc39ef3 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -711,14 +711,19 @@ class Qwen3LLMModel(Qwen3Model): @@ -3016,7 +3507,7 @@ index d641826e33..3abc39ef32 100644 positions, hidden_states, diff --git a/python/sglang/srt/multimodal/processors/glm4v.py b/python/sglang/srt/multimodal/processors/glm4v.py -index 33cce6fe25..0970c4550d 100644 +index 33cce6fe2..0970c4550 100644 --- a/python/sglang/srt/multimodal/processors/glm4v.py +++ b/python/sglang/srt/multimodal/processors/glm4v.py @@ -1,6 +1,9 @@ @@ -3076,7 +3567,7 @@ index 33cce6fe25..0970c4550d 100644 self, image_data: List[Union[str, bytes]], diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py -index 4395654e4e..f9b5ea4abb 100644 +index 4395654e4..f9b5ea4ab 100644 --- a/python/sglang/srt/multimodal/processors/qwen_vl.py +++ b/python/sglang/srt/multimodal/processors/qwen_vl.py @@ -317,7 +317,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor): @@ -3089,7 +3580,7 @@ index 4395654e4e..f9b5ea4abb 100644 image_data=image_data, video_data=request_obj.video_data, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py -index b080aeb168..5b29ebf566 100644 +index b080aeb16..5b29ebf56 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -635,6 +635,7 @@ class ServerArgs: @@ -3114,7 +3605,7 @@ index b080aeb168..5b29ebf566 100644 "--enable-fused-qk-norm-rope", action="store_true", diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py -index 5fe45086ca..b283d2e9bd 100644 +index 5fe45086c..b283d2e9b 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -341,7 +341,10 @@ class EAGLEDraftCudaGraphRunner: @@ -3145,7 +3636,7 @@ index 5fe45086ca..b283d2e9bd 100644 self.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices) diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py -index ac629c7ee5..c039d23508 100644 +index ac629c7ee..c039d2350 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -774,6 +774,10 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): @@ -3188,7 +3679,7 @@ index ac629c7ee5..c039d23508 100644 @dataclass diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py -index 4636128fa7..a9b61df393 100644 +index 4636128fa..a9b61df39 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -2359,6 +2359,8 @@ class SafeUnpickler(pickle.Unpickler): @@ -3201,7 +3692,7 @@ index 4636128fa7..a9b61df393 100644 DENY_CLASSES = { diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py -index 3be16446e0..1b2371c839 100644 +index 3be16446e..1b2371c83 100644 --- a/python/sglang/srt/utils/weight_checker.py +++ b/python/sglang/srt/utils/weight_checker.py @@ -69,6 +69,9 @@ def _check_tensors( diff --git a/docs/en/advanced/external-rollout-engines.md b/docs/en/advanced/external-rollout-engines.md index 498afa0c5f..7712987595 100644 --- a/docs/en/advanced/external-rollout-engines.md +++ b/docs/en/advanced/external-rollout-engines.md @@ -13,6 +13,7 @@ This page is a roadmap. Use it to decide when to use `--rollout-external-engine- | Trainer and external engines can form an NCCL group | Default `--update-weight-mode full --update-weight-transport nccl` | | Trainer and external engines cannot form an NCCL group, but can see the same filesystem path | `--update-weight-mode full --update-weight-transport disk` | | Full checkpoints are too heavy for large-model cross-cluster or cross-DC sync | `--update-weight-mode delta --update-weight-transport disk` | +| slime owns the SGLang engines and you want Mooncake TransferEngine full-weight sync | [Mooncake TransferEngine Weight Sync](mooncake-weight-sync.md) | | Rollout serving can use an independent SGLang environment, or even different GPU models/vendors | external engines + disk transport | | You want to validate delta wire/apply logic inside one datacenter | `--update-weight-mode delta --update-weight-transport nccl` | | You need frozen reference, reward, or tool-side models | Prefer `update_weights: false` in [SGLang Config](sglang-config.md#3-multi-model-serving) | @@ -110,6 +111,7 @@ For encoding choices, wire layout, receiver-side selective overwrite, and tuning - Disk transport requires trainer and SGLang engines to see the same `--update-weight-disk-dir` path; a path visible only to the trainer is not enough. - External engines are not recovered by slime fault tolerance; their lifecycle belongs to the external deployment system. - `--sglang-config` and `--rollout-external-engine-addrs` are mutually exclusive. +- Mooncake TransferEngine weight sync is only for slime-managed engines today; it is rejected with `--rollout-external-engine-addrs`. - Delta mode does not support `--colocate`, because colocated sync uses CUDA IPC handles and delta encoding does not reduce the actual transfer. ## Related Work diff --git a/docs/en/advanced/mooncake-weight-sync.md b/docs/en/advanced/mooncake-weight-sync.md new file mode 100644 index 0000000000..0f8a0bc61e --- /dev/null +++ b/docs/en/advanced/mooncake-weight-sync.md @@ -0,0 +1,87 @@ +# Mooncake TransferEngine Weight Sync + +Mooncake TransferEngine transport is an experimental full-weight sync path for +slime-managed SGLang engines. Instead of building an NCCL update group or +writing a checkpoint to a shared filesystem, slime asks each SGLang engine to +allocate a local Mooncake receiver buffer, writes each weight bucket into that +buffer with Mooncake TransferEngine, and then asks SGLang to load the bucket from +its local buffer. + +Use this only when slime owns the SGLang engine lifecycle. If your engines are +pre-launched by another system through `--rollout-external-engine-addrs`, use +NCCL or disk transport instead. + +## Quick Start + +```bash +--update-weight-mode full +--update-weight-transport mooncake +--mooncake-metadata-server P2PHANDSHAKE +--mooncake-protocol tcp +--mooncake-buffer-size $((8 * 1024 * 1024 * 1024)) +--mooncake-buffer-count 1 +``` + +`P2PHANDSHAKE` is the only supported metadata mode today. The SGLang receiver +returns a per-engine target name during initialization, so slime does not require +an external Mooncake metadata service for this path. + +For RDMA deployments, switch the protocol and pass the device name used by the +Mooncake runtime: + +```bash +--mooncake-protocol rdma +--mooncake-device-name mlx5_0 +``` + +## Current Limits + +Mooncake transport intentionally starts with a narrow contract: + +- It supports only `--update-weight-mode full`; delta mode is rejected. +- It supports only slime-managed SGLang engines; external rollout engines are + rejected. +- It does not support `--colocate`; colocated sync uses CUDA IPC tensor handles. +- It currently supports one GPU per rollout engine. +- Megatron pipeline model parallel size must be 1. +- SGLang DP size and EP size must be 1. +- SGLang DP attention is not supported. +- `--mooncake-metadata-server` must stay `P2PHANDSHAKE`. +- `--mooncake-rpc-port-base` is reserved and rejected by argument validation. +- `--mooncake-buffer-count` is reserved for future double buffering and must be + `1`; the sender uses slot 0 for every update. + +For `--sglang-config`, these limits are checked only for server groups that +receive training weights (`update_weights: true`). Frozen reference/reward +models and placeholder groups do not need Mooncake receivers. + +## Buffer Sizing + +`--mooncake-buffer-size` must be large enough for the largest weight bucket sent +by slime. If it is unset, slime uses `--update-weight-buffer-size`. + +If a bucket is larger than the receiver buffer, slime fails the sync before +writing into Mooncake. Increase either `--mooncake-buffer-size` or reduce +`--update-weight-buffer-size` so every bucket fits the receiver buffer. + +`--mooncake-buffer-count` is present only to keep the receiver API shape ready +for future double buffering. It does not increase concurrency in the current +implementation, and values other than `1` are rejected. + +## Relationship With Other Transports + +`--update-weight-mode` decides what gets sent, and `--update-weight-transport` +decides how it reaches SGLang: + +| mode | transport | behavior | +|---|---|---| +| `full` | `nccl` | broadcast every HF weight chunk over a trainer-engine NCCL group | +| `full` | `disk` | write a complete HF checkpoint, then call `update_weights_from_disk` | +| `full` | `mooncake` | write each bucket into SGLang receiver buffers through Mooncake TransferEngine | +| `delta` | `nccl` | broadcast sparse changed positions and values over NCCL | +| `delta` | `disk` | write sparse safetensors, then call `update_weights_from_disk(load_format="delta")` | + +Use Mooncake transport only for the `full` + slime-managed SGLang case above. +For external engines, heterogeneous GPU fleets, or cross-cluster deployments +with a shared filesystem, see [External Rollout Engines](external-rollout-engines.md) +and [Delta Weight Sync](delta-weight-sync.md). diff --git a/docs/en/get_started/usage.md b/docs/en/get_started/usage.md index 5deeda7980..90ac695b56 100644 --- a/docs/en/get_started/usage.md +++ b/docs/en/get_started/usage.md @@ -147,6 +147,7 @@ Note: - By default, SGLang reads the maximum context length from the `config.json` in the Hugging Face checkpoint. You can use the `--sglang-context-length` parameter to override this value to support longer inference. - During co-located training and inference, although Megatron and SGLang will offload sequentially, they still need to leave some memory for each other. You need to adjust SGLang's total VRAM usage by reducing `--sglang-mem-fraction-static`. - slime supports passing through sgl-router parameters by adding a `router` prefix to the original parameter name. For example, sgl-router's `--balance-abs-threshold` parameter should be set as `--router-balance-abs-threshold`. Since sgl-router uses cache-aware routing by default, it may cause uneven request distribution. You can set `--router-balance-abs-threshold 0` to force balanced distribution, but this may affect prefix cache hit rate in multi-turn conversation scenarios. + - For slime-managed SGLang engines, full-weight sync can also use Mooncake TransferEngine with `--update-weight-mode full --update-weight-transport mooncake`. This path currently requires `--mooncake-metadata-server P2PHANDSHAKE`, one GPU per rollout engine, no `--colocate`, no external rollout engines, and `--mooncake-buffer-size` large enough for the largest weight bucket. See [Mooncake TransferEngine Weight Sync](../advanced/mooncake-weight-sync.md). - If SGLang engines are pre-launched by an external system, connect to them with `--rollout-external-engine-addrs host1:port host2:port`. When the trainer and engines cannot form an NCCL weight-update group, use `--update-weight-mode full --update-weight-transport disk --update-weight-disk-dir /shared/fs/updates`; slime writes a complete HF checkpoint and asks SGLang to hot-load it through `update_weights_from_disk`. For large models or cross-cluster deployments, use `--update-weight-mode delta --update-weight-transport disk` instead. See [External Rollout Engines Roadmap](../advanced/external-rollout-engines.md) and [Delta Weight Sync](../advanced/delta-weight-sync.md). For details on some of SGLang's customizations and the principles behind how slime incorporates SGLang, please see the "How to Use SGLang" section. diff --git a/docs/en/index.rst b/docs/en/index.rst index d066a36837..aceb6b8934 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -41,6 +41,7 @@ Start by Use Case - Use PD disaggregation: :doc:`advanced/pd-disaggregation` - Use BF16 training with FP8 rollout or FP8 KV cache: :doc:`advanced/low-precision` - Use delta weight sync: :doc:`advanced/delta-weight-sync` +- Use Mooncake TransferEngine weight sync: :doc:`advanced/mooncake-weight-sync` - Understand CI and reliability coverage: :doc:`developer_guide/ci` - Debug, trace, and profile long-running jobs: :doc:`developer_guide/debug`, :doc:`developer_guide/trace`, :doc:`developer_guide/profiling` @@ -85,6 +86,7 @@ Start by Use Case advanced/pd-disaggregation.md advanced/external-rollout-engines.md advanced/delta-weight-sync.md + advanced/mooncake-weight-sync.md advanced/sglang-config.md advanced/megatron-config.md advanced/arch-support-beyond-megatron.md diff --git a/docs/zh/advanced/external-rollout-engines.md b/docs/zh/advanced/external-rollout-engines.md index 9aae0ef5ec..a111166155 100644 --- a/docs/zh/advanced/external-rollout-engines.md +++ b/docs/zh/advanced/external-rollout-engines.md @@ -13,6 +13,7 @@ External rollout engine 指的是：SGLang engine 不由 slime 训练任务启 | 训练器和 external engine 可以建立 NCCL group | 默认的 `--update-weight-mode full --update-weight-transport nccl` | | 训练器和 external engine 不能建立 NCCL group，但能共享同一路径的文件系统 | `--update-weight-mode full --update-weight-transport disk` | | 大模型跨集群或跨数据中心同步，full checkpoint 太重 | `--update-weight-mode delta --update-weight-transport disk` | +| slime 负责启动 SGLang engine，并希望使用 Mooncake TransferEngine 做全量权重同步 | [Mooncake TransferEngine 权重同步](mooncake-weight-sync.md) | | rollout serving 想使用独立 SGLang 环境，甚至不同型号或不同厂家的 GPU | external engine + disk transport | | 想验证 delta wire/apply 逻辑，但仍在同一数据中心内 | `--update-weight-mode delta --update-weight-transport nccl` | | 需要 reference、reward、tool-side model 等冻结模型 | 优先用 [SGLang Config](sglang-config.md#3-多模型服务) 的 `update_weights: false` | @@ -110,6 +111,7 @@ delta update 面向大模型、跨集群或跨数据中心训推解耦。它不 - disk transport 要求训练端和 SGLang engine 看到同一个 `--update-weight-disk-dir` 路径；路径只在训练端可见是不够的。 - external engine 当前不支持 slime 的 fault tolerance 恢复流程；engine 生命周期由外部系统负责。 - `--sglang-config` 与 `--rollout-external-engine-addrs` 互斥。 +- Mooncake TransferEngine 权重同步当前只支持 slime-managed engine；与 `--rollout-external-engine-addrs` 同用会被拒绝。 - delta mode 不支持 `--colocate`，因为 colocated 权重同步通过 CUDA IPC 传句柄，delta 编码不会节省实际传输量。 ## 参考工作 diff --git a/docs/zh/advanced/mooncake-weight-sync.md b/docs/zh/advanced/mooncake-weight-sync.md new file mode 100644 index 0000000000..3a7b7a6356 --- /dev/null +++ b/docs/zh/advanced/mooncake-weight-sync.md @@ -0,0 +1,81 @@ +# Mooncake TransferEngine 权重同步 + +Mooncake TransferEngine transport 是面向 slime-managed SGLang engine 的实验性 +全量权重同步路径。它不建立 NCCL 权重更新组，也不把 checkpoint 写到共享文件系统； +slime 会先让每个 SGLang engine 分配本地 Mooncake receiver buffer，再通过 Mooncake +TransferEngine 把每个权重 bucket 写入 receiver buffer，最后通知 SGLang 从本地 buffer +加载权重。 + +这条路径只适用于 slime 负责 SGLang engine 生命周期的场景。如果 engine 通过 +`--rollout-external-engine-addrs` 由外部系统预启动，请继续使用 NCCL 或 disk transport。 + +## 快速开始 + +```bash +--update-weight-mode full +--update-weight-transport mooncake +--mooncake-metadata-server P2PHANDSHAKE +--mooncake-protocol tcp +--mooncake-buffer-size $((8 * 1024 * 1024 * 1024)) +--mooncake-buffer-count 1 +``` + +当前只支持 `P2PHANDSHAKE` metadata 模式。SGLang receiver 初始化时会返回每个 +engine 的 target name，因此这条路径不依赖外部 Mooncake metadata service。 + +如果使用 RDMA，把 protocol 切到 `rdma`，并传入 Mooncake runtime 使用的设备名： + +```bash +--mooncake-protocol rdma +--mooncake-device-name mlx5_0 +``` + +## 当前限制 + +Mooncake transport 当前刻意收敛在一组窄契约内： + +- 只支持 `--update-weight-mode full`；delta mode 会被拒绝。 +- 只支持 slime-managed SGLang engine；external rollout engine 会被拒绝。 +- 不支持 `--colocate`；colocated 权重同步使用 CUDA IPC tensor handle。 +- 当前只支持每个 rollout engine 1 张 GPU。 +- Megatron pipeline model parallel size 必须为 1。 +- SGLang DP size 和 EP size 必须为 1。 +- 不支持 SGLang DP attention。 +- `--mooncake-metadata-server` 必须保持 `P2PHANDSHAKE`。 +- `--mooncake-rpc-port-base` 是预留参数，参数校验会拒绝设置它。 +- `--mooncake-buffer-count` 是为未来 double buffering 预留的参数，当前必须是 + `1`；发送端每次更新都使用 slot 0。 + +如果使用 `--sglang-config`，这些限制只会检查接收训练权重的 server group +（`update_weights: true`）。冻结的 reference/reward model 和 placeholder group +不需要 Mooncake receiver。 + +## Buffer 大小 + +`--mooncake-buffer-size` 必须能容纳 slime 发送的最大权重 bucket。如果不设置， +slime 会使用 `--update-weight-buffer-size`。 + +如果某个 bucket 大于 receiver buffer，slime 会在写入 Mooncake 前让本次同步失败。 +可以增大 `--mooncake-buffer-size`，或减小 `--update-weight-buffer-size`，保证每个 +bucket 都能放进 receiver buffer。 + +`--mooncake-buffer-count` 只是为了让 receiver API 保持未来 double buffering 的形状。 +当前实现不会因为设置多个 buffer 而提升并发，且参数值不是 `1` 时会被拒绝。 + +## 与其他传输方式的关系 + +`--update-weight-mode` 决定发送什么，`--update-weight-transport` 决定如何送到 +SGLang： + +| mode | transport | 行为 | +|---|---|---| +| `full` | `nccl` | 通过 trainer-engine NCCL group 广播所有 HF 权重 chunk | +| `full` | `disk` | 写完整 HF checkpoint，然后调用 `update_weights_from_disk` | +| `full` | `mooncake` | 通过 Mooncake TransferEngine 把每个 bucket 写入 SGLang receiver buffer | +| `delta` | `nccl` | 通过 NCCL 广播稀疏变化位置和值 | +| `delta` | `disk` | 写稀疏 safetensors，然后调用 `update_weights_from_disk(load_format="delta")` | + +Mooncake transport 只用于上表中的 `full` + slime-managed SGLang 场景。External +engine、异构 GPU 集群或共享文件系统跨集群部署，见 +[External Rollout Engines](external-rollout-engines.md) 和 +[Delta 权重同步](delta-weight-sync.md)。 diff --git a/docs/zh/get_started/usage.md b/docs/zh/get_started/usage.md index a501f94634..5d1511dd2d 100644 --- a/docs/zh/get_started/usage.md +++ b/docs/zh/get_started/usage.md @@ -151,6 +151,7 @@ sglang 的加载非常简单，只需要： - sglang 默认会从 huggingface ckpt 中 `config.json` 读取模型的最大 context length，可以使用 `--sglang-context-length` 参数来对这个值进行覆盖，从而支持进行更长的推理； - 在训推一体的训练过程中，虽然 megatron 和 sglang 会先后 offload，但是还是需要为对方留有一些空间，需要通过减小 `--sglang-mem-fraction-static` 来调整 sglang 的显存占用总量。 - slime 支持透传 sgl-router 的参数，方式是在原参数名前加上 `router` 前缀。例如，sgl-router 的 `--balance-abs-threshold` 参数需要设置为 `--router-balance-abs-threshold`。由于 sgl-router 默认使用 cache-aware routing，可能会导致请求分配不均衡的问题。可以通过设置 `--router-balance-abs-threshold 0` 来强制均衡分配，但这可能会影响多轮对话场景下 prefix cache 的命中率。 +- 对于 slime-managed SGLang engine，全量权重同步也可以使用 Mooncake TransferEngine：`--update-weight-mode full --update-weight-transport mooncake`。这条路径当前要求 `--mooncake-metadata-server P2PHANDSHAKE`、每个 rollout engine 1 张 GPU、不使用 `--colocate`、不使用 external rollout engine，并且 `--mooncake-buffer-size` 要能容纳最大的权重 bucket。详见 [Mooncake TransferEngine 权重同步](../advanced/mooncake-weight-sync.md)。 - 如果 SGLang engine 已经由外部系统预启动，可以通过 `--rollout-external-engine-addrs host1:port host2:port` 连接。此时如果训练器和 engine 无法建立 NCCL 权重同步 group，可以使用 `--update-weight-mode full --update-weight-transport disk --update-weight-disk-dir /shared/fs/updates`，slime 会写完整 HF checkpoint 并调用 SGLang 的 `update_weights_from_disk` 热加载；大模型或跨集群场景可进一步使用 `--update-weight-mode delta --update-weight-transport disk`。详见 [External Rollout Engines 配置路线图](../advanced/external-rollout-engines.md) 和 [Delta 权重同步](../advanced/delta-weight-sync.md)。 对于一些 sglang 的自定义以及 slime 引入 sglang 的原理，请见 sglang 使用方法一节。 diff --git a/docs/zh/index.rst b/docs/zh/index.rst index c062cde2a2..d39800ee00 100644 --- a/docs/zh/index.rst +++ b/docs/zh/index.rst @@ -41,6 +41,7 @@ slime 的设计目标，是让这两大能力彼此强化，同时避免把系 - 使用 PD disaggregation：:doc:`advanced/pd-disaggregation` - 使用 BF16 训练 + FP8 rollout 或 FP8 KV cache：:doc:`advanced/low-precision` - 使用 delta weight sync：:doc:`advanced/delta-weight-sync` +- 使用 Mooncake TransferEngine 权重同步：:doc:`advanced/mooncake-weight-sync` - 了解 CI 和可靠性覆盖：:doc:`developer_guide/ci` - 调试、trace 和 profiling 长时间任务：:doc:`developer_guide/debug`、:doc:`developer_guide/trace`、:doc:`developer_guide/profiling` @@ -85,6 +86,7 @@ slime 的设计目标，是让这两大能力彼此强化，同时避免把系 advanced/pd-disaggregation.md advanced/external-rollout-engines.md advanced/delta-weight-sync.md + advanced/mooncake-weight-sync.md advanced/sglang-config.md advanced/megatron-config.md advanced/arch-support-beyond-megatron.md diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 6bdc2dd3fe..fb92df10fd 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -36,6 +36,7 @@ from .update_weight.common import named_params_and_buffers from .update_weight.update_weight_from_disk import UpdateWeightFromDisk from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed +from .update_weight.update_weight_from_mooncake import UpdateWeightFromMooncake from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor logging.getLogger("megatron").setLevel(logging.WARNING) @@ -152,6 +153,8 @@ def init( assert self.args.update_weight_mode == "full" if self.args.update_weight_transport == "disk": update_weight_cls = UpdateWeightFromDisk + elif self.args.update_weight_transport == "mooncake": + update_weight_cls = UpdateWeightFromMooncake else: assert ( self.args.update_weight_mode == "full" and self.args.update_weight_transport == "nccl" diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_mooncake.py b/slime/backends/megatron_utils/update_weight/update_weight_from_mooncake.py new file mode 100644 index 0000000000..0155c21668 --- /dev/null +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_mooncake.py @@ -0,0 +1,502 @@ +from __future__ import annotations + +import logging +import socket +import sys +import time +from argparse import Namespace +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass +from typing import Any + +logger = logging.getLogger(__name__) + +try: + import ray + from megatron.core import mpu + + from .update_weight_from_distributed import UpdateWeightFromDistributed +except ImportError: # pragma: no cover - keeps pure helpers importable in minimal unit environments. + ray = None + UpdateWeightFromDistributed = object + + class _MissingMPU: + def get_data_parallel_rank(self, with_context_parallel: bool = True) -> int: + raise RuntimeError("megatron.core.mpu is required for Mooncake weight updates.") + + def get_tensor_model_parallel_rank(self) -> int: + raise RuntimeError("megatron.core.mpu is required for Mooncake weight updates.") + + def get_pipeline_model_parallel_rank(self) -> int: + raise RuntimeError("megatron.core.mpu is required for Mooncake weight updates.") + + mpu = _MissingMPU() + + +def _tensor_nbytes(tensor: Any) -> int: + return int(tensor.numel() * tensor.element_size()) + + +def _dtype_name(dtype: Any) -> str: + return str(dtype).replace("torch.", "") + + +def _find_buffer(descriptor: Mapping[str, Any], buffer_slot: int) -> Mapping[str, Any]: + for buffer in descriptor.get("buffers", []): + if int(buffer.get("slot", 0)) == buffer_slot: + return buffer + raise ValueError(f"Mooncake receiver descriptor has no buffer slot {buffer_slot}: {descriptor!r}") + + +def _format_target_name(server_name: str, rpc_port: int | str) -> str: + if ":" in server_name and not server_name.startswith("["): + return f"[{server_name}]:{rpc_port}" + return f"{server_name}:{rpc_port}" + + +def _target_name_from_descriptor(descriptor: Mapping[str, Any]) -> str: + target_name = descriptor.get("target_name") + if target_name: + return str(target_name) + server_name = descriptor.get("server_name") + rpc_port = descriptor.get("rpc_port") + if server_name and rpc_port is not None: + return _format_target_name(str(server_name), rpc_port) + raise ValueError( + "Mooncake receiver descriptor must include target_name, or both server_name and rpc_port: " f"{descriptor!r}" + ) + + +def _serialize_delta(delta: Any) -> dict[str, Any]: + return { + "encoding": delta.encoding.value, + "params": [asdict(param) for param in delta.params], + "checksum": delta.checksum, + } + + +def _validate_single_gpu_rollout_engines( + *, + args: Namespace, + engine_gpu_counts: Sequence[int] | None, + rollout_engine_count: int, +) -> None: + if engine_gpu_counts is None: + engine_gpu_counts = [getattr(args, "rollout_num_gpus_per_engine", 1)] * rollout_engine_count + if any(int(count) != 1 for count in engine_gpu_counts): + raise NotImplementedError( + "--update-weight-transport=mooncake currently supports only one GPU per rollout engine. " + "Multi-TP rollout engines need one Mooncake receiver descriptor per TP rank." + ) + + +def default_mooncake_local_server_name( + *, + metadata_server: str, + pp_rank: int, + host_getter=socket.gethostname, +) -> str: + if metadata_server == "P2PHANDSHAKE": + hostname = host_getter() + try: + return socket.gethostbyname(hostname) + except socket.gaierror: + return hostname + return f"slime-train-pp{pp_rank}" + + +def build_mooncake_manifest( + *, + descriptor: Mapping[str, Any], + named_tensors: Sequence[tuple[str, Any]], + weight_version: int, + load_format: str | None, + buffer_slot: int = 0, + delta: Any | None = None, +) -> dict[str, Any]: + buffer = _find_buffer(descriptor, buffer_slot) + capacity = int(buffer["capacity"]) + remote_addr = int(buffer["addr"]) + buffer_device = str(buffer.get("device", descriptor.get("buffer_device", "cuda"))) + target_name = _target_name_from_descriptor(descriptor) + + offset = 0 + entries = [] + for name, tensor in named_tensors: + nbytes = _tensor_nbytes(tensor) + entries.append( + { + "name": name, + "dtype": _dtype_name(tensor.dtype), + "shape": list(tensor.shape), + "offset": offset, + "nbytes": nbytes, + } + ) + offset += nbytes + + if offset > capacity: + raise ValueError( + f"Weight bucket of {offset} bytes exceeds Mooncake receiver buffer slot " + f"{buffer_slot} capacity {capacity} bytes for {target_name}." + ) + + manifest: dict[str, Any] = { + "weight_version": str(weight_version), + "target_name": target_name, + "buffer_slot": buffer_slot, + "remote_addr": remote_addr, + "buffer_device": buffer_device, + "nbytes": offset, + "tensors": entries, + } + if load_format is not None: + manifest["load_format"] = load_format + if delta is not None: + manifest["delta"] = _serialize_delta(delta) + return manifest + + +@dataclass(frozen=True) +class MooncakeTransferCall: + target_name: str + remote_addr: int + nbytes: int + tensor_names: tuple[str, ...] + + +class FakeMooncakeTransferClient: + def __init__(self) -> None: + self.calls: list[MooncakeTransferCall] = [] + + def write_bucket( + self, + *, + receiver_descriptors: Sequence[Mapping[str, Any]], + named_tensors: Sequence[tuple[str, Any]], + weight_version: int, + load_format: str | None, + delta: Any | None, + buffer_slot: int = 0, + ) -> list[dict[str, Any]]: + manifests = [] + for descriptor in receiver_descriptors: + manifest = build_mooncake_manifest( + descriptor=descriptor, + named_tensors=named_tensors, + weight_version=weight_version, + load_format=load_format, + buffer_slot=buffer_slot, + delta=delta, + ) + self.calls.append( + MooncakeTransferCall( + target_name=manifest["target_name"], + remote_addr=manifest["remote_addr"], + nbytes=manifest["nbytes"], + tensor_names=tuple(entry["name"] for entry in manifest["tensors"]), + ) + ) + manifests.append(manifest) + return manifests + + def close(self) -> None: + return None + + +class MooncakeTransferClient: + def __init__( + self, + *, + local_server_name: str, + metadata_server: str, + protocol: str, + device_name: str, + ) -> None: + from mooncake.engine import TransferEngine + + self.engine = TransferEngine() + ret = self.engine.initialize(local_server_name, metadata_server, protocol, device_name) + if ret != 0: + raise RuntimeError(f"Mooncake TransferEngine.initialize failed with code {ret}") + self.local_server_name = local_server_name + + @classmethod + def from_args(cls, args: Namespace, *, pp_rank: int) -> MooncakeTransferClient: + return cls( + local_server_name=default_mooncake_local_server_name( + metadata_server=args.mooncake_metadata_server, + pp_rank=pp_rank, + ), + metadata_server=args.mooncake_metadata_server, + protocol=args.mooncake_protocol, + device_name=args.mooncake_device_name or "", + ) + + def write_bucket( + self, + *, + receiver_descriptors: Sequence[Mapping[str, Any]], + named_tensors: Sequence[tuple[str, Any]], + weight_version: int, + load_format: str | None, + delta: Any | None, + buffer_slot: int = 0, + ) -> list[dict[str, Any]]: + manifests = [ + build_mooncake_manifest( + descriptor=descriptor, + named_tensors=named_tensors, + weight_version=weight_version, + load_format=load_format, + buffer_slot=buffer_slot, + delta=delta, + ) + for descriptor in receiver_descriptors + ] + buffer_devices = {str(manifest.get("buffer_device", "cuda")) for manifest in manifests} + if buffer_devices not in ({"cuda"}, set()): + raise RuntimeError(f"Mooncake weight transport requires CUDA receiver buffers: {buffer_devices}") + + registered_ptrs: list[int] = [] + try: + for _name, tensor in named_tensors: + if not tensor.is_cuda: + raise TypeError("Mooncake GPU-direct weight transport requires CUDA tensors.") + if not tensor.is_contiguous(): + raise ValueError("Mooncake weight transport requires contiguous tensors.") + ptr = int(tensor.data_ptr()) + ret = self.engine.register_memory(ptr, _tensor_nbytes(tensor)) + if ret != 0: + raise RuntimeError(f"Mooncake register_memory failed with code {ret}") + registered_ptrs.append(ptr) + + for manifest in manifests: + self._write_manifest(manifest, named_tensors) + return manifests + finally: + unregister_errors = [] + for ptr in registered_ptrs: + ret = self.engine.unregister_memory(ptr) + if ret != 0: + unregister_errors.append(f"unregister_memory({ptr})={ret}") + if unregister_errors: + message = "Mooncake unregister_memory failed: " + "; ".join(unregister_errors) + if sys.exc_info()[0] is None: + raise RuntimeError(message) + logger.warning(message) + + def _write_manifest( + self, + manifest: Mapping[str, Any], + named_tensors: Sequence[tuple[str, Any]], + ) -> None: + target_name = manifest["target_name"] + remote_base = int(manifest["remote_addr"]) + tensor_by_name = dict(named_tensors) + import torch + + stream_by_device: dict[Any, Any] = {} + + for entry in manifest["tensors"]: + tensor = tensor_by_name[entry["name"]] + device = tensor.device + stream = stream_by_device.get(device) + if stream is None: + stream = torch.cuda.Stream(device=device) + stream_by_device[device] = stream + + current_stream = torch.cuda.current_stream(device) + event = current_stream.record_event() + stream.wait_event(event) + ret = self.engine.transfer_write_on_cuda( + target_name, + int(tensor.data_ptr()), + remote_base + int(entry["offset"]), + int(entry["nbytes"]), + stream.cuda_stream, + ) + if ret != 0: + raise RuntimeError(f"Mooncake transfer_write_on_cuda failed with code {ret}") + + for stream in stream_by_device.values(): + stream.synchronize() + + def close(self) -> None: + free_engine = getattr(self.engine, "freeEngine", None) + if callable(free_engine): + free_engine() + + +class UpdateWeightFromMooncake(UpdateWeightFromDistributed): + def __init__( + self, + args: Namespace, + model: Sequence[Any], + weights_getter, + *, + model_name: str, + quantization_config: dict[str, int | str | list[str]] | None, + transfer_client: MooncakeTransferClient | FakeMooncakeTransferClient | None = None, + ) -> None: + if UpdateWeightFromDistributed is object: + self.args = args + self.model = model + self.model_name = model_name + self.quantization_config = quantization_config + self.weight_version = 0 + self.update_weight_metrics = {} + else: + super().__init__( + args, + model, + weights_getter, + model_name=model_name, + quantization_config=quantization_config, + ) + self._transfer_client = transfer_client + self._receiver_descriptors: list[Mapping[str, Any]] = [] + self._buffer_slot = 0 + self._connected = False + + def _destroy_mooncake_receivers(self, rollout_engines: Sequence[Any] | None = None) -> Exception | None: + if ray is None: + return None + engines = rollout_engines if rollout_engines is not None else getattr(self, "rollout_engines", None) + if not getattr(self, "_is_pp_src_rank", False) or not engines: + return None + try: + ray.get([engine.destroy_mooncake_weight_receiver.remote() for engine in engines]) + except Exception as exc: + return exc + return None + + def _close_transfer_client(self) -> Exception | None: + if self._transfer_client is None: + return None + try: + self._transfer_client.close() + except Exception as exc: + return exc + finally: + self._transfer_client = None + return None + + def _reset_mooncake_state(self, *, raise_destroy_error: bool) -> None: + destroy_error = self._destroy_mooncake_receivers() + close_error = self._close_transfer_client() + self._receiver_descriptors = [] + self._connected = False + if raise_destroy_error and destroy_error is not None: + raise destroy_error + if raise_destroy_error and close_error is not None: + raise close_error + + def connect_rollout_engines( + self, + rollout_engines: Sequence[Any], + rollout_engine_lock: Any, + engine_gpu_counts: Sequence[int] | None = None, + engine_gpu_offsets: Sequence[int] | None = None, + ) -> None: + self.rollout_engines = rollout_engines + self.rollout_engine_lock = rollout_engine_lock + self._engine_gpu_counts = engine_gpu_counts + self._is_pp_src_rank = ( + mpu.get_data_parallel_rank(with_context_parallel=True) == 0 and mpu.get_tensor_model_parallel_rank() == 0 + ) + pp_rank = mpu.get_pipeline_model_parallel_rank() + self._group_name = f"slime-mooncake-pp_{pp_rank}" + _validate_single_gpu_rollout_engines( + args=self.args, + engine_gpu_counts=engine_gpu_counts, + rollout_engine_count=len(rollout_engines), + ) + + if not self._is_pp_src_rank: + return + + if self._connected: + self._reset_mooncake_state(raise_destroy_error=False) + + config = { + "metadata_server": self.args.mooncake_metadata_server, + "protocol": self.args.mooncake_protocol, + "device_name": self.args.mooncake_device_name or "", + "buffer_size": getattr(self.args, "mooncake_buffer_size", None) or self.args.update_weight_buffer_size, + "buffer_count": self.args.mooncake_buffer_count, + "rpc_port_base": getattr(self.args, "mooncake_rpc_port_base", None), + "group_name": self._group_name, + } + try: + self._receiver_descriptors = ray.get( + [engine.init_mooncake_weight_receiver.remote(config) for engine in rollout_engines] + ) + if self._transfer_client is None: + self._transfer_client = MooncakeTransferClient.from_args(self.args, pp_rank=pp_rank) + self._connected = True + except Exception: + self._destroy_mooncake_receivers(rollout_engines) + self._close_transfer_client() + self._receiver_descriptors = [] + raise + + def disconnect_rollout_engines(self) -> None: + self._reset_mooncake_state(raise_destroy_error=True) + + def _update_bucket_weights_from_distributed( + self, + converted_named_tensors: list[tuple[str, Any]], + pbar=None, + load_format: str | None = None, + delta: Any | None = None, + ) -> None: + self._update_bucket_weights_from_mooncake( + converted_named_tensors, + pbar=pbar, + load_format=load_format, + delta=delta, + ) + + def _update_bucket_weights_from_mooncake( + self, + converted_named_tensors: list[tuple[str, Any]], + pbar=None, + load_format: str | None = None, + delta: Any | None = None, + ) -> None: + if not getattr(self, "_is_pp_src_rank", False) or not converted_named_tensors: + converted_named_tensors.clear() + return + if self._transfer_client is None: + raise RuntimeError("Mooncake transfer client is not initialized.") + + while not ray.get(self.rollout_engine_lock.acquire.remote()): + time.sleep(0.1) + + try: + manifests = self._transfer_client.write_bucket( + receiver_descriptors=self._receiver_descriptors, + named_tensors=converted_named_tensors, + weight_version=self.weight_version, + load_format=load_format, + delta=delta, + buffer_slot=self._buffer_slot, + ) + refs = [ + engine.update_weights_from_mooncake.remote(manifest) + for engine, manifest in zip(self.rollout_engines, manifests, strict=True) + ] + ray.get(refs) + total_bytes = sum(manifest["nbytes"] for manifest in manifests) + self.update_weight_metrics["mooncake/write_bucket_count"] = ( + self.update_weight_metrics.get("mooncake/write_bucket_count", 0) + 1 + ) + self.update_weight_metrics["mooncake/write_bytes"] = ( + self.update_weight_metrics.get("mooncake/write_bytes", 0) + total_bytes + ) + converted_named_tensors.clear() + if pbar is not None: + pbar.update(1) + finally: + ray.get(self.rollout_engine_lock.release.remote()) diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py index 15c4dd7231..752b031d21 100644 --- a/slime/backends/sglang_utils/sglang_engine.py +++ b/slime/backends/sglang_utils/sglang_engine.py @@ -468,6 +468,25 @@ def update_weights_from_distributed( payload, ) + def init_mooncake_weight_receiver(self, config: dict): + if self.node_rank != 0: + return None + return self._make_request("init_mooncake_weight_receiver", config) + + def update_weights_from_mooncake(self, manifest: dict): + if self.node_rank != 0: + return None + return self._make_request("update_weights_from_mooncake", manifest) + + def destroy_mooncake_weight_receiver(self): + if self.node_rank != 0: + return None + try: + return self._make_request("destroy_mooncake_weight_receiver") + except (requests.exceptions.ConnectionError, requests.exceptions.Timeout): + # The engine may have restarted before receiver initialization completed. + pass + def pause_generation(self): response = requests.post(f"http://{self.server_host}:{self.server_port}/pause_generation", json={}) response.raise_for_status() diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 2e8044c705..637eebbf69 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -145,14 +145,59 @@ def add_train_arguments(parser): ) parser.add_argument( "--update-weight-transport", - choices=["nccl", "disk"], + choices=["nccl", "disk", "mooncake"], default="nccl", help=( "Carrier for weight sync. In full mode, 'nccl' broadcasts chunks and " "'disk' writes a complete HF checkpoint under --update-weight-disk-dir " - "before engines reload it. In delta mode, 'nccl' broadcasts sparse deltas; " - "'disk' writes sparse safetensors under --update-weight-disk-dir and pushes " - "once at end-of-sync." + "before engines reload it; 'mooncake' sends chunks through Mooncake Transfer " + "Engine. In delta mode, 'nccl' broadcasts sparse deltas; 'disk' writes sparse " + "safetensors under --update-weight-disk-dir and pushes once at end-of-sync." + ), + ) + parser.add_argument( + "--mooncake-metadata-server", + type=str, + default="P2PHANDSHAKE", + help="Mooncake Transfer Engine metadata server for --update-weight-transport=mooncake.", + ) + parser.add_argument( + "--mooncake-protocol", + choices=["tcp", "rdma"], + default="tcp", + help="Mooncake Transfer Engine protocol for weight sync.", + ) + parser.add_argument( + "--mooncake-device-name", + type=str, + default="", + help="Mooncake Transfer Engine device name, for example an RDMA device.", + ) + parser.add_argument( + "--mooncake-rpc-port-base", + type=int, + default=None, + help=( + "Reserved base RPC port for SGLang Mooncake receivers. The current Mooncake " + "Python binding auto-assigns ports, so setting this option is rejected." + ), + ) + parser.add_argument( + "--mooncake-buffer-size", + type=int, + default=None, + help=( + "Mooncake receiver buffer size in bytes. Defaults to --update-weight-buffer-size " "when unset." + ), + ) + parser.add_argument( + "--mooncake-buffer-count", + type=int, + default=1, + help=( + "Reserved Mooncake receiver buffer slot count. The current sender uses one " + "slot per SGLang engine, so this must stay 1 until double-buffered updates " + "are implemented." ), ) parser.add_argument( @@ -237,7 +282,7 @@ def add_train_arguments(parser): type=str, nargs="*", default=None, - help="""List of regex patterns of parameter names to TRAIN. All other parameters will be FROZEN. + help=r"""List of regex patterns of parameter names to TRAIN. All other parameters will be FROZEN. Supports Python regex syntax (re.search). Examples: @@ -1639,6 +1684,8 @@ def parse_megatron_role_args(base_args, megatron_config_path, role): ) role_args = _apply_megatron_role_overrides(base_args, overrides, role) + if role == "actor" and hasattr(role_args, "update_weight_transport"): + _validate_update_weight_args(role_args) logger.info( f"Parsed megatron config for role={role} from {megatron_config_path}: overrides = {list(overrides.keys())}" ) @@ -1728,9 +1775,166 @@ def _resolve_update_weight_disk_dir(args) -> None: args.update_weight_delta_dir = disk_dir +def _validate_mooncake_external_args(args) -> None: + if getattr(args, "update_weight_transport", None) == "mooncake": + if ( + getattr(args, "rollout_external", False) + or getattr(args, "rollout_external_engine_addrs", None) is not None + ): + raise ValueError( + "--update-weight-transport=mooncake currently supports only slime-managed " + "SGLang engines; external engines must include the Mooncake receiver endpoint " + "before this transport can be enabled." + ) + + +_MOONCAKE_SGLANG_OVERRIDE_ALIASES = { + "dp_size": ("dp_size", "data_parallel_size"), + "tp_size": ("tp_size", "tensor_parallel_size"), + "pp_size": ("pp_size", "pipeline_parallel_size"), + "ep_size": ("ep_size", "expert_parallel_size"), + "enable_dp_attention": ("enable_dp_attention",), +} + + +def _normalize_sglang_overrides(overrides): + return {str(key).replace("-", "_"): value for key, value in overrides.items()} + + +def _get_sglang_override(overrides, field_name): + for alias in _MOONCAKE_SGLANG_OVERRIDE_ALIASES[field_name]: + if alias in overrides: + return overrides[alias] + return None + + +def _iter_mooncake_sglang_config_parallelism(args): + config_path = getattr(args, "sglang_config", None) + if not config_path: + return + + from slime.backends.sglang_utils.sglang_config import SglangConfig + + config = SglangConfig.from_yaml(config_path) + for model in config.models: + model.resolve(args) + # Mooncake P1 updates only the actor/weight-receiving model. Frozen + # reference/reward models and placeholder groups do not need receivers. + if not model.update_weights: + continue + for group in model.server_groups: + if group.worker_type == "placeholder": + continue + overrides = _normalize_sglang_overrides(group.overrides) + yield { + "num_gpus_per_engine": group.num_gpus_per_engine, + "dp_size": _get_sglang_override(overrides, "dp_size"), + "tp_size": _get_sglang_override(overrides, "tp_size"), + "pp_size": _get_sglang_override(overrides, "pp_size"), + "ep_size": _get_sglang_override(overrides, "ep_size"), + "enable_dp_attention": _get_sglang_override(overrides, "enable_dp_attention"), + } + + +def _validate_mooncake_parallelism_args(args) -> None: + rollout_gpus_per_engine = getattr(args, "rollout_num_gpus_per_engine", 1) + if int(rollout_gpus_per_engine) != 1: + raise ValueError( + "--update-weight-transport=mooncake currently supports only one GPU per rollout engine; " + f"got --rollout-num-gpus-per-engine={rollout_gpus_per_engine}." + ) + + pipeline_model_parallel_size = getattr(args, "pipeline_model_parallel_size", 1) + if int(pipeline_model_parallel_size) != 1: + raise ValueError( + "--update-weight-transport=mooncake currently supports only Megatron pipeline model parallel size 1; " + f"got pipeline_model_parallel_size={pipeline_model_parallel_size}." + ) + + sglang_dp_size = getattr(args, "sglang_dp_size", getattr(args, "sglang_data_parallel_size", 1)) + if int(sglang_dp_size) != 1: + raise ValueError( + "--update-weight-transport=mooncake currently supports only SGLang DP size 1; " + f"got sglang_dp_size={sglang_dp_size}." + ) + + sglang_ep_size = getattr(args, "sglang_ep_size", getattr(args, "sglang_expert_parallel_size", 1)) + if int(sglang_ep_size) != 1: + raise ValueError( + "--update-weight-transport=mooncake currently supports only SGLang expert parallel size 1; " + f"got sglang_ep_size={sglang_ep_size}." + ) + + if getattr(args, "sglang_enable_dp_attention", False): + raise ValueError("--update-weight-transport=mooncake currently does not support SGLang DP attention.") + + for group_parallelism in _iter_mooncake_sglang_config_parallelism(args): + override_dp_size = group_parallelism["dp_size"] + if override_dp_size is not None and int(override_dp_size) != 1: + raise ValueError( + "--update-weight-transport=mooncake currently supports only SGLang DP size 1; " + f"sglang_config overrides contain dp_size={override_dp_size}." + ) + + for key in ("tp_size", "pp_size"): + override_size = group_parallelism[key] + if override_size is not None and int(override_size) != 1: + raise ValueError( + "--update-weight-transport=mooncake currently supports only one GPU per rollout engine; " + f"sglang_config overrides contain {key}={override_size}." + ) + + override_ep_size = group_parallelism["ep_size"] + if override_ep_size is not None and int(override_ep_size) != 1: + raise ValueError( + "--update-weight-transport=mooncake currently supports only SGLang expert parallel size 1; " + f"sglang_config overrides contain ep_size={override_ep_size}." + ) + + if group_parallelism["enable_dp_attention"]: + raise ValueError( + "--update-weight-transport=mooncake currently does not support SGLang DP attention; " + "sglang_config overrides contain enable_dp_attention=True." + ) + + gpus_per_engine = group_parallelism["num_gpus_per_engine"] + if int(gpus_per_engine) != 1: + raise ValueError( + "--update-weight-transport=mooncake currently supports only one GPU per rollout engine; " + f"sglang_config contains num_gpus_per_engine={gpus_per_engine}." + ) + + def _validate_update_weight_args(args) -> None: _resolve_update_weight_disk_dir(args) + if args.update_weight_transport == "mooncake": + _validate_mooncake_external_args(args) + _validate_mooncake_parallelism_args(args) + if args.update_weight_mode != "full": + raise ValueError( + "--update-weight-mode=delta with --update-weight-transport=mooncake is not supported yet." + ) + if getattr(args, "mooncake_metadata_server", "P2PHANDSHAKE") != "P2PHANDSHAKE": + raise ValueError( + "--update-weight-transport=mooncake currently supports only " + "--mooncake-metadata-server=P2PHANDSHAKE. External Mooncake metadata servers " + "need globally unique receiver names before they can be enabled." + ) + if args.colocate: + raise ValueError( + "--update-weight-transport=mooncake is not supported with --colocate; colocated " + "weight sync uses CUDA IPC tensor handles." + ) + if getattr(args, "mooncake_buffer_count", 1) != 1: + raise ValueError("--mooncake-buffer-count is reserved for future double buffering and must be 1.") + if getattr(args, "mooncake_buffer_size", None) is not None and args.mooncake_buffer_size <= 0: + raise ValueError("--mooncake-buffer-size must be positive when set.") + if getattr(args, "mooncake_rpc_port_base", None) is not None: + raise ValueError( + "--mooncake-rpc-port-base is reserved but not supported by the current Mooncake " + "Python TransferEngine binding." + ) if args.update_weight_mode == "delta": if args.update_weight_transport not in ("nccl", "disk"): raise ValueError( @@ -1879,6 +2083,7 @@ def slime_validate_args(args): args.debug_train_only = True args.rollout_external = args.rollout_external_engine_addrs is not None + _validate_mooncake_external_args(args) if args.rollout_external and not args.debug_train_only: apply_external_engine_info_to_args(args, logger=logger) diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py index db5c0bfd81..5c3c2b15bd 100644 --- a/tests/test_megatron_argument_validation.py +++ b/tests/test_megatron_argument_validation.py @@ -351,6 +351,27 @@ def test_slime_validate_args_preserves_zero_rollout_gpus_without_colocate(monkey assert args.offload_rollout is False +@pytest.mark.unit +def test_slime_validate_args_rejects_mooncake_with_external_rollout_before_discovery(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + called = False + + def _apply_external_engine_info_to_args(*_args, **_kwargs): + nonlocal called + called = True + + monkeypatch.setattr(module, "apply_external_engine_info_to_args", _apply_external_engine_info_to_args) + args = make_slime_validate_args( + update_weight_transport="mooncake", + rollout_external_engine_addrs=["127.0.0.1:30000"], + ) + + with pytest.raises(ValueError, match="slime-managed SGLang"): + module.slime_validate_args(args) + + assert called is False + + @pytest.mark.unit def test_update_weight_delta_rejects_colocate(monkeypatch): module = load_slime_arguments_module(monkeypatch) @@ -381,5 +402,344 @@ def test_update_weight_delta_rejects_unknown_transport(monkeypatch): module._validate_update_weight_args(args) +@pytest.mark.unit +def test_update_weight_full_allows_mooncake_without_disk_dir(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + ) + + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_multi_gpu_rollout_engine(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + rollout_num_gpus_per_engine=2, + ) + + with pytest.raises(ValueError, match="one GPU per rollout engine"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_sglang_dp(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + rollout_num_gpus_per_engine=1, + sglang_dp_size=2, + ) + + with pytest.raises(ValueError, match="SGLang DP size 1"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_sglang_ep(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + rollout_num_gpus_per_engine=1, + sglang_dp_size=1, + sglang_ep_size=2, + ) + + with pytest.raises(ValueError, match="SGLang expert parallel size 1"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_sglang_dp_attention(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + rollout_num_gpus_per_engine=1, + sglang_dp_size=1, + sglang_enable_dp_attention=True, + ) + + with pytest.raises(ValueError, match="SGLang DP attention"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_megatron_pipeline_parallelism(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + rollout_num_gpus_per_engine=1, + sglang_dp_size=1, + pipeline_model_parallel_size=2, + ) + + with pytest.raises(ValueError, match="Megatron pipeline model parallel size 1"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_actor_megatron_config_pipeline_parallelism_override(monkeypatch, tmp_path): + module = load_slime_arguments_module(monkeypatch) + config_path = tmp_path / "megatron.yaml" + config_path.write_text( + """ +megatron: + - role: actor + overrides: + pipeline_model_parallel_size: 2 +""", + encoding="utf-8", + ) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + rollout_num_gpus_per_engine=1, + sglang_dp_size=1, + pipeline_model_parallel_size=1, + ) + + with pytest.raises(ValueError, match="Megatron pipeline model parallel size 1"): + module.parse_megatron_role_args(args, str(config_path), role="actor") + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_sglang_config_multi_gpu_group(monkeypatch, tmp_path): + module = load_slime_arguments_module(monkeypatch) + config_path = tmp_path / "sglang.yaml" + config_path.write_text( + """ +sglang: + - name: actor + update_weights: true + server_groups: + - worker_type: regular + num_gpus: 2 + num_gpus_per_engine: 2 +""", + encoding="utf-8", + ) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + rollout_num_gpus_per_engine=1, + sglang_dp_size=1, + sglang_config=str(config_path), + hf_checkpoint="/actor", + ) + + with pytest.raises(ValueError, match="one GPU per rollout engine"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "override_yaml, message", + [ + ("dp_size: 2", "SGLang DP size 1"), + ("data-parallel-size: 2", "SGLang DP size 1"), + ("tp-size: 2", "one GPU per rollout engine"), + ("tensor_parallel_size: 2", "one GPU per rollout engine"), + ("pp_size: 2", "one GPU per rollout engine"), + ("pipeline-parallel-size: 2", "one GPU per rollout engine"), + ("ep_size: 2", "SGLang expert parallel size 1"), + ("expert-parallel-size: 2", "SGLang expert parallel size 1"), + ("enable_dp_attention: true", "SGLang DP attention"), + ], +) +def test_update_weight_mooncake_rejects_sglang_config_parallelism_overrides( + monkeypatch, tmp_path, override_yaml, message +): + module = load_slime_arguments_module(monkeypatch) + config_path = tmp_path / "sglang.yaml" + config_path.write_text( + f""" +sglang: + - name: actor + update_weights: true + server_groups: + - worker_type: regular + num_gpus: 1 + overrides: + {override_yaml} +""", + encoding="utf-8", + ) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + rollout_num_gpus_per_engine=1, + sglang_dp_size=1, + sglang_config=str(config_path), + hf_checkpoint="/actor", + ) + + with pytest.raises(ValueError, match=message): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_allows_inferred_frozen_sglang_config_multi_gpu_group(monkeypatch, tmp_path): + module = load_slime_arguments_module(monkeypatch) + config_path = tmp_path / "sglang.yaml" + config_path.write_text( + """ +sglang: + - name: actor + server_groups: + - worker_type: regular + num_gpus: 1 + - name: ref + model_path: /ref + server_groups: + - worker_type: regular + num_gpus: 2 + num_gpus_per_engine: 2 +""", + encoding="utf-8", + ) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + rollout_num_gpus_per_engine=1, + sglang_dp_size=1, + sglang_config=str(config_path), + hf_checkpoint="/actor", + ) + + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_non_p2p_metadata_server(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="127.0.0.1:2379", + ) + + with pytest.raises(ValueError, match="metadata-server=P2PHANDSHAKE"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_delta_rejects_mooncake_until_delta_carrier_exists(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="delta", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + ) + + with pytest.raises(ValueError, match="delta.*mooncake.*not supported"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_reserved_rpc_port_base(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + mooncake_rpc_port_base=18000, + ) + + with pytest.raises(ValueError, match="rpc-port-base.*reserved"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_reserved_buffer_count(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + mooncake_buffer_count=2, + ) + + with pytest.raises(ValueError, match="buffer-count.*reserved.*must be 1"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_mooncake_rejects_external_rollout(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="full", + update_weight_transport="mooncake", + update_weight_disk_dir=None, + update_weight_delta_dir=None, + colocate=False, + mooncake_metadata_server="P2PHANDSHAKE", + rollout_external=False, + rollout_external_engine_addrs=["127.0.0.1:30000"], + ) + + with pytest.raises(ValueError, match="slime-managed SGLang"): + module._validate_update_weight_args(args) + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_mooncake_weight_transport.py b/tests/test_mooncake_weight_transport.py new file mode 100644 index 0000000000..96a4150556 --- /dev/null +++ b/tests/test_mooncake_weight_transport.py @@ -0,0 +1,640 @@ +from __future__ import annotations + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + + +class FakeTensor: + def __init__(self, shape, dtype: str, element_size: int): + self.shape = tuple(shape) + self.dtype = dtype + self._element_size = element_size + + def numel(self): + out = 1 + for dim in self.shape: + out *= dim + return out + + def element_size(self): + return self._element_size + + +def load_mooncake_module(): + module_path = ( + Path(__file__).resolve().parents[1] + / "slime" + / "backends" + / "megatron_utils" + / "update_weight" + / "update_weight_from_mooncake.py" + ) + module_name = "test_update_weight_from_mooncake_module" + sys.modules.pop(module_name, None) + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +class RemoteMethod: + def __init__(self, fn): + self._fn = fn + + def remote(self, *args, **kwargs): + return self._fn(*args, **kwargs) + + +class RecordingClient: + def __init__(self): + self.close_count = 0 + self.calls = [] + + def write_bucket(self, **kwargs): + client = load_mooncake_module().FakeMooncakeTransferClient() + manifests = client.write_bucket(**kwargs) + self.calls.extend(client.calls) + return manifests + + def close(self): + self.close_count += 1 + + +def ray_get(value): + if isinstance(value, list): + out = [] + for item in value: + if isinstance(item, Exception): + raise item + out.append(item) + return out + if isinstance(value, Exception): + raise value + return value + + +def patch_pp_source_rank(monkeypatch, module): + monkeypatch.setattr(module, "ray", types.SimpleNamespace(get=ray_get)) + monkeypatch.setattr(module, "time", types.SimpleNamespace(sleep=lambda _seconds: None)) + monkeypatch.setattr( + module.mpu, + "get_data_parallel_rank", + lambda with_context_parallel=True: 0, + ) + monkeypatch.setattr(module.mpu, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(module.mpu, "get_pipeline_model_parallel_rank", lambda: 0) + + +@pytest.mark.unit +def test_build_manifest_assigns_offsets_and_preserves_tensor_metadata(): + module = load_mooncake_module() + + tensors = [ + ("layer.0.weight", FakeTensor((2, 3), "bfloat16", 2)), + ("layer.0.bias", FakeTensor((3,), "float32", 4)), + ] + descriptor = { + "target_name": "127.0.0.1:18000", + "buffers": [{"slot": 0, "addr": 4096, "capacity": 64}], + } + + manifest = module.build_mooncake_manifest( + descriptor=descriptor, + named_tensors=tensors, + weight_version=7, + load_format=None, + buffer_slot=0, + ) + + assert manifest["weight_version"] == "7" + assert manifest["target_name"] == "127.0.0.1:18000" + assert manifest["buffer_slot"] == 0 + assert manifest["remote_addr"] == 4096 + assert manifest["nbytes"] == 24 + assert manifest["tensors"] == [ + { + "name": "layer.0.weight", + "dtype": "bfloat16", + "shape": [2, 3], + "offset": 0, + "nbytes": 12, + }, + { + "name": "layer.0.bias", + "dtype": "float32", + "shape": [3], + "offset": 12, + "nbytes": 12, + }, + ] + + +@pytest.mark.unit +def test_build_manifest_preserves_cuda_receiver_buffer_device(): + module = load_mooncake_module() + + manifest = module.build_mooncake_manifest( + descriptor={ + "target_name": "127.0.0.1:18000", + "buffer_device": "cuda", + "buffers": [{"slot": 0, "addr": 4096, "capacity": 64, "device": "cuda"}], + }, + named_tensors=[("layer.weight", FakeTensor((4,), "float16", 2))], + weight_version=1, + load_format=None, + ) + + assert manifest["buffer_device"] == "cuda" + + +@pytest.mark.unit +def test_build_manifest_composes_target_name_from_server_name_and_rpc_port(): + module = load_mooncake_module() + + manifest = module.build_mooncake_manifest( + descriptor={ + "server_name": "127.0.0.1", + "rpc_port": 18000, + "buffers": [{"slot": 0, "addr": 4096, "capacity": 64}], + }, + named_tensors=[("layer.weight", FakeTensor((4,), "float16", 2))], + weight_version=1, + load_format=None, + ) + + assert manifest["target_name"] == "127.0.0.1:18000" + + +@pytest.mark.unit +def test_build_manifest_rejects_bucket_larger_than_receiver_buffer(): + module = load_mooncake_module() + + descriptor = { + "target_name": "127.0.0.1:18000", + "buffers": [{"slot": 0, "addr": 4096, "capacity": 8}], + } + + with pytest.raises(ValueError, match="exceeds Mooncake receiver buffer"): + module.build_mooncake_manifest( + descriptor=descriptor, + named_tensors=[("too.large", FakeTensor((3,), "float32", 4))], + weight_version=1, + load_format=None, + buffer_slot=0, + ) + + +@pytest.mark.unit +def test_fake_mooncake_transfer_client_records_one_write_per_receiver(): + module = load_mooncake_module() + + client = module.FakeMooncakeTransferClient() + tensors = [("layer.weight", FakeTensor((4,), "float16", 2))] + descriptors = [ + {"target_name": "engine-a:18000", "buffers": [{"slot": 0, "addr": 1000, "capacity": 16}]}, + {"target_name": "engine-b:18001", "buffers": [{"slot": 0, "addr": 2000, "capacity": 16}]}, + ] + + manifests = client.write_bucket( + receiver_descriptors=descriptors, + named_tensors=tensors, + weight_version=3, + load_format=None, + delta=None, + buffer_slot=0, + ) + + assert [m["target_name"] for m in manifests] == ["engine-a:18000", "engine-b:18001"] + assert [call.target_name for call in client.calls] == ["engine-a:18000", "engine-b:18001"] + assert [call.remote_addr for call in client.calls] == [1000, 2000] + assert all(call.nbytes == 8 for call in client.calls) + assert all(call.tensor_names == ("layer.weight",) for call in client.calls) + + +@pytest.mark.unit +def test_mooncake_transfer_client_raises_when_unregister_fails(monkeypatch): + module = load_mooncake_module() + + class FakeCudaTensor(FakeTensor): + is_cuda = True + device = "cuda:0" + + def is_contiguous(self): + return True + + def data_ptr(self): + return 1234 + + class FakeEngine: + def initialize(self, *_args): + return 0 + + def register_memory(self, *_args): + return 0 + + def unregister_memory(self, *_args): + return -5 + + engine_mod = types.ModuleType("mooncake.engine") + engine_mod.TransferEngine = FakeEngine + monkeypatch.setitem(sys.modules, "mooncake", types.ModuleType("mooncake")) + monkeypatch.setitem(sys.modules, "mooncake.engine", engine_mod) + + client = module.MooncakeTransferClient( + local_server_name="trainer", + metadata_server="P2PHANDSHAKE", + protocol="tcp", + device_name="", + ) + monkeypatch.setattr(client, "_write_manifest", lambda _manifest, _named_tensors: None) + + with pytest.raises(RuntimeError, match="unregister_memory\\(1234\\)=-5"): + client.write_bucket( + receiver_descriptors=[ + {"target_name": "engine:18000", "buffers": [{"slot": 0, "addr": 1000, "capacity": 16}]} + ], + named_tensors=[("layer.weight", FakeCudaTensor((4,), "float16", 2))], + weight_version=1, + load_format=None, + delta=None, + ) + + +@pytest.mark.unit +def test_mooncake_transfer_client_rejects_non_cuda_receiver_buffer(monkeypatch): + module = load_mooncake_module() + + class FakeEngine: + def initialize(self, *_args): + return 0 + + def register_memory(self, *_args): + return 0 + + def unregister_memory(self, *_args): + return 0 + + engine_mod = types.ModuleType("mooncake.engine") + engine_mod.TransferEngine = FakeEngine + monkeypatch.setitem(sys.modules, "mooncake", types.ModuleType("mooncake")) + monkeypatch.setitem(sys.modules, "mooncake.engine", engine_mod) + + client = module.MooncakeTransferClient( + local_server_name="trainer", + metadata_server="P2PHANDSHAKE", + protocol="tcp", + device_name="", + ) + + tensor = FakeTensor((4,), "float16", 2) + with pytest.raises(RuntimeError, match="CUDA receiver buffers"): + client.write_bucket( + receiver_descriptors=[ + { + "target_name": "engine:18000", + "buffer_device": "cpu", + "buffers": [{"slot": 0, "addr": 1000, "capacity": 16, "device": "cpu"}], + } + ], + named_tensors=[("layer.weight", tensor)], + weight_version=1, + load_format=None, + delta=None, + ) + + +@pytest.mark.unit +def test_default_local_server_name_uses_host_ip_for_p2p_handshake(monkeypatch): + module = load_mooncake_module() + + monkeypatch.setattr(module.socket, "gethostbyname", lambda hostname: f"ip-for-{hostname}") + + assert ( + module.default_mooncake_local_server_name( + metadata_server="P2PHANDSHAKE", + pp_rank=3, + host_getter=lambda: "trainer-host", + ) + == "ip-for-trainer-host" + ) + + +@pytest.mark.unit +def test_default_local_server_name_uses_stable_name_for_metadata_server(): + module = load_mooncake_module() + + assert ( + module.default_mooncake_local_server_name( + metadata_server="127.0.0.1:2379", + pp_rank=3, + host_getter=lambda: "trainer-host", + ) + == "slime-train-pp3" + ) + + +@pytest.mark.unit +def test_update_weight_from_mooncake_publishes_manifests_to_rollout_engines(monkeypatch): + module = load_mooncake_module() + + acquire = types.SimpleNamespace(remote=lambda: True) + release = types.SimpleNamespace(remote=lambda: None) + lock = types.SimpleNamespace(acquire=acquire, release=release) + published = [] + + class _Engine: + def __init__(self, descriptor): + self.init_mooncake_weight_receiver = RemoteMethod(lambda _config: descriptor) + self.update_weights_from_mooncake = RemoteMethod(lambda manifest: published.append(manifest)) + self.destroy_mooncake_weight_receiver = RemoteMethod(lambda: None) + + patch_pp_source_rank(monkeypatch, module) + + updater = module.UpdateWeightFromMooncake( + args=types.SimpleNamespace( + mooncake_metadata_server="P2PHANDSHAKE", + mooncake_protocol="tcp", + mooncake_device_name="", + mooncake_buffer_size=16, + mooncake_buffer_count=1, + ), + model=[], + weights_getter=lambda: {}, + model_name="qwen", + quantization_config=None, + transfer_client=module.FakeMooncakeTransferClient(), + ) + updater.connect_rollout_engines( + [ + _Engine({"target_name": "engine-a:18000", "buffers": [{"slot": 0, "addr": 1000, "capacity": 16}]}), + _Engine({"target_name": "engine-b:18001", "buffers": [{"slot": 0, "addr": 2000, "capacity": 16}]}), + ], + lock, + ) + + updater._update_bucket_weights_from_mooncake( + [("layer.weight", FakeTensor((4,), "float16", 2))], + pbar=None, + ) + + assert [manifest["target_name"] for manifest in published] == ["engine-a:18000", "engine-b:18001"] + assert updater.update_weight_metrics["mooncake/write_bucket_count"] == 1 + assert updater.update_weight_metrics["mooncake/write_bytes"] == 16 + + +@pytest.mark.unit +def test_connect_reinitializes_existing_receivers_before_reconnect(monkeypatch): + module = load_mooncake_module() + patch_pp_source_rank(monkeypatch, module) + + class Engine: + def __init__(self): + self.init_count = 0 + self.destroy_count = 0 + self.init_mooncake_weight_receiver = RemoteMethod(self._init) + self.destroy_mooncake_weight_receiver = RemoteMethod(self._destroy) + + def _init(self, _config): + self.init_count += 1 + return { + "target_name": f"engine:{18000 + self.init_count}", + "buffers": [{"slot": 0, "addr": 1000, "capacity": 16}], + } + + def _destroy(self): + self.destroy_count += 1 + + first_client = RecordingClient() + second_client = RecordingClient() + monkeypatch.setattr( + module.MooncakeTransferClient, + "from_args", + classmethod(lambda _cls, _args, pp_rank: second_client), + ) + engine = Engine() + updater = module.UpdateWeightFromMooncake( + args=types.SimpleNamespace( + mooncake_metadata_server="P2PHANDSHAKE", + mooncake_protocol="tcp", + mooncake_device_name="", + mooncake_buffer_size=16, + mooncake_buffer_count=1, + ), + model=[], + weights_getter=lambda: {}, + model_name="qwen", + quantization_config=None, + transfer_client=first_client, + ) + + updater.connect_rollout_engines([engine], types.SimpleNamespace()) + updater.connect_rollout_engines([engine], types.SimpleNamespace()) + + assert engine.init_count == 2 + assert engine.destroy_count == 1 + assert first_client.close_count == 1 + assert updater._transfer_client is second_client + + +@pytest.mark.unit +def test_connect_cleans_receiver_when_transfer_client_init_fails(monkeypatch): + module = load_mooncake_module() + patch_pp_source_rank(monkeypatch, module) + + class Engine: + def __init__(self): + self.destroy_count = 0 + self.init_mooncake_weight_receiver = RemoteMethod( + lambda _config: {"target_name": "engine:18000", "buffers": [{"slot": 0, "addr": 1000, "capacity": 16}]} + ) + self.destroy_mooncake_weight_receiver = RemoteMethod(self._destroy) + + def _destroy(self): + self.destroy_count += 1 + + monkeypatch.setattr( + module.MooncakeTransferClient, + "from_args", + classmethod(lambda _cls, _args, pp_rank: (_ for _ in ()).throw(RuntimeError("client init failed"))), + ) + engine = Engine() + updater = module.UpdateWeightFromMooncake( + args=types.SimpleNamespace( + mooncake_metadata_server="P2PHANDSHAKE", + mooncake_protocol="tcp", + mooncake_device_name="", + mooncake_buffer_size=16, + mooncake_buffer_count=1, + ), + model=[], + weights_getter=lambda: {}, + model_name="qwen", + quantization_config=None, + ) + + with pytest.raises(RuntimeError, match="client init failed"): + updater.connect_rollout_engines([engine], types.SimpleNamespace()) + + assert engine.destroy_count == 1 + assert updater._receiver_descriptors == [] + assert updater._connected is False + + +@pytest.mark.unit +def test_connect_rejects_multi_gpu_rollout_engine_until_tp_receiver_contract_exists(monkeypatch): + module = load_mooncake_module() + patch_pp_source_rank(monkeypatch, module) + + updater = module.UpdateWeightFromMooncake( + args=types.SimpleNamespace( + mooncake_metadata_server="P2PHANDSHAKE", + mooncake_protocol="tcp", + mooncake_device_name="", + mooncake_buffer_size=16, + mooncake_buffer_count=1, + rollout_num_gpus_per_engine=2, + ), + model=[], + weights_getter=lambda: {}, + model_name="qwen", + quantization_config=None, + transfer_client=module.FakeMooncakeTransferClient(), + ) + + with pytest.raises(NotImplementedError, match="one GPU per rollout engine"): + updater.connect_rollout_engines( + rollout_engines=[types.SimpleNamespace()], + rollout_engine_lock=types.SimpleNamespace(), + engine_gpu_counts=[2], + ) + + +@pytest.mark.unit +def test_disconnect_closes_client_even_when_destroy_rpc_fails(monkeypatch): + module = load_mooncake_module() + patch_pp_source_rank(monkeypatch, module) + client = RecordingClient() + + engine = types.SimpleNamespace( + destroy_mooncake_weight_receiver=RemoteMethod(lambda: RuntimeError("destroy failed")), + ) + updater = module.UpdateWeightFromMooncake( + args=types.SimpleNamespace(), + model=[], + weights_getter=lambda: {}, + model_name="qwen", + quantization_config=None, + transfer_client=client, + ) + updater._is_pp_src_rank = True + updater.rollout_engines = [engine] + updater._receiver_descriptors = [{"target_name": "engine:18000", "buffers": []}] + updater._connected = True + + with pytest.raises(RuntimeError, match="destroy failed"): + updater.disconnect_rollout_engines() + + assert client.close_count == 1 + assert updater._transfer_client is None + assert updater._receiver_descriptors == [] + assert updater._connected is False + + +@pytest.mark.unit +def test_update_failure_still_releases_rollout_lock(monkeypatch): + module = load_mooncake_module() + patch_pp_source_rank(monkeypatch, module) + + class Lock: + def __init__(self): + self.release_count = 0 + self.acquire = RemoteMethod(lambda: True) + self.release = RemoteMethod(self._release) + + def _release(self): + self.release_count += 1 + + lock = Lock() + engine = types.SimpleNamespace( + update_weights_from_mooncake=RemoteMethod(lambda _manifest: RuntimeError("publish failed")), + ) + updater = module.UpdateWeightFromMooncake( + args=types.SimpleNamespace(), + model=[], + weights_getter=lambda: {}, + model_name="qwen", + quantization_config=None, + transfer_client=module.FakeMooncakeTransferClient(), + ) + updater._is_pp_src_rank = True + updater.rollout_engine_lock = lock + updater.rollout_engines = [engine] + updater._receiver_descriptors = [ + {"target_name": "engine:18000", "buffers": [{"slot": 0, "addr": 1000, "capacity": 16}]} + ] + updater.weight_version = 1 + + with pytest.raises(RuntimeError, match="publish failed"): + updater._update_bucket_weights_from_mooncake([("layer.weight", FakeTensor((4,), "float16", 2))]) + + assert lock.release_count == 1 + + +@pytest.mark.unit +@pytest.mark.parametrize("patch_dir", ["latest", "v0.5.12.post1", "v0.5.9"]) +def test_sglang_patch_exposes_mooncake_weight_receiver_contract(patch_dir): + patch_path = Path(__file__).resolve().parents[1] / "docker" / "patch" / patch_dir / "sglang.patch" + patch_text = patch_path.read_text() + + common_snippets = [ + "class InitMooncakeWeightReceiverReqInput(BaseReq):", + "class UpdateWeightsFromMooncakeReqInput(BaseReq):", + "class DestroyMooncakeWeightReceiverReqInput(BaseReq):", + '@app.post("/init_mooncake_weight_receiver")', + '@app.post("/update_weights_from_mooncake")', + '@app.post("/destroy_mooncake_weight_receiver")', + "self.tp_worker.init_mooncake_weight_receiver", + "self.model_runner.init_mooncake_weight_receiver", + "def update_weights_from_mooncake(", + "def destroy_mooncake_weight_receiver(", + "buffer_device: Optional[str] = None", + '"buffer_device": result.buffer_device', + '"device": buf.device.type', + "cleanup_ok, cleanup_message = self.destroy_mooncake_weight_receiver()", + 'getattr(buf, "is_cuda", False)', + 'unsupported = "Mooncake weight receiver currently supports only dp_size=1."', + ] + version_snippets = { + "latest": [ + '("init_mooncake_weight_receiver", InitMooncakeWeightReceiverReqOutput)', + "self.weight_updater.init_mooncake_weight_receiver", + ], + "v0.5.12.post1": [ + '("init_mooncake_weight_receiver", InitMooncakeWeightReceiverReqOutput)', + "self.init_mooncake_weight_receiver_communicator(obj)", + ], + "v0.5.9": [ + "self.init_mooncake_weight_receiver_communicator = _Communicator", + "self.init_mooncake_weight_receiver_communicator.handle_recv", + ], + } + required_snippets = common_snippets + version_snippets[patch_dir] + missing = [snippet for snippet in required_snippets if snippet not in patch_text] + + assert missing == [] + assert "assert self.server_args.dp_size == 1" not in patch_text + assert ( + 'if ret != 0:\n+ free_engine = getattr(engine, "freeEngine", None)\n' + "+ if callable(free_engine):\n+ free_engine()\n+ last_error =" + in patch_text + ) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__]))