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

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
@@ -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"])
@@ -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
)