config: resolve the draft worker's config per runner, not on a copy

The v2 spec workers got a published `ServerArgs` copy carrying two values: the
target's context length and `--speculative-draft-load-format`. Neither is a
process-wide config change — each is consumed by exactly one constructor — so
the copy, the publish switch around the draft build, and the replay of the
target's resolved overrides onto it all go away, and the values travel to the
runner that owns them:

- **Context length.** `TpModelWorker` already takes it (`context_length=None`
  keeps `server_args.context_length`); the four v2 draft workers and
  `build_draft_tp_worker` pass the target's, which every one of them has in
  scope as `target_worker` / `target_model_config`.
- **Load format.** `ModelRunner._draft_load_format()` resolves it for a draft
  runner and `build_load_config` takes it, so the `LoadConfig` is per-runner.
  Model code also reads it off the bag while it builds — Inkling replaces
  per-element noise in its shared-expert scales under dummy loading — so the
  load is wrapped in a scoped bag override that puts the target's value back.
- `skip_tokenizer_init` was on the copy for nobody: `TpModelWorker` already
  short-circuits the tokenizer for a draft worker (`or self.is_draft_worker`).

`PrefillCudaGraphRunner._max_addressable_prefix_len` capped the prefix by
`server_args.context_length`, which the copy used to carry for the draft; it now
reads the runner's own `model_config.context_len`. That is also more accurate for
the target, whose `--context-length` may be unset while the resolved context is
shorter than the token table.

What stays a variant is the dflash/dspark path's attention backend: backend
selection reads it off the config object the draft runner holds, and the
resolved gate has to survive the variant's publish. `draft_server_args_overrides`
now carries only those fields and says why.
This commit is contained in:
Cheng Wan
2026-08-05 19:31:23 -07:00
committed by GitHub
parent 99cfc90658
commit 64eeb153df
16 changed files with 281 additions and 304 deletions
+7 -14
View File
@@ -911,27 +911,20 @@ 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_server_args = draft_server_args_copy(
server_args=self.server_args,
target_model_config=self.tp_worker.model_runner.model_config,
)
# Launch a draft worker for speculative decoding. It builds its draft
# from this process's own config: what differs for the draft — the
# target's context length, the draft load format, its attention backend
# — is resolved per runner, not on a config copy.
draft_worker_kwargs = dict(
server_args=draft_server_args,
server_args=self.server_args,
gpu_id=self.ps.gpu_id,
ps=self.ps,
nccl_port=self.nccl_port,
target_worker=self.tp_worker,
)
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)
DraftWorkerClass = self.spec_algorithm.create_worker(self.server_args)
self.draft_worker = DraftWorkerClass(**draft_worker_kwargs)
if self.spec_algorithm.is_ngram():
from sglang.srt.speculative.external_corpus_manager import (
@@ -176,6 +176,7 @@ from sglang.srt.runtime_context import (
get_model,
get_parallel,
get_schedule,
get_spec,
set_global_dwdp_manager,
)
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
@@ -292,6 +293,10 @@ class ModelRunner:
self.dist_port = nccl_port
self.server_args = server_args
self.is_draft_worker = is_draft_worker
# This runner's own load format, resolved before anything keys off it:
# the remote-instance transfer engine is initialized at the top of
# initialize(), long before the weights are loaded.
self.draft_load_format = self._resolve_draft_load_format()
self.is_generation = model_config.is_generation
self.device_timer = None
self.is_multimodal = model_config.is_multimodal
@@ -639,7 +644,9 @@ class ModelRunner:
)
def maybe_init_remote_instance_transfer_engine(self):
if self.server_args.remote_instance_weight_loader_use_transfer_engine():
if self.server_args.remote_instance_weight_loader_use_transfer_engine(
load_format=self.draft_load_format
):
self.remote_instance_weight_transporter.init_engine()
def maybe_init_expert_location_metadata(self):
@@ -1027,8 +1034,10 @@ class ModelRunner:
set_cuda_arch()
draft_load_format = self.draft_load_format
self.load_config = build_load_config(
server_args=self.server_args,
load_format=draft_load_format,
tp_rank=self.ps.tp_rank,
remote_instance_weight_transporter_engine=self.remote_instance_weight_transporter.engine,
remote_instance_weight_transporter_session_id=self.remote_instance_weight_transporter.session_id,
@@ -1055,15 +1064,16 @@ class ModelRunner:
server_args=self.server_args, tp_rank=self.ps.tp_rank
)
loaded = load_model_with_memory_saver(
server_args=self.server_args,
model_config=self.model_config,
load_config=self.load_config,
device=self.device,
gpu_id=self.gpu_id,
memory_saver_adapter=self.memory_saver_adapter,
is_draft_worker=self.is_draft_worker,
)
with self._load_format_scope(draft_load_format):
loaded = load_model_with_memory_saver(
server_args=self.server_args,
model_config=self.model_config,
load_config=self.load_config,
device=self.device,
gpu_id=self.gpu_id,
memory_saver_adapter=self.memory_saver_adapter,
is_draft_worker=self.is_draft_worker,
)
self.loader = loaded.loader
self.model = loaded.model
if loaded.remote_instance_weight_info is not None:
@@ -1200,6 +1210,31 @@ class ModelRunner:
else:
return self.max_total_num_tokens
def _load_format_scope(self, load_format: Optional[str]):
"""Make this runner's load format the published one while it loads.
Model code reads it off the bag during construction (Inkling replaces
per-element noise in its shared-expert scales under dummy loading), so a
draft loading a different way than the target needs its own value live
for the load, and the target's back afterwards.
"""
if load_format is None:
return contextlib.nullcontext()
return get_model().override(load_format=load_format)
def _resolve_draft_load_format(self) -> Optional[str]:
"""``--speculative-draft-load-format``, for a draft runner only.
The draft loads its own checkpoint, so its load format is this runner's
own resolved value; the target keeps ``--load-format``.
"""
if not self.is_draft_worker:
return None
load_format = get_spec().speculative_draft_load_format
if load_format is not None:
logger.info(f"Using draft model load_format: '{load_format}'")
return load_format
def configure_kv_cache_dtype(self):
spec_algorithm = getattr(self, "spec_algorithm", None)
resolved_kv_cache_dtype, self.kv_cache_dtype = (
@@ -79,10 +79,13 @@ def maybe_downgrade_dtype_for_legacy_gpu(
def maybe_trigger_remote_instance_nccl_send_group(
*, server_args: ServerArgs, tp_rank: int
*, server_args: ServerArgs, tp_rank: int, load_format: Optional[str] = None
) -> None:
"""``load_format`` is this runner's effective format: a draft loading under
``--speculative-draft-draft-load-format`` needs its own send group, and the
target's format cannot answer for it."""
if (
server_args.load_format == LoadFormat.REMOTE_INSTANCE
(load_format or server_args.load_format) == LoadFormat.REMOTE_INSTANCE
and server_args.remote_instance_weight_loader_backend
== RemoteInstanceWeightLoaderBackend.NCCL
):
@@ -184,6 +187,7 @@ def build_load_config(
*,
server_args: ServerArgs,
tp_rank: int,
load_format: Optional[str] = None,
remote_instance_weight_transporter_engine: Any,
remote_instance_weight_transporter_session_id: str,
draft_model_idx: Optional[int],
@@ -201,7 +205,7 @@ def build_load_config(
)
return LoadConfig(
load_format=server_args.load_format,
load_format=load_format or server_args.load_format,
download_dir=server_args.download_dir,
model_loader_extra_config=server_args.model_loader_extra_config,
tp_rank=tp_rank,
@@ -777,7 +777,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
@staticmethod
def _max_addressable_prefix_len(model_runner) -> int:
table_width = model_runner.req_to_token_pool.req_to_token.shape[1]
configured_context = model_runner.server_args.context_length
# This runner's own resolved context (a draft runs at the target's
# positions), not the launcher's flag, which may be unset.
configured_context = model_runner.model_config.context_len
return (
min(table_width, configured_context)
if configured_context is not None and configured_context > 0
+7 -2
View File
@@ -636,8 +636,13 @@ class _ConfigBag:
@contextmanager
def override(self, **kwargs):
"""Scoped, transactional test-only override of this bag's own leaves
(keys validated before any write; restored on exit)."""
"""Scoped, transactional override of this bag's own leaves (keys
validated before any write; restored on exit).
For a window where one runner's value differs from the process's — a
draft model loading under ``--speculative-draft-load-format`` while the
target keeps ``--load-format`` — and for tests forcing a code path.
A permanent change goes through ``get_context().override``."""
fields = object.__getattribute__(self, "_fields")
unknown = set(kwargs) - set(fields)
if unknown:
+4 -2
View File
@@ -9194,12 +9194,14 @@ class ServerArgs:
"""Transport backend for modelexpress."""
return self._parsed_modelexpress_config.get("transport", "nixl")
def remote_instance_weight_loader_use_transfer_engine(self):
def remote_instance_weight_loader_use_transfer_engine(self, load_format=None):
"""``load_format`` overrides the seed's: a draft runner loading under
``--speculative-draft-load-format`` needs its own transfer engine."""
# Use TransferEngine as seed backend.
if self.remote_instance_weight_loader_start_seed_via_transfer_engine:
return True
# Use TransferEngine as client backend.
if self.load_format == "remote_instance" and (
if (load_format or self.load_format) == "remote_instance" and (
self.remote_instance_weight_loader_backend == "transfer_engine"
or (
self.remote_instance_weight_loader_backend == "modelexpress"
@@ -9,7 +9,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, get_spec
from sglang.srt.runtime_context import get_context, get_schedule
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.dflash_info import DFlashVerifyInput
from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
@@ -63,60 +63,23 @@ 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(draft_backend) -> dict:
"""The fields a draft variant must carry: its attention backend.
def draft_server_args_overrides(target_model_config, draft_backend) -> dict:
"""Pre-publish field adjustments for a draft ``ServerArgs`` copy.
Downstream draft-worker logic keys on ``speculative_draft_attention_backend``
(backend selection in ``_get_attention_backend``, the fa4-draft KV dtype
override in ``configure_kv_cache_dtype``); ``context_length`` keeps the
draft aligned with the target; ``disable_chunked_prefix_cache`` is the
target's resolved gate (bag-only, absent from the pristine copy).
Backend selection reads them off the config object the draft runner holds --
``speculative_draft_attention_backend`` in ``_resolve_attention_backend_strs``
and ``configure_kv_cache_dtype``, ``attention_backend`` in the non-hybrid
branch of the backend build, and the split pair must not shadow either with
the target's. ``disable_chunked_prefix_cache`` is the target's resolved gate,
which lives in the bags only: publishing the variant re-projects the bags
from it, so the value has to travel on the variant.
"""
return dict(
skip_tokenizer_init=True,
speculative_draft_attention_backend=draft_backend,
prefill_attention_backend=None,
decode_attention_backend=None,
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)
return server_args.derive(
"draft_worker.copy",
**{
**resolved,
"context_length": target_model_config.context_len,
**_draft_load_format_fields(),
},
)
@@ -139,8 +102,7 @@ def build_draft_tp_worker(
)
)
draft_server_args = server_args.derive(
"draft_worker.build",
**draft_server_args_overrides(target_model_config, draft_backend),
"draft_worker.build", **draft_server_args_overrides(draft_backend)
)
# The draft's layers must resolve config from the draft's own bags.
@@ -152,6 +114,8 @@ def build_draft_tp_worker(
ps=ps,
nccl_port=nccl_port,
is_draft_worker=True,
# The draft runs at absolute target positions.
context_length=target_model_config.context_len,
)
draft_model_runner = draft_worker.model_runner
@@ -169,6 +169,8 @@ class EagleDraftWorker(EagleDraftWorkerBase):
ps=replace(ps, pp_rank=0),
nccl_port=nccl_port,
is_draft_worker=True,
# The draft runs at absolute target positions.
context_length=target_worker.model_runner.model_config.context_len,
)
# Alias for better readability
@@ -138,6 +138,8 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
ps=replace(ps, pp_rank=0),
nccl_port=nccl_port,
is_draft_worker=True,
# The draft runs at absolute target positions.
context_length=self.target_worker.model_runner.model_config.context_len,
)
embed, head = self.target_worker.model_runner.model.get_embed_and_head()
@@ -159,6 +159,8 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
nccl_port=nccl_port,
is_draft_worker=True,
is_multi_layer_eagle=True,
# The draft runs at absolute target positions.
context_length=target_worker.model_runner.model_config.context_len,
)
# Alias for better readability
@@ -75,6 +75,8 @@ class StandaloneDraftWorker(EagleDraftWorker):
ps=replace(ps, pp_rank=0),
nccl_port=nccl_port,
is_draft_worker=True,
# The draft runs at absolute target positions.
context_length=target_worker.model_runner.model_config.context_len,
)
# Alias for better readability