minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

config: resolve the draft worker's config per runner, not on a copy

Browse Source 
The v2 spec workers got a published `ServerArgs` copy carrying two values: the
target's context length and `--speculative-draft-load-format`. Neither is a
process-wide config change — each is consumed by exactly one constructor — so
the copy, the publish switch around the draft build, and the replay of the
target's resolved overrides onto it all go away, and the values travel to the
runner that owns them:

- **Context length.** `TpModelWorker` already takes it (`context_length=None`
  keeps `server_args.context_length`); the four v2 draft workers and
  `build_draft_tp_worker` pass the target's, which every one of them has in
  scope as `target_worker` / `target_model_config`.
- **Load format.** `ModelRunner._draft_load_format()` resolves it for a draft
  runner and `build_load_config` takes it, so the `LoadConfig` is per-runner.
  Model code also reads it off the bag while it builds — Inkling replaces
  per-element noise in its shared-expert scales under dummy loading — so the
  load is wrapped in a scoped bag override that puts the target's value back.
- `skip_tokenizer_init` was on the copy for nobody: `TpModelWorker` already
  short-circuits the tokenizer for a draft worker (`or self.is_draft_worker`).

`PrefillCudaGraphRunner._max_addressable_prefix_len` capped the prefix by
`server_args.context_length`, which the copy used to carry for the draft; it now
reads the runner's own `model_config.context_len`. That is also more accurate for
the target, whose `--context-length` may be unset while the resolved context is
shorter than the token table.

What stays a variant is the dflash/dspark path's attention backend: backend
selection reads it off the config object the draft runner holds, and the
resolved gate has to survive the variant's publish. `draft_server_args_overrides`
now carries only those fields and says why.
This commit is contained in:
Cheng Wan
2026-08-05 19:31:23 -07:00
committed by GitHub
parent 99cfc90658
commit 64eeb153df
16 changed files with 281 additions and 304 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+7 -14
View File
@@ -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 (
python/sglang/srt/model_executor/model_runner.py
+45 -10
View File
@@ -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 = (
python/sglang/srt/model_executor/model_runner_components/load_model_utils.py
+7 -3
View File
@@ -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,
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+3 -1
View File
@@ -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
python/sglang/srt/runtime_context.py
+7 -2
View File
@@ -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:
python/sglang/srt/server_args.py
+4 -2
View File
@@ -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"
python/sglang/srt/speculative/draft_worker_common.py
+13 -49
View File
@@ -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
python/sglang/srt/speculative/eagle_worker_v2.py
+2
View File
@@ -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
python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py
+2
View File
@@ -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()
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+2
View File
@@ -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
python/sglang/srt/speculative/standalone_worker_v2.py
+2
View File
@@ -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
test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py
+4 -8
View File
@@ -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"])
test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py
+13 -1
View File
@@ -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
)
test/registered/unit/spec/test_draft_per_runner_config.py
+170
View File
@@ -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()
test/registered/unit/spec/test_draft_server_args_copy.py
-84
View File
@@ -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()
test/registered/unit/spec/test_spec_worker_draft_isolation.py
-130
View File
@@ -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()