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
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()