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
@@ -237,10 +237,13 @@ Never module-skip a test "until the migration settles" — seed the context inst
|
||||
## Hard-won pitfalls (check these before/while refactoring)
|
||||
|
||||
- **Moving code drops first-line guards**: early returns (`if self.is_draft_worker: return`)
|
||||
are the easiest thing to lose when relocating a method body. Only drafts built through
|
||||
`build_draft_tp_worker()` get private bags (a preserved publish of the rewritten copy);
|
||||
drafts constructed directly with `is_draft_worker=True` skip publish and **share the
|
||||
target's bags** — a draft-side write there poisons the target.
|
||||
are the easiest thing to lose when relocating a method body. Every draft is built
|
||||
under a preserved publish of its own config: the scheduler makes the copy with
|
||||
`draft_server_args_copy()` (seeded from the resolved config, so load-time overrides
|
||||
carry) and publishes it around the worker factory, and `build_draft_tp_worker()`
|
||||
nests the same shape for dflash / dspark. The publish ends when construction does —
|
||||
anything the draft reads later (`alloc_memory_pool`, `init_attention_backends`,
|
||||
cuda-graph capture) is back on the target's bags.
|
||||
- **Registry-completeness timing**: a gate that consults an extensible list is only correct
|
||||
after the registrars ran (platform `init_backend()` at module import). See "load-time vs
|
||||
resolution-time".
|
||||
|
||||
@@ -868,30 +868,26 @@ 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)
|
||||
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():
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
"""The draft's ServerArgs is a copy; the target's stays as the launcher left it.
|
||||
|
||||
Regression: the v2 spec workers wrote the draft's context_length (and the
|
||||
scheduler the draft's load_format) onto the ServerArgs instance they share with
|
||||
the target worker, so every later reader of that instance saw draft values.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.srt.speculative.draft_worker_common import (
|
||||
draft_server_args_copy,
|
||||
draft_server_args_overrides,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
TARGET_MODEL_CONFIG = SimpleNamespace(context_len=4096)
|
||||
|
||||
|
||||
class TestDraftServerArgsCopy(CustomTestCase):
|
||||
def _seed(self, **fields):
|
||||
override = get_context().override_server_args(**fields)
|
||||
server_args = override.install()
|
||||
self.addCleanup(override.restore)
|
||||
return server_args
|
||||
|
||||
def test_the_draft_context_length_follows_the_target(self):
|
||||
target = self._seed(context_length=None)
|
||||
draft = draft_server_args_copy(target, TARGET_MODEL_CONFIG)
|
||||
self.assertEqual(draft.context_length, 4096)
|
||||
|
||||
def test_the_target_instance_is_left_alone(self):
|
||||
target = self._seed(context_length=None, load_format="auto")
|
||||
draft = draft_server_args_copy(target, TARGET_MODEL_CONFIG)
|
||||
self.assertIsNot(draft, target)
|
||||
self.assertIsNone(target.context_length)
|
||||
self.assertEqual(target.load_format, "auto")
|
||||
|
||||
def test_the_draft_load_format_applies_only_when_configured(self):
|
||||
target = self._seed(load_format="auto", speculative_draft_load_format="dummy")
|
||||
self.assertEqual(
|
||||
draft_server_args_copy(target, TARGET_MODEL_CONFIG).load_format, "dummy"
|
||||
)
|
||||
self.assertEqual(target.load_format, "auto")
|
||||
|
||||
target = self._seed(load_format="auto")
|
||||
self.assertEqual(
|
||||
draft_server_args_copy(target, TARGET_MODEL_CONFIG).load_format, "auto"
|
||||
)
|
||||
|
||||
def test_load_time_overrides_reach_the_draft(self):
|
||||
target = self._seed(disable_chunked_prefix_cache=False)
|
||||
# What the target runner resolved before the draft is built — e.g. the
|
||||
# chunked-prefix gate for an attention backend that cannot serve it.
|
||||
get_context().override("test.gate", disable_chunked_prefix_cache=True)
|
||||
|
||||
draft = draft_server_args_copy(target, TARGET_MODEL_CONFIG)
|
||||
self.assertTrue(draft.disable_chunked_prefix_cache)
|
||||
self.assertFalse(target.disable_chunked_prefix_cache)
|
||||
|
||||
def test_the_draft_specific_fields_win_over_the_resolved_ones(self):
|
||||
target = self._seed(context_length=None, load_format="auto")
|
||||
get_context().override("test.late", context_length=128, load_format="npcache")
|
||||
|
||||
draft = draft_server_args_copy(target, TARGET_MODEL_CONFIG)
|
||||
self.assertEqual(draft.context_length, 4096)
|
||||
|
||||
def test_the_built_draft_overrides_carry_the_load_format_too(self):
|
||||
self._seed(speculative_draft_load_format="dummy")
|
||||
fields = draft_server_args_overrides(TARGET_MODEL_CONFIG, "triton")
|
||||
self.assertEqual(fields["load_format"], "dummy")
|
||||
|
||||
self._seed()
|
||||
self.assertNotIn(
|
||||
"load_format", draft_server_args_overrides(TARGET_MODEL_CONFIG, "triton")
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,130 @@
|
||||
"""The scheduler hands every draft worker a ServerArgs copy.
|
||||
|
||||
Regression: the v2 spec workers wrote the draft's context_length onto the
|
||||
instance they share with the target worker, and the scheduler wrote the draft's
|
||||
load_format onto that same object, so the target's config carried draft values
|
||||
for the rest of the process. The copy is made once, before the worker factory,
|
||||
so plugin algorithms registered through SpeculativeAlgorithm.register get it too.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class _StopConstruction(Exception):
|
||||
"""Cuts the draft worker off once its ServerArgs is captured."""
|
||||
|
||||
|
||||
def _scheduler(server_args):
|
||||
scheduler = Scheduler.__new__(Scheduler)
|
||||
scheduler.server_args = server_args
|
||||
model_config = SimpleNamespace(context_len=4096)
|
||||
scheduler.tp_worker = SimpleNamespace(
|
||||
model_runner=SimpleNamespace(model_config=model_config)
|
||||
)
|
||||
scheduler.ps = SimpleNamespace(gpu_id=0)
|
||||
scheduler.nccl_port = 0
|
||||
return scheduler
|
||||
|
||||
|
||||
class TestSchedulerDraftServerArgs(CustomTestCase):
|
||||
def _seed(self, **fields):
|
||||
override = get_context().override_server_args(
|
||||
speculative_algorithm="EAGLE", **fields
|
||||
)
|
||||
server_args = override.install()
|
||||
self.addCleanup(override.restore)
|
||||
return server_args
|
||||
|
||||
def _captured_draft_args(self, server_args):
|
||||
seen = {}
|
||||
|
||||
def worker_class(**kwargs):
|
||||
seen["server_args"] = kwargs["server_args"]
|
||||
raise _StopConstruction
|
||||
|
||||
scheduler = _scheduler(server_args)
|
||||
scheduler.spec_algorithm = SimpleNamespace(
|
||||
is_none=lambda: False,
|
||||
is_ngram=lambda: False,
|
||||
create_worker=lambda _sa: worker_class,
|
||||
)
|
||||
with self.assertRaises(_StopConstruction):
|
||||
scheduler.maybe_init_draft_worker()
|
||||
return seen["server_args"]
|
||||
|
||||
def test_the_draft_gets_a_copy_carrying_the_target_context_length(self):
|
||||
server_args = self._seed(context_length=None)
|
||||
draft = self._captured_draft_args(server_args)
|
||||
self.assertIsNot(draft, server_args)
|
||||
self.assertEqual(draft.context_length, 4096)
|
||||
self.assertIsNone(server_args.context_length)
|
||||
|
||||
def test_the_draft_config_is_published_while_the_draft_is_built(self):
|
||||
from sglang.srt.runtime_context import get_model
|
||||
|
||||
server_args = self._seed(
|
||||
load_format="auto", speculative_draft_load_format="dummy"
|
||||
)
|
||||
seen = {}
|
||||
|
||||
def worker_class(**kwargs):
|
||||
seen["published"] = get_model().load_format
|
||||
raise _StopConstruction
|
||||
|
||||
scheduler = _scheduler(server_args)
|
||||
scheduler.spec_algorithm = SimpleNamespace(
|
||||
is_none=lambda: False,
|
||||
is_ngram=lambda: False,
|
||||
create_worker=lambda _sa: worker_class,
|
||||
)
|
||||
with self.assertRaises(_StopConstruction):
|
||||
scheduler.maybe_init_draft_worker()
|
||||
|
||||
# Model-level weight loading reads the bags, not the instance it was
|
||||
# handed, so the draft's config has to be the published one while it
|
||||
# builds — and the target's has to be back afterwards.
|
||||
self.assertEqual(seen["published"], "dummy")
|
||||
self.assertEqual(get_model().load_format, "auto")
|
||||
|
||||
def test_the_worker_factory_sees_the_draft_config(self):
|
||||
server_args = self._seed(
|
||||
load_format="auto", speculative_draft_load_format="dummy"
|
||||
)
|
||||
seen = {}
|
||||
|
||||
def create_worker(factory_server_args):
|
||||
seen["load_format"] = factory_server_args.load_format
|
||||
raise _StopConstruction
|
||||
|
||||
scheduler = _scheduler(server_args)
|
||||
scheduler.spec_algorithm = SimpleNamespace(
|
||||
is_none=lambda: False,
|
||||
is_ngram=lambda: False,
|
||||
create_worker=create_worker,
|
||||
)
|
||||
with self.assertRaises(_StopConstruction):
|
||||
scheduler.maybe_init_draft_worker()
|
||||
|
||||
# A registered algorithm may pick its worker class from the config it
|
||||
# is handed, so the factory and the worker must see the same one.
|
||||
self.assertEqual(seen["load_format"], "dummy")
|
||||
|
||||
def test_a_configured_draft_load_format_never_reaches_the_target(self):
|
||||
server_args = self._seed(
|
||||
load_format="auto", speculative_draft_load_format="dummy"
|
||||
)
|
||||
draft = self._captured_draft_args(server_args)
|
||||
self.assertEqual(draft.load_format, "dummy")
|
||||
self.assertEqual(server_args.load_format, "auto")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -49,7 +49,7 @@ _EXCLUDED = (
|
||||
"multimodal_gen",
|
||||
)
|
||||
|
||||
_BASELINE = 31
|
||||
_BASELINE = 26
|
||||
|
||||
|
||||
class TestServerArgsWriterRatchet(CustomTestCase):
|
||||
|
||||
Reference in New Issue
Block a user