Revert #27749 and re-impl mtp update weight from distributed

python/sglang/srt/managers/io_struct.py

@@ -1496,8 +1496,9 @@ class UpdateWeightsFromDistributedReqInput(BaseReq):
     weight_version: Optional[str] = None
     # Optional format specification for loading
     load_format: Optional[str] = None
-    # Optional: Determine whether to disable updating the draft model
-    disable_draft_model: Optional[bool] = None
+    # Which model runners to update: "target" (target model only), "draft" (draft
+    # worker(s) only), or "all" (default).
+    selector: Literal["target", "draft", "all"] = "all"
     # Whether to call torch.cuda.empty_cache() during flush
     torch_empty_cache: bool = False
 
@@ -1525,8 +1526,9 @@ class UpdateWeightsFromTensorReqInput(BaseReq):
     abort_all_requests: bool = False
     # Optional: Update weight version along with weights
     weight_version: Optional[str] = None
-    # Optional: Determine whether to disable updating the draft model
-    disable_draft_model: Optional[bool] = None
+    # Which model runners to update: "target" (target model only), "draft" (draft
+    # worker(s) only), or "all" (default).
+    selector: Literal["target", "draft", "all"] = "all"
     # Whether to call torch.cuda.empty_cache() during flush
     torch_empty_cache: bool = False
 

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

@@ -6,7 +6,7 @@
 import traceback
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, Iterator, Optional, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple
 
 import torch
 
@@ -41,6 +41,8 @@
     UpdateWeightsFromTensorReqInput,
     UpdateWeightsFromTensorReqOutput,
 )
+from sglang.srt.utils import MultiprocessingSerializer
+from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
 
 logger = logging.getLogger(__name__)
 
@@ -72,6 +74,17 @@ def _merge_checksum_payloads(target: Dict, draft: Dict) -> Dict:
     return target
 
 
+def _parse_runner_selector(selector: str) -> Set[str]:
+    """Map a {target, draft, all} weight-op selector to the set of roles it covers."""
+    if selector == "all":
+        return {"target", "draft"}
+    if selector in ("target", "draft"):
+        return {selector}
+    raise ValueError(
+        f"invalid selector {selector!r}; expected 'target', 'draft', or 'all'"
+    )
+
+
 @dataclass(kw_only=True, slots=True)
 class SchedulerWeightUpdaterManager:
     tp_worker: Any
@@ -135,31 +148,65 @@ def destroy_weights_update_group(
         success, message = self.tp_worker.destroy_weights_update_group(recv_req)
         return DestroyWeightsUpdateGroupReqOutput(success, message)
 
+    def get_model_runners(self, selector: str = "all") -> List[Tuple[str, Any]]:
+        """Resolve a {target, draft} selector to (role, ModelRunner) pairs, target
+        first. role is "" for the target runner; draft roles come from the draft
+        worker's iter_draft_runners()."""
+        parsed = _parse_runner_selector(selector)
+        runners: List[Tuple[str, Any]] = []
+        if "target" in parsed:
+            runners.append(("", self.tp_worker.model_runner))
+        if "draft" in parsed and self.draft_worker is not None:
+            runners += self.draft_worker.iter_draft_runners()
+        return runners
+
     def update_weights_from_distributed(
         self,
         recv_req: UpdateWeightsFromDistributedReqInput,
     ) -> Tuple[bool, str]:
-        """Update the online model parameter."""
+        """Update the online model parameter, fanning out to the selected runners."""
         with self._observe_weight_load("distributed"):
-            if recv_req.disable_draft_model:
-                worker = self.tp_worker
-            else:
-                worker = self.draft_worker or self.tp_worker
-            success, message = worker.update_weights_from_distributed(recv_req)
+            # The target (main) model owns this process's connection to the training
+            # engine, so it receives the broadcast once; the received weights are then
+            # loaded into each selected runner locally.
+            try:
+                weights = self.tp_worker.model_runner.receive_weights_from_distributed(
+                    recv_req.names,
+                    recv_req.dtypes,
+                    recv_req.shapes,
+                    recv_req.group_name,
+                    recv_req.load_format,
+                )
+                for _, runner in self.get_model_runners(recv_req.selector):
+                    runner.load_weights(weights)
+                success, message = True, "Succeeded to update parameter online."
+            except Exception as e:
+                success = False
+                message = (
+                    f"Failed to update parameter online: {e}. The full weights of the "
+                    "ModelRunner are partially updated. Please discard the whole weights."
+                )
+                logger.error(message)
             if success:
                 self.flush_cache_after_weight_update(recv_req)
-            else:
-                logger.error(message)
             return UpdateWeightsFromDistributedReqOutput(success, message)
 
     def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
-        """Update the online model parameter from tensors."""
+        """Update the online model parameter from tensors, fanning out to the
+        selected runners."""
         with self._observe_weight_load("tensor"):
-            if recv_req.disable_draft_model:
-                worker = self.tp_worker
-            else:
-                worker = self.draft_worker or self.tp_worker
-            success, message = worker.update_weights_from_tensor(recv_req)
+            monkey_patch_torch_reductions()
+            named_tensors = MultiprocessingSerializer.deserialize(
+                recv_req.serialized_named_tensors[self.tp_worker.tp_rank]
+            )
+            success, message = True, "Success"
+            for _, runner in self.get_model_runners(recv_req.selector):
+                success, message = runner.update_weights_from_tensor(
+                    named_tensors=named_tensors,
+                    load_format=recv_req.load_format,
+                )
+                if not success:
+                    break
             if success:
                 self.flush_cache_after_weight_update(recv_req)
             else:

python/sglang/srt/managers/tp_worker.py

@@ -34,9 +34,7 @@
     SendWeightsToRemoteInstanceReqInput,
     UnloadLoRAAdapterReqInput,
     UpdateWeightFromDiskReqInput,
-    UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromIPCReqInput,
-    UpdateWeightsFromTensorReqInput,
 )
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
@@ -145,29 +143,6 @@ def send_weights_to_remote_instance(
         )
         return success, message
 
-    def update_weights_from_distributed(
-        self, recv_req: UpdateWeightsFromDistributedReqInput
-    ):
-        success, message = self.model_runner.update_weights_from_distributed(
-            recv_req.names,
-            recv_req.dtypes,
-            recv_req.shapes,
-            recv_req.group_name,
-            recv_req.load_format,
-        )
-        return success, message
-
-    def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
-
-        monkey_patch_torch_reductions()
-        success, message = self.model_runner.update_weights_from_tensor(
-            named_tensors=MultiprocessingSerializer.deserialize(
-                recv_req.serialized_named_tensors[self.tp_rank]
-            ),
-            load_format=recv_req.load_format,
-        )
-        return success, message
-
     def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput):
         """Update weights from IPC for checkpoint-engine integration."""
         success, message = self.model_runner.update_weights_from_ipc(recv_req)
@@ -428,6 +403,16 @@ def _init_dllm_algorithm(self):
     def model_runner(self) -> "ModelRunner":
         return self._model_runner
 
+    def iter_draft_runners(self) -> List[Tuple[str, "ModelRunner"]]:
+        # The target worker shares this class (is_draft_worker=False) and returns [].
+        if not self.is_draft_worker:
+            return []
+        if self.model_runner_list:
+            return [
+                (f"draft_step_{i}", r) for i, r in enumerate(self.model_runner_list)
+            ]
+        return [("draft", self.model_runner)]
+
     def register_hicache_layer_transfer_counter(self, counter: LayerDoneCounter):
         self.hicache_layer_transfer_counter = counter
 

python/sglang/srt/model_executor/model_runner.py

@@ -1944,62 +1944,27 @@ def destroy_weights_update_group(self, group_name):
             logger.error(message)
             return False, message
 
-    def update_weights_from_distributed(
-        self,
-        names,
-        dtypes,
-        shapes,
-        group_name,
-        load_format: Optional[str] = None,
-    ):
-        """
-        Update specific parameter in the model weights online
-        through `_model_update_group` process group.
-
-        Args:
-            name: the name of the parameter to be updated.
-            dtype: the data type of the parameter to be updated.
-            shape: the shape of the parameter to be updated.
-        """
+    def load_weights(self, weights) -> None:
+        """Load an in-memory list of (name, tensor) weights into this runner's model."""
+        self.model.load_weights(weights)
 
-        return self.update_weights_from_distributed_to_model_runners(
-            [self], names, dtypes, shapes, group_name, load_format
-        )
-
-    def update_weights_from_distributed_to_model_runners(
+    def receive_weights_from_distributed(
         self,
-        model_runners,
         names,
         dtypes,
         shapes,
         group_name,
         load_format: Optional[str] = None,
     ):
+        """Receive one weight broadcast from the training engine over this runner's
+        `_model_update_group`. Only the runner that joined the group (the target /
+        main model) can receive; the caller loads the result into each runner."""
+
         assert group_name in self._model_update_group, (
             f"Group {group_name} not in {list(self._model_update_group.keys())}. "
             "Please call `init_weights_update_group` first."
         )
 
-        try:
-            weights = self._receive_weights_from_distributed(
-                names, dtypes, shapes, group_name, load_format
-            )
-            for model_runner in model_runners:
-                model_runner.model.load_weights(weights)
-            return True, "Succeeded to update parameter online."
-
-        except Exception as e:
-            error_msg = (
-                f"Failed to update parameter online: {e}. "
-                f"The full weights of the ModelRunner are partially updated. "
-                f"Please discard the whole weights."
-            )
-            logger.error(error_msg)
-            return False, error_msg
-
-    def _receive_weights_from_distributed(
-        self, names, dtypes, shapes, group_name, load_format: Optional[str] = None
-    ):
         if load_format == "flattened_bucket":
             return self._receive_bucketed_weights_from_distributed(
                 names, dtypes, shapes, group_name

python/sglang/srt/speculative/dflash_worker.py

@@ -1,12 +1,11 @@
 import logging
 import math
 from copy import deepcopy
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
 from sglang.srt.distributed import get_tp_group
-from sglang.srt.managers.io_struct import UpdateWeightsFromDistributedReqInput
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.managers.tp_worker import TpModelWorker
@@ -35,6 +34,10 @@
 _is_npu = is_npu()
 
 
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
 logger = logging.getLogger(__name__)
 
 _FusedKVMaterializeHelper = None
@@ -346,23 +349,18 @@ def __getattr__(self, name):
         # Delegate anything not implemented yet to the target worker.
         return getattr(self.target_worker, name)
 
-    def update_weights_from_distributed(
-        self, recv_req: UpdateWeightsFromDistributedReqInput
-    ):
-        # Spelled out instead of falling through `__getattr__` so the gap is
-        # visible: this only updates the target. The DFlash draft model
-        # (`self.draft_model_runner`) is NOT refreshed and goes stale after a
-        # distributed weight update.
-        # TODO(dflash): fan out to the draft runner like EAGLEWorker does.
-        return self.target_worker.update_weights_from_distributed(recv_req)
-
     def clear_cache_pool(self):
         # The target worker owns the shared KV allocator/cache. For the compact
         # sliding-window path, the draft req->token view is rebuilt from committed
         # target state before each draft forward, so there is nothing persistent
         # to flush here.
         pass
 
+    def iter_draft_runners(self) -> list[tuple[str, "ModelRunner"]]:
+        # Explicit so DFlash's __getattr__ doesn't delegate this to the target
+        # (whose model_runner is aliased here) and miss the real draft.
+        return [("draft", self.draft_model_runner)]
+
     def _gather_req_to_token_masked(
         self,
         *,