[Bug][Spec] fix startup crash and reduce CUDA graph memory usage for speculative adaptive (#35275)

Co-authored-by: alphabetc1 <2508695655@qq.com>
Co-authored-by: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/srt/model_executor/model_runner.py

@@ -182,6 +182,7 @@ from sglang.srt.runtime_context import (
     get_spec,
     is_ep_joiner,
     is_ep_scale_joiner,
+    max_speculative_num_draft_tokens,
     remote_instance_transfer_engine_enabled,
     set_global_dwdp_manager,
 )
@@ -193,6 +194,9 @@ from sglang.srt.server_args import (  # noqa: F401  (re-export)
     add_chunked_prefix_cache_attention_backend,
     get_global_server_args,
 )
+from sglang.srt.speculative.adaptive_spec_params import (
+    resolve_candidate_steps_from_config,
+)
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.speculative.spec_utils import resolve_num_tokens_per_req
 from sglang.srt.state_capturer.base import TopkCaptureOutput
@@ -819,9 +823,26 @@ class ModelRunner:
 
     def max_decode_logits_rows(self) -> int:
         """Rows the shared logits buffer needs."""
-        num_tokens_per_req = self.decode_num_tokens_per_req()
-        capture_bs, _ = get_batch_sizes_to_capture(self, num_tokens_per_req)
-        return max(capture_bs) * num_tokens_per_req
+        # Resolution can turn speculative_adaptive off, so the effective value
+        # lives in the bags while the startup record keeps the CLI input.
+        spec = get_spec()
+        draft_token_counts = [max_speculative_num_draft_tokens()]
+        if spec.speculative_adaptive:
+            draft_token_counts.extend(
+                steps + 1
+                for steps in resolve_candidate_steps_from_config(
+                    spec.speculative_adaptive_config
+                )
+            )
+
+        max_rows = 0
+        for draft_tokens in draft_token_counts:
+            num_tokens_per_req = self.decode_num_tokens_per_req(
+                num_draft_tokens=draft_tokens
+            )
+            capture_bs, _ = get_batch_sizes_to_capture(self, num_tokens_per_req)
+            max_rows = max(max_rows, max(capture_bs) * num_tokens_per_req)
+        return max_rows
 
     def alloc_memory_pool(self, memory_pool_config: Optional[MemoryPoolConfig] = None):
         """Allocate KV cache memory pools only (no backends or cuda graphs)."""

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -98,6 +98,9 @@ from sglang.srt.model_executor.runner_utils.capture_mode import (
 from sglang.srt.model_executor.runner_utils.deepep_adapter import (
     DeepEPCudaGraphRunnerAdapter,
 )
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_capture_stream,
+)
 from sglang.srt.model_executor.runner_utils.shared_read_event import make_external_event
 from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups
 from sglang.srt.runtime_context import (
@@ -1061,7 +1064,12 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         # can reuse the memory pool allocated for the large shapes.
         with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc):
             if not self.enable_pdmux:
-                with graph_capture() as graph_capture_context, profile_context as prof:
+                with (
+                    graph_capture(
+                        stream=get_or_create_global_graph_capture_stream()
+                    ) as graph_capture_context,
+                    profile_context as prof,
+                ):
                     self.stream = graph_capture_context.stream
                     with self.backend.capture_session(self.stream):
                         self._capture_one_stream()
@@ -1239,8 +1247,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                     variant_label,
                     dsa_variant,
                 )
+                # Adaptive runners may own a different backend than model_runner.
                 post_warmup_hook = getattr(
-                    self.model_runner.attn_backend,
+                    attn_backend,
                     "on_after_cuda_graph_warmup",
                     None,
                 )

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -117,6 +117,9 @@ from sglang.srt.model_executor.runner_utils import (
 from sglang.srt.model_executor.runner_utils.buffers import (
     PrefillInputBuffers,
 )
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_capture_stream,
+)
 from sglang.srt.model_loader.utils import resolve_language_model
 from sglang.srt.runtime_context import (
     get_exec,
@@ -1341,7 +1344,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # decode + prefill runners; see BaseRunner.warmup).
         self.warmup()
         with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc):
-            with graph_capture() as graph_capture_context:
+            with graph_capture(
+                stream=get_or_create_global_graph_capture_stream()
+            ) as graph_capture_context:
                 self.stream = graph_capture_context.stream
                 with self.backend.capture_session(self.stream):
                     self._capture_one_stream()

python/sglang/srt/model_executor/runner_utils/__init__.py

@@ -24,6 +24,7 @@ from sglang.srt.model_executor.runner_utils.deepep_adapter import (  # noqa: F40
 )
 from sglang.srt.model_executor.runner_utils.pool import (  # noqa: F401
     get_global_graph_memory_pool,
+    get_or_create_global_graph_capture_stream,
     set_global_graph_memory_pool,
 )
 from sglang.srt.model_executor.runner_utils.shared_read_event import (  # noqa: F401

python/sglang/srt/model_executor/runner_utils/pool.py

@@ -14,6 +14,7 @@
 """Process-wide CUDA graph memory pool shared across the prefill and
 decode graph backends. The two phases never replay concurrently, so
 sharing one pool reserves only the larger phase's capture footprint.
+Serial capture passes also share one stream to reuse allocator scratch.
 """
 
 from __future__ import annotations
@@ -25,7 +26,7 @@ from typing import Any, Iterator, Optional
 import torch
 
 from sglang.srt.environ import envs
-from sglang.srt.runtime_context import get_resources
+from sglang.srt.runtime_context import get_resources, get_stream
 from sglang.srt.utils import is_cuda
 from sglang.srt.utils.cuda_vmm_utils import BumpArenaStub
 
@@ -38,6 +39,8 @@ _borrow_static_runs: Optional[list[tuple[int, int]]] = None
 _borrow_extents_total = 0
 _largest_logged_graph_pool_borrow = 0
 
+_CAPTURE_STREAM_NAME = "cuda_graph_capture"
+
 
 def disable_graph_pool_borrow(reason: str) -> None:
     """Disable borrowing when graph storage is managed outside the shared pool."""
@@ -81,6 +84,15 @@ def get_or_create_global_graph_memory_pool(device_module: Any) -> Any:
     return resources.graph_memory_pool
 
 
+def get_or_create_global_graph_capture_stream() -> Any:
+    """Return the shared graph capture stream, creating it on first use so every
+    capture pass reserves the pool's scratch once instead of per stream.
+
+    CUDA only — the NPU / XPU / CPU graph runners keep their own streams.
+    """
+    return get_stream(_CAPTURE_STREAM_NAME)
+
+
 def graph_pool_borrow_enabled() -> bool:
     if (
         _borrow_disabled_reason is not None

python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py

@@ -278,9 +278,10 @@ def _seeded_rng(seed: int, *, device: str | torch.device):
 
 
 @contextmanager
-def _single_rank_graph_capture():
-    stream = torch.cuda.Stream()
-    yield SimpleNamespace(stream=stream)
+def _single_rank_graph_capture(stream=None):
+    # Mirrors `graph_capture(stream=None)`: capture runs on the caller's stream
+    # when it leases one, so the shim stays valid for both call shapes.
+    yield SimpleNamespace(stream=stream if stream is not None else torch.cuda.Stream())
 
 
 def _reset_cuda_graph_test_buffers() -> None:

test/registered/unit/model_executor/test_model_runner_decode_rows.py

@@ -0,0 +1,51 @@
+import tempfile
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class _FakeModelRunner:
+    max_decode_logits_rows = ModelRunner.max_decode_logits_rows
+
+    def __init__(self, *, initial_width: int, cuda_graph_bs: list[int]):
+        self.initial_width = initial_width
+        self.cuda_graph_bs = cuda_graph_bs
+
+    def decode_num_tokens_per_req(self, *, num_draft_tokens=None):
+        return self.initial_width if num_draft_tokens is None else num_draft_tokens
+
+
+def _alignment_8_capture_bs(runner, width):
+    return ([bs for bs in runner.cuda_graph_bs if bs * width % 8 == 0], [])
+
+
+class TestModelRunnerDecodeRows(unittest.TestCase):
+    def test_adaptive_sizing_covers_a_wider_candidate_width(self):
+        """The shared logits buffer is sized for the widest adaptive candidate
+        width: bs 12 at width 6 needs 72 rows."""
+        runner = _FakeModelRunner(initial_width=4, cuda_graph_bs=[4, 8, 12])
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
+            f.write('{"1":{"candidate_steps":[3,5]}}')
+            f.flush()
+            spec = SimpleNamespace(
+                speculative_adaptive=True, speculative_adaptive_config=f.name
+            )
+            with patch(
+                "sglang.srt.model_executor.model_runner.get_spec", return_value=spec
+            ), patch(
+                "sglang.srt.model_executor.model_runner.max_speculative_num_draft_tokens",
+                return_value=6,
+            ), patch(
+                "sglang.srt.model_executor.model_runner.get_batch_sizes_to_capture",
+                side_effect=_alignment_8_capture_bs,
+            ):
+                self.assertEqual(runner.max_decode_logits_rows(), 72)
+
+
+if __name__ == "__main__":
+    unittest.main()