From b8a6adadfe8c292fbb673291dbd2a5d23232f500 Mon Sep 17 00:00:00 2001 From: GoldPancake <56388518+Deleter-D@users.noreply.github.com> Date: Thu, 27 Aug 2026 15:46:13 +0800 Subject: [PATCH] [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) --- .../sglang/srt/model_executor/model_runner.py | 27 ++++++++-- .../runner/decode_cuda_graph_runner.py | 13 ++++- .../runner/prefill_cuda_graph_runner.py | 7 ++- .../model_executor/runner_utils/__init__.py | 1 + .../srt/model_executor/runner_utils/pool.py | 14 ++++- .../runner_modes/speculative_draft_runner.py | 7 +-- .../test_model_runner_decode_rows.py | 51 +++++++++++++++++++ 7 files changed, 110 insertions(+), 10 deletions(-) create mode 100644 test/registered/unit/model_executor/test_model_runner_decode_rows.py diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d4d6e3e54..6738521f4 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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).""" diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index c3d311e4f..f3c009117 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/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, ) diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 922b2d21a..84085a287 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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() diff --git a/python/sglang/srt/model_executor/runner_utils/__init__.py b/python/sglang/srt/model_executor/runner_utils/__init__.py index a8541e9ea..c2529cb20 100644 --- a/python/sglang/srt/model_executor/runner_utils/__init__.py +++ b/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 diff --git a/python/sglang/srt/model_executor/runner_utils/pool.py b/python/sglang/srt/model_executor/runner_utils/pool.py index 2ff115188..1e86cd3e9 100644 --- a/python/sglang/srt/model_executor/runner_utils/pool.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index dd76e9187..cbebd85c0 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/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: diff --git a/test/registered/unit/model_executor/test_model_runner_decode_rows.py b/test/registered/unit/model_executor/test_model_runner_decode_rows.py new file mode 100644 index 000000000..b96cf1e47 --- /dev/null +++ b/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()