From 64eeb153dfaa5f6080039ce76bfd2bfe2a7097b1 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:31:23 -0700 Subject: [PATCH] config: resolve the draft worker's config per runner, not on a copy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- python/sglang/srt/managers/scheduler.py | 21 +-- .../sglang/srt/model_executor/model_runner.py | 55 ++++-- .../load_model_utils.py | 10 +- .../runner/prefill_cuda_graph_runner.py | 4 +- python/sglang/srt/runtime_context.py | 9 +- python/sglang/srt/server_args.py | 6 +- .../srt/speculative/draft_worker_common.py | 62 ++----- .../sglang/srt/speculative/eagle_worker_v2.py | 2 + .../speculative/frozen_kv_mtp_worker_v2.py | 2 + .../multi_layer_eagle_worker_v2.py | 2 + .../srt/speculative/standalone_worker_v2.py | 2 + .../test_chunked_prefix_cache_gate.py | 12 +- .../test_prefill_cuda_graph_runner.py | 14 +- .../unit/spec/test_draft_per_runner_config.py | 170 ++++++++++++++++++ .../unit/spec/test_draft_server_args_copy.py | 84 --------- .../spec/test_spec_worker_draft_isolation.py | 130 -------------- 16 files changed, 281 insertions(+), 304 deletions(-) create mode 100644 test/registered/unit/spec/test_draft_per_runner_config.py delete mode 100644 test/registered/unit/spec/test_draft_server_args_copy.py delete mode 100644 test/registered/unit/spec/test_spec_worker_draft_isolation.py 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()