diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8365066b5..4f023f7a8 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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 ( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6c6755469..cc029ae5b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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 = ( diff --git a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py index f3ada403b..582307804 100644 --- a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py @@ -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, diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 26d23c549..2e7a0bf53 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -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 diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index f958b3d3e..9c742a864 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -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: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 99bdfda2f..67eed9672 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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" diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index adc23c728..857178a6a 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/python/sglang/srt/speculative/draft_worker_common.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index e00b522d7..8c385b353 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -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 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 3ff8c8a28..c32402333 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -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() 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 5f085c17d..d1b347ce8 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -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 diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index f32d97522..d830b9848 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -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 diff --git a/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py b/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py index 9f78f7b08..f20223985 100644 --- a/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py +++ b/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py @@ -51,20 +51,16 @@ class TestChunkedPrefixCacheGate(CustomTestCase): get_context().set_server_args(sa) # what a later republish would do self.assertFalse(get_schedule().disable_chunked_prefix_cache) - def test_draft_copy_overrides_carry_the_gate(self): - # The draft copy comes from the pristine instance, which never sees - # the bag-only gate; the copy's pre-publish overrides carry it. - from types import SimpleNamespace - + def test_draft_variant_fields_carry_the_gate(self): + # Publishing the draft variant re-projects the bags from it, so the + # gate — which lives in the bags only — has to travel on the variant. from sglang.srt.speculative.draft_worker_common import ( draft_server_args_overrides, ) self._seed(attention_backend="triton") maybe_disable_chunked_prefix_cache(use_mla_backend=True, is_draft_worker=False) - fields = draft_server_args_overrides( - SimpleNamespace(context_len=64), draft_backend="fa3" - ) + fields = draft_server_args_overrides(draft_backend="fa3") self.assertTrue(fields["disable_chunked_prefix_cache"]) diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py index 515d7bd12..e52b08e0c 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py +++ b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py @@ -88,13 +88,14 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase): model_runner = SimpleNamespace( server_args=SimpleNamespace( chunked_prefill_size=16, - context_length=None, cuda_graph_config=SimpleNamespace( prefill=SimpleNamespace( full_prefill_prefix_chunk_tokens=None, max_bs=8 ) ), ), + # Wider than the token table, so the table is the binding limit. + model_config=SimpleNamespace(context_len=4096), req_to_token_pool=SimpleNamespace( req_to_token=torch.empty((1, 32), dtype=torch.int32) ), @@ -138,6 +139,17 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase): (1, 4), ) + # A context shorter than the token table binds instead: a draft runner + # capped at the target's context, or a short --context-length. + model_runner.model_config.context_len = 8 + model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = ( + 256 + ) + self.assertEqual( + PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4), + (8, 32), + ) + model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = ( 0 ) diff --git a/test/registered/unit/spec/test_draft_per_runner_config.py b/test/registered/unit/spec/test_draft_per_runner_config.py new file mode 100644 index 000000000..325ab1d79 --- /dev/null +++ b/test/registered/unit/spec/test_draft_per_runner_config.py @@ -0,0 +1,170 @@ +"""What differs for a draft worker is resolved per runner, not on a config copy. + +The v2 spec workers used to write the draft's `context_length` onto the +`ServerArgs` they share with the target, and the scheduler the draft's +`load_format`; then both moved to a published copy of the config. Neither is a +process-wide config change: the draft's context length is the target model's, +its load format is `--speculative-draft-load-format`, and both are consumed by +one constructor each — so they travel as arguments to the runner that owns them. +""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.model_executor.model_runner_components.load_model_utils import ( + build_load_config, +) +from sglang.srt.runtime_context import get_context, get_model +from sglang.srt.speculative.draft_worker_common import 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") + + +class _StopConstruction(Exception): + """Cuts the draft worker off once its ServerArgs is captured.""" + + +class TestDraftPerRunnerConfig(CustomTestCase): + def _seed(self, **fields): + override = get_context().override_server_args(**fields) + server_args = override.install() + self.addCleanup(override.restore) + return server_args + + # -- the draft load format is the draft runner's own resolved value -------- + + def _load_format_of(self, *, is_draft_worker: bool): + runner = ModelRunner.__new__(ModelRunner) + runner.is_draft_worker = is_draft_worker + return runner._resolve_draft_load_format() + + def test_the_draft_load_format_applies_to_the_draft_runner_only(self): + self._seed(load_format="auto", speculative_draft_load_format="dummy") + self.assertEqual(self._load_format_of(is_draft_worker=True), "dummy") + self.assertIsNone(self._load_format_of(is_draft_worker=False)) + + def test_an_unset_draft_load_format_leaves_the_load_config_alone(self): + self._seed(load_format="auto") + self.assertIsNone(self._load_format_of(is_draft_worker=True)) + + def test_the_draft_format_is_published_while_the_draft_loads(self): + """Model code reads the load format off the bag as it builds.""" + self._seed(load_format="auto", speculative_draft_load_format="dummy") + runner = ModelRunner.__new__(ModelRunner) + runner.is_draft_worker = True + + with runner._load_format_scope(runner._resolve_draft_load_format()): + self.assertEqual(get_model().load_format, "dummy") + self.assertEqual(get_model().load_format, "auto") + + def test_the_target_load_never_shifts_the_published_format(self): + self._seed(load_format="auto", speculative_draft_load_format="dummy") + runner = ModelRunner.__new__(ModelRunner) + runner.is_draft_worker = False + + with runner._load_format_scope(runner._resolve_draft_load_format()): + self.assertEqual(get_model().load_format, "auto") + + def test_the_transfer_engine_gate_answers_for_the_runner(self): + """The engine is initialized at the top of initialize(), long before the + weights load, so the gate has to see the draft's format.""" + server_args = self._seed( + load_format="auto", + speculative_draft_load_format="remote_instance", + remote_instance_weight_loader_backend="transfer_engine", + ) + self.assertFalse( + server_args.remote_instance_weight_loader_use_transfer_engine() + ) + self.assertTrue( + server_args.remote_instance_weight_loader_use_transfer_engine( + load_format="remote_instance" + ) + ) + + def test_the_load_config_takes_the_per_runner_format_when_given(self): + server_args = self._seed(load_format="auto") + common = dict( + server_args=server_args, + tp_rank=0, + remote_instance_weight_transporter_engine=None, + remote_instance_weight_transporter_session_id=None, + draft_model_idx=None, + weight_cache_mode="disable", + weight_cache_socket=None, + ) + self.assertEqual(build_load_config(**common).load_format, "auto") + self.assertEqual( + build_load_config(load_format="dummy", **common).load_format, "dummy" + ) + + # -- the variant left in the dflash / dspark path carries backends only ---- + + def test_the_draft_variant_carries_the_backend_family_only(self): + self._seed(disable_chunked_prefix_cache=False) + fields = draft_server_args_overrides("triton") + + self.assertEqual(fields["attention_backend"], "triton") + self.assertEqual(fields["speculative_draft_attention_backend"], "triton") + self.assertIsNone(fields["prefill_attention_backend"]) + self.assertIsNone(fields["decode_attention_backend"]) + self.assertNotIn("context_length", fields) + self.assertNotIn("load_format", fields) + self.assertNotIn("skip_tokenizer_init", fields) + + def test_the_variant_carries_the_targets_resolved_gate(self): + """Publishing the variant re-projects the bags, so the gate travels.""" + self._seed(disable_chunked_prefix_cache=False) + get_context().override("test.gate", disable_chunked_prefix_cache=True) + self.assertTrue( + draft_server_args_overrides("triton")["disable_chunked_prefix_cache"] + ) + + # -- the scheduler hands over the process's own config --------------------- + + def _scheduler(self, server_args, create_worker): + scheduler = Scheduler.__new__(Scheduler) + scheduler.server_args = server_args + scheduler.tp_worker = SimpleNamespace( + model_runner=SimpleNamespace(model_config=SimpleNamespace(context_len=4096)) + ) + scheduler.ps = SimpleNamespace(gpu_id=0) + scheduler.nccl_port = 0 + scheduler.spec_algorithm = SimpleNamespace( + is_none=lambda: False, + is_ngram=lambda: False, + create_worker=create_worker, + ) + return scheduler + + def test_the_draft_worker_and_its_factory_get_the_published_config(self): + server_args = self._seed(speculative_algorithm="EAGLE", load_format="auto") + seen = {} + + def worker_class(**kwargs): + seen["worker"] = kwargs["server_args"] + seen["published_while_building"] = get_model().load_format + raise _StopConstruction + + def create_worker(factory_server_args): + seen["factory"] = factory_server_args + return worker_class + + with self.assertRaises(_StopConstruction): + self._scheduler(server_args, create_worker).maybe_init_draft_worker() + + # No copy, and no publish switch: a registered algorithm picking its + # worker class from the config it is handed sees the same object the + # worker does, and the bags stay the target's throughout. + self.assertIs(seen["factory"], server_args) + self.assertIs(seen["worker"], server_args) + self.assertEqual(seen["published_while_building"], "auto") + self.assertIs(get_context().server_args, server_args) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/spec/test_draft_server_args_copy.py b/test/registered/unit/spec/test_draft_server_args_copy.py deleted file mode 100644 index bda3db280..000000000 --- a/test/registered/unit/spec/test_draft_server_args_copy.py +++ /dev/null @@ -1,84 +0,0 @@ -"""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 deleted file mode 100644 index 860f07d3d..000000000 --- a/test/registered/unit/spec/test_spec_worker_draft_isolation.py +++ /dev/null @@ -1,130 +0,0 @@ -"""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()