refactor: centralize speculative weight-update fan-out in the scheduler - #28548
Merged
Conversation
guapisolo
requested review from
Qiaolin-Yu,
Ying1123,
hnyls2002,
kpham-sgl,
merrymercy and
xiezhq-hermann
June 17, 2026 18:29
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
guapisolo
force-pushed
the
refactor/spec-update-weights-fanout
branch
from
June 17, 2026 19:26
fd31e41 to
7c1f636
Compare
guapisolo
requested review from
CatherineSue,
Fridge003,
JustinTong0323,
ispobock,
lifuhuang,
slin1237 and
yushengsu-thu
June 17, 2026 19:26
…duler Replace the per-worker update_weights_from_distributed / _from_tensor forwarding added in #27749 with a single scheduler-side fan-out. The scheduler resolves the runner list via get_model_runners(selector) (target + drafts via each worker's iter_draft_runners); the target runner owns the distributed update group, receives the broadcast once, and loads the weights into every selected runner. The tensor path deserializes once on the scheduler and loads into each selected runner. Delete the per-worker forwarding methods in the 7 speculative workers; they now only expose iter_draft_runners (TpModelWorker provides the base implementation; draft-capable workers override it). Replace the ad-hoc disable_draft_model field on the distributed/tensor request structs with the unified {target,draft} selector (default both; "target" = target-only); no caller sent disable_draft_model. This also closes the DFlash stale-draft gap — its iter_draft_runners now fans the update out to the draft runner. Rewrite test_distributed_weight_update_spec_worker.py for the new scheduler fan-out contract. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
guapisolo
force-pushed
the
refactor/spec-update-weights-fanout
branch
from
June 17, 2026 19:56
7c1f636 to
cc58960
Compare
The fan-out commit routed update_weights_from_distributed through ModelRunner.update_weights_from_distributed_to_model_runners(model_runners, ...), which made one ModelRunner iterate over sibling runners and load into each — runner-set enumeration leaking into the runner layer (tech debt from #27749). Move the loop up to the scheduler so ModelRunner exposes only single-runner ops, symmetric with the tensor path. The target (main) model owns this process's update group, so the scheduler has it receive the broadcast once (receive_weights_from_distributed) and then loads the received weights into each selected runner (load_weights). Delete the multi-runner update_weights_from_distributed_to_model_runners, the now-dead single ModelRunner.update_weights_from_distributed, and the caller-less BaseTpWorker.update_weights_from_{distributed,tensor} forwarders left by the fan-out commit. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
yueming-yuan
approved these changes
Jun 17, 2026
This was referenced Jun 17, 2026
guapisolo
added a commit
to guapisolo/sglang
that referenced
this pull request
Jun 18, 2026
Re-implement distributed weight-update for speculative draft worker(s) on the v2 worker layout, on top of the preceding revert of sgl-project#27749's earlier design. Centralize the weight-update fan-out in SchedulerWeightUpdaterManager: an UpdateWeightsFrom{Distributed,Tensor}ReqInput `selector` ({target, draft, both}) resolves to the (role, ModelRunner) pairs to update. For distributed updates the target model receives the broadcast once and the weights are loaded into every selected runner locally; tensor updates deserialize once and fan out the same way. Draft workers expose their runner(s) via iter_draft_runners(); the per-worker update_weights_from_{distributed,tensor} methods and the `disable_draft_model` flag are gone. Combines sglang-miles sgl-project#27749 and sgl-project#28548. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Centralize speculative weight-update fan-out in the scheduler via
get_model_runners.Motivation
update_weights_from_distributed/update_weights_from_tensordispatched toself.draft_worker or self.tp_worker, and each of the 7 speculative workers re-implemented forwarding that hard-coded which runners to update — duplicating the runner-set knowledgeiter_draft_runners()already encodes (introduced in #27749). The online weight-update behavior for the target and draft runners is unchanged; only where runner selection lives changes.Before / After
eagle_worker,eagle_worker_v2,multi_layer_eagle_worker,multi_layer_eagle_worker_v2,frozen_kv_mtp_worker,ngram_worker,dflash_worker) built its own[draft..., target]runner list forupdate_weights_from_distributed_to_model_runners.update_weights_from_distributed/_from_tensorinSchedulerUpdateWeightsMixinresolve the runner list once viaget_model_runners(selector).model_runner.update_weights_from_distributed_to_model_runnersreceives the broadcast once for all selected runners.serialized_named_tensors[self.tp_rank]once for all selected runners.update_weights_from_distributed/update_weights_from_tensormethods.iter_draft_runners().disable_draft_modelon the twoUpdateWeights*ReqInputstructs becomes the unified{target, draft}selector(default both).Behavior Preservation
test_distributed_weight_update_spec_worker.pyasserts one distributed receive loads into target plus draft.target-only selector skips the draft worker.test_check_weights_draft_fanout.py,test_weight_checker.py) still pass.Verification
_cuda5, distributed-spec-worker 3, weight-checker 63).ruff --select F401,F821,black --check,isort --checkclean on the touched files.Review Focus
update_weights_from_tensorinscheduler_update_weights_mixin.pyfor the once-only deserialize across selected runners.disable_draft_modeltoselectorswap on the request structs.get_model_runnerskeeps the target first so the target owns the distributed update group.CI States
Latest PR Test (Base): ❌ Missing
run-cilabel -- add it to run CI tests.Latest PR Test (Extra): ❌ Blocked --
run-ciis required first.