diff --git a/.claude/skills/sglang-runtime-context/SKILL.md b/.claude/skills/sglang-runtime-context/SKILL.md index 10b6c547159d..95f699a6af10 100644 --- a/.claude/skills/sglang-runtime-context/SKILL.md +++ b/.claude/skills/sglang-runtime-context/SKILL.md @@ -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". diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 589d652329a9..09e838e5727a 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -868,31 +868,27 @@ def maybe_init_draft_worker(self): 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 ( diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index 0aeb6e9e8837..abf643b712f8 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/python/sglang/srt/speculative/draft_worker_common.py @@ -10,7 +10,7 @@ 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,7 +85,39 @@ 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( diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index f6e86eeb37ad..4a32b6f07f8e 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -264,12 +264,6 @@ def init_token_map(self): 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 @@ def __init__( 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, diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 2e251bfc1450..845036b26f86 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -679,12 +679,6 @@ def __init__( 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, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 3c8568feee9b..293bb176de30 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -907,12 +907,6 @@ def __init__( 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, diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index ce04a28efeb3..18fa75a658b1 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -150,12 +150,6 @@ def __init__( 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, diff --git a/test/registered/unit/spec/test_draft_server_args_copy.py b/test/registered/unit/spec/test_draft_server_args_copy.py new file mode 100644 index 000000000000..bda3db28072b --- /dev/null +++ b/test/registered/unit/spec/test_draft_server_args_copy.py @@ -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() diff --git a/test/registered/unit/spec/test_spec_worker_draft_isolation.py b/test/registered/unit/spec/test_spec_worker_draft_isolation.py new file mode 100644 index 000000000000..860f07d3df12 --- /dev/null +++ b/test/registered/unit/spec/test_spec_worker_draft_isolation.py @@ -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() diff --git a/test/registered/unit/test_server_args_writer_ratchet.py b/test/registered/unit/test_server_args_writer_ratchet.py index 3d2bde7e2ed4..f8d26bc21edc 100644 --- a/test/registered/unit/test_server_args_writer_ratchet.py +++ b/test/registered/unit/test_server_args_writer_ratchet.py @@ -49,7 +49,7 @@ "multimodal_gen", ) -_BASELINE = 31 +_BASELINE = 26 class TestServerArgsWriterRatchet(CustomTestCase):