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:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user