spec: build every draft worker from a draft ServerArgs copy (#33335)

EAGLEWorkerV2, StandaloneWorkerV2, MultiLayerEagleWorkerV2 and
FrozenKVMTPWorkerV2 wrote the draft's context_length onto the ServerArgs
instance they share with the target worker, and the scheduler wrote the draft's
load_format onto that same object just before creating them. The target's config
carried draft values from then on, and anything constructed later in the process
inherited them.

Scheduler.maybe_init_draft_worker now makes one draft copy through
draft_server_args_copy() and hands it to both the worker factory and the worker,
so every algorithm gets it — the four built-ins, dflash/dspark (which deepcopy
it again inside build_draft_tp_worker), and anything registered through
SpeculativeAlgorithm.register. The copy starts from the config the process
resolved, not from the pristine seed, so load-time overrides made before this
point (the chunked-prefix gate, the SM100 GDN prefill default) are part of what
the draft sees; context_length and load_format are applied on top.

The construction runs under a preserved publish of that copy, the shape
build_draft_tp_worker already used. Weight loading reads the bags rather than the
instance it was handed — Inkling's ModelOpt scale normalization keys on
load_format — so the draft has to be built with its own config published, and
the target's is back in the slot when construction returns.

The EAGLE hot-token-map write is deleted, not moved. init_token_map runs from
alloc_memory_pool, long after the draft's TpModelWorker built its ModelConfig,
and hot_vocab_size is only ever read off model_config.hf_config, which
json_model_override_args reaches at ModelConfig construction. The write could not
affect the draft model; only the shared instance saw it. hot_token_id is
unchanged, so a draft checkpoint that declares hot_vocab_size behaves as before.

Tests: draft_server_args_copy carries the target context_length, a configured
draft load_format and any load-time override while leaving the target's instance
alone; and the scheduler handoff pins that the factory and the worker both
receive the copy, that the copy is the published config during construction, and
that the target's is restored afterwards.

Writer ratchet 31 -> 26.
This commit is contained in:
Cheng Wan
2026-08-02 21:22:52 -07:00
committed by GitHub
parent ebb1c88d23
commit 9bc8848fcf
10 changed files with 275 additions and 53 deletions
+13 -17
View File
@@ -868,31 +868,27 @@ class Scheduler(
self.external_corpus_manager = None
return
from sglang.srt.speculative.draft_worker_common import (
draft_server_args_copy,
)
# Launch a draft worker for speculative decoding
draft_worker_kwargs = dict(
draft_server_args = draft_server_args_copy(
server_args=self.server_args,
target_model_config=self.tp_worker.model_runner.model_config,
)
draft_worker_kwargs = dict(
server_args=draft_server_args,
gpu_id=self.ps.gpu_id,
ps=self.ps,
nccl_port=self.nccl_port,
target_worker=self.tp_worker,
)
if get_spec().speculative_draft_load_format is not None:
# Write the draft load_format onto server_args (not just the bag):
# the draft worker is built from a copy of self.server_args and
# build_load_config reads server_args.load_format, so a bag-only
# override would be ignored and the draft would load in the target's
# format.
self.server_args.override(
"scheduler.draft_load_format",
load_format=get_spec().speculative_draft_load_format,
)
logger.info(
f"Using draft model load_format: '{get_spec().speculative_draft_load_format}'"
)
DraftWorkerClass = self.spec_algorithm.create_worker(self.server_args)
self.draft_worker = DraftWorkerClass(**draft_worker_kwargs)
DraftWorkerClass = self.spec_algorithm.create_worker(draft_server_args)
with get_context().preserve_config():
get_context().set_server_args(draft_server_args)
self.draft_worker = DraftWorkerClass(**draft_worker_kwargs)
if self.spec_algorithm.is_ngram():
from sglang.srt.speculative.external_corpus_manager import (
@@ -10,7 +10,7 @@ import torch
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
from sglang.srt.runtime_context import get_context, get_schedule
from sglang.srt.runtime_context import get_context, get_schedule, get_spec
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.dflash_info import DFlashVerifyInput
from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
@@ -61,6 +61,13 @@ def _resolve_draft_attention_backend_fallback(
return draft_backend
def _draft_load_format_fields() -> dict:
draft_load_format = get_spec().speculative_draft_load_format
if draft_load_format is None:
return {}
return dict(load_format=draft_load_format)
def draft_server_args_overrides(target_model_config, draft_backend) -> dict:
"""Pre-publish field adjustments for a draft ``ServerArgs`` copy.
@@ -78,9 +85,41 @@ def draft_server_args_overrides(target_model_config, draft_backend) -> dict:
attention_backend=draft_backend,
context_length=target_model_config.context_len,
disable_chunked_prefix_cache=get_schedule().disable_chunked_prefix_cache,
**_draft_load_format_fields(),
)
def draft_server_args_copy(server_args: ServerArgs, target_model_config) -> ServerArgs:
"""A draft-only ``ServerArgs`` for the workers that build their own draft.
Starts from the config the process resolved, not from the pristine seed:
the copy is published while the draft builds, and load-time overrides made
before this point (the chunked-prefix gate, the SM100 GDN prefill default)
are part of what the draft's layers must see. On top of that,
``context_length`` follows the target (the draft reads target KV) and
``load_format`` follows ``--speculative-draft-load-format``. The target's
own instance is untouched.
"""
draft_load_format = get_spec().speculative_draft_load_format
if draft_load_format is not None:
logger.info(f"Using draft model load_format: '{draft_load_format}'")
resolved = {}
for _source, fields in get_context().overrides_log():
resolved.update(fields)
draft_server_args = deepcopy(server_args)
draft_server_args.override(
"draft_worker.copy",
**{
**resolved,
"context_length": target_model_config.context_len,
**_draft_load_format_fields(),
},
)
return draft_server_args
def build_draft_tp_worker(
*,
server_args: ServerArgs,
@@ -264,12 +264,6 @@ class EagleDraftWorker(EagleDraftWorkerBase):
self.hot_token_id = None
elif get_spec().speculative_token_map is not None:
self.hot_token_id = load_token_map(get_spec().speculative_token_map)
self.server_args.override(
"eagle_worker.hot_token_map",
json_model_override_args=(
f'{{"hot_vocab_size": {len(self.hot_token_id)}}}'
),
)
else:
self.hot_token_id = None
@@ -1010,12 +1004,6 @@ class EAGLEWorkerV2(BaseSpecWorker):
server_args.speculative_algorithm
)
# Override the context length of the draft model to be the same as the target model.
server_args.override(
"spec_worker.match_target_context_length",
context_length=target_worker.model_runner.model_config.context_len,
)
self._draft_worker = EagleDraftWorker(
server_args,
gpu_id,
@@ -679,12 +679,6 @@ class FrozenKVMTPWorkerV2(EAGLEWorkerV2):
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
# Match the draft context length to the target (assistant reads target KV).
server_args.override(
"spec_worker.match_target_context_length",
context_length=target_worker.model_runner.model_config.context_len,
)
self._draft_worker = FrozenKVMTPDraftWorker(
server_args,
gpu_id,
@@ -907,12 +907,6 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
server_args.speculative_algorithm
)
# Override the context length of the draft model to be the same as the target model.
server_args.override(
"spec_worker.match_target_context_length",
context_length=target_worker.model_runner.model_config.context_len,
)
self._draft_worker = MultiLayerEagleDraftWorker(
server_args,
gpu_id,
@@ -150,12 +150,6 @@ class StandaloneWorkerV2(EAGLEWorkerV2):
server_args.speculative_algorithm
)
# Override the context length of the draft model to be the same as the target model.
server_args.override(
"spec_worker.match_target_context_length",
context_length=target_worker.model_runner.model_config.context_len,
)
# Create our custom draft worker that doesn't share embeddings/lm_head
self._draft_worker = StandaloneDraftWorker(
server_args,