[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

@@ -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,

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

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