From 5a8e360e705fc7b8046f6b060ba4fc557ff606c7 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Sun, 9 Aug 2026 20:56:00 -0700 Subject: [PATCH] [PD] Skip speculative verify scratch on prefill servers (saves num_draft_tokens x mamba pool per rank) (#34191) --- python/sglang/srt/mem_cache/kv_cache_configurator.py | 9 ++++++++- .../model_runner_components/cuda_graph_setup.py | 8 ++++++++ python/sglang/srt/model_executor/runner/base_runner.py | 8 +++++++- 3 files changed, 23 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 5697f26ec..c2642914a 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -805,7 +805,14 @@ class KVCacheConfigurator: ), enable_mamba_extra_buffer=mamba_extra_buffer_enabled(), enable_mamba_extra_buffer_lazy=mamba_extra_buffer_lazy_enabled(), - speculative_num_draft_tokens=max_speculative_num_draft_tokens(), + # A PD prefill server never runs TARGET_VERIFY, so skip the + # verify-only per-draft-token state snapshots (see the draft-head + # case above: None => the pool skips SpeculativeState). + speculative_num_draft_tokens=( + None + if get_disagg().disaggregation_mode == "prefill" + else max_speculative_num_draft_tokens() + ), speculative_eagle_topk=get_spec().speculative_eagle_topk, enable_overlap_schedule=not get_schedule().disable_overlap_schedule, start_layer=self.layer_info.start_layer, diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index 475038f54..a41a320f9 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -414,6 +414,14 @@ def capture_decode_graph(*, model_runner: ModelRunner) -> GraphCapture: capture_time=0, ) + # A PD prefill server never replays the target-verify graph, and its pool + # is built without the spec-verify scratch the capture would need. + if ( + model_runner.spec_algorithm.is_speculative() + and not model_runner.is_draft_worker + and model_runner.server_args.disaggregation_mode == "prefill" + ): + return no_capture if not model_runner.is_generation: # TODO: Currently, cuda graph only captures decode steps, which only exists for generation models return no_capture diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index a7bc4d6f8..1597344ca 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -404,7 +404,13 @@ class BaseRunner(ABC): else get_server_return_hidden_states_mode(mr.server_args) ) num_tokens_per_req = 1 - if mr.spec_algorithm.is_speculative(): + # A PD prefill target worker's pool has no SpeculativeState, so a + # TARGET_VERIFY dummy forward would trip the linear-attn backend's + # pool-type assert. Warm up in plain DECODE instead. + _is_pd_prefill_target = ( + mr.server_args.disaggregation_mode == "prefill" and not mr.is_draft_worker + ) + if mr.spec_algorithm.is_speculative() and not _is_pd_prefill_target: if mr.is_draft_worker: assert ( mr.spec_algorithm.supports_target_verify_for_draft()