From ca0ee3f1a84ea97f0a06036d30e854d0f0ded37a Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 14 Jul 2026 18:39:46 -0700 Subject: [PATCH] [Spec] Consolidate spec-worker weight updates into BaseSpecWorker via draft_runners (#31078) --- .../srt/speculative/base_spec_worker.py | 28 ++++++++++++++++ .../sglang/srt/speculative/eagle_worker_v2.py | 28 +--------------- .../multi_layer_eagle_worker_v2.py | 33 ++++--------------- 3 files changed, 36 insertions(+), 53 deletions(-) diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index 4268fea18..453039de2 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/python/sglang/srt/speculative/base_spec_worker.py @@ -13,6 +13,10 @@ if _is_cpu: from sgl_kernel import assign_draft_cache_locs_contiguous_cpu if TYPE_CHECKING: + from sglang.srt.managers.io_struct import ( + UpdateWeightFromDiskReqInput, + UpdateWeightsFromIPCReqInput, + ) from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.mem_cache.memory_pool import ReqToTokenPool @@ -82,6 +86,12 @@ class EagleDraftWorkerBase(ABC): def draft_extend(): pass + @property + def draft_runners(self) -> list[ModelRunner]: + """All draft model runners; multi-layer eagle overrides with its + per-step runner list.""" + return [self.draft_runner] + def alloc_memory_pool(self, **kwargs): pass @@ -343,6 +353,24 @@ class BaseSpecWorker(ABC): if self.draft_worker is not None: self.draft_worker.init_cuda_graphs() + def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): + for runner in self.draft_worker.draft_runners: + success, message = runner.weight_updater.update_weights_from_disk( + recv_req.model_path, + recv_req.load_format, + recapture_cuda_graph=recv_req.recapture_cuda_graph, + ) + if not success: + return success, message + return True, "Succeeded to update model weights." + + def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput): + for runner in self.draft_worker.draft_runners: + success, message = runner.weight_updater.update_weights_from_ipc(recv_req) + if not success: + return success, message + return True, "Succeeded to update model weights." + def on_verify_complete_cpu( self, num_correct_drafts_per_req: list[int], batch_size: int = 0 ) -> None: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index c8e575c98..93260868e 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -29,11 +29,7 @@ from sglang.srt.layers.moe.utils import ( speculative_moe_backend_context, ) from sglang.srt.layers.utils.logprob import compute_spec_v2_logprobs -from sglang.srt.managers.io_struct import ( - UpdateWeightFromDiskReqInput, - UpdateWeightsFromIPCReqInput, - UpdateWeightsFromTensorReqInput, -) +from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker @@ -1719,28 +1715,6 @@ class EAGLEWorkerV2(BaseSpecWorker): out.view(bs, nd, *x.shape[1:])[:, :s1] = gathered.view(bs, s1, *x.shape[1:]) return out - def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): - success, message = ( - self._draft_worker.draft_runner.weight_updater.update_weights_from_disk( - recv_req.model_path, - recv_req.load_format, - recapture_cuda_graph=recv_req.recapture_cuda_graph, - ) - ) - if not success: - return success, message - return True, "Succeeded to update model weights." - - def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput): - success, message = ( - self._draft_worker.draft_runner.weight_updater.update_weights_from_ipc( - recv_req - ) - ) - if not success: - return success, message - return True, "Succeeded to update model weights." - def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): monkey_patch_torch_reductions() named_tensors = MultiprocessingSerializer.deserialize( diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 6ecf96708..85c60483f 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -12,6 +12,8 @@ # limitations under the License. # ============================================================================== +from __future__ import annotations + import logging from dataclasses import replace from typing import TYPE_CHECKING, List @@ -26,10 +28,6 @@ from sglang.srt.hardware_backend.npu.graph_runner.multi_layer_eagle_draft_extend ) from sglang.srt.layers.moe.utils import speculative_moe_backend_context from sglang.srt.layers.utils.logprob import compute_spec_v2_logprobs -from sglang.srt.managers.io_struct import ( - UpdateWeightFromDiskReqInput, - UpdateWeightsFromIPCReqInput, -) from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker @@ -154,6 +152,11 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): self.tree_mask_mode = default_tree_mask_mode() self.plan_stream, self.plan_stream_ctx = get_plan_stream(self.device) + @property + def draft_runners(self) -> List[ModelRunner]: + # One runner per draft step (len == speculative_num_steps). + return self.draft_runner_list + def alloc_memory_pool( self, memory_pool_config=None, @@ -863,25 +866,3 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): indexer_topk_output=forward_batch_output.indexer_topk_output, extra_keep_alive_refs=[verify_forward_batch], ) - - def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): - for i in range(self.speculative_num_steps): - success, message = self._draft_worker.draft_runner_list[ - i - ].weight_updater.update_weights_from_disk( - recv_req.model_path, - recv_req.load_format, - recapture_cuda_graph=recv_req.recapture_cuda_graph, - ) - if not success: - return success, message - return True, "Succeeded to update model weights." - - def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput): - for i in range(self.speculative_num_steps): - success, message = self._draft_worker.draft_runner_list[ - i - ].weight_updater.update_weights_from_ipc(recv_req) - if not success: - return success, message - return True, "Succeeded to update model weights."