From 95d8bd63fdd189a19f5ec2f88eb327a55d6449ce Mon Sep 17 00:00:00 2001 From: Hanming Lu <69857889+hanming-lu@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:11:19 -0700 Subject: [PATCH] [cuda_graph] Size dp-local decode graphs by the runner's own gather requirement get_batch_sizes_to_capture filtered the decode buckets by get_cuda_graph_batch_size_alignment(), which multiplies the alignment by attn_tp_size whenever the global require_gathered_buffer() is true. DecodeCudaGraphRunner already exempts a dp-local forward from the MLP gather through _forward_is_dp_local, but the bucket filter never saw that exemption: under DP attention a dp-local draft kept only buckets whose rows were a multiple of attn_tp_size, padded a one- or two-request local batch up to that bucket, and a small user-supplied bucket list failed startup with `capture_bs=[]`. Thread the requirement through instead of re-deriving it globally: - get_cuda_graph_batch_size_alignment / get_batch_sizes_to_capture take an optional keyword gathered_buffer_required; omitted, both behave as before. - DecodeCudaGraphRunner.decode_graph_gather_requirements(model_runner) owns the (mlp_tp_gather, attn_tp_gather) pair, including the dp-local exemption. __init__ derives its instance predicates from it and passes require_gathered_buffer to both bucket derivations. - ModelRunner.decode_graph_gather_requirements() delegates to the runner class so max_decode_logits_rows and BaseRunner._widest_decode_rows size the logits buffer and the reduction workspace with the alignment the runner captures with. - get_batch_sizes_to_capture pads num_max_requests with the alignment it computed rather than the global get_cuda_graph_max_batch_size. - capture_decode_graph reports bs= from the constructed runner's capture_bs in the end-of-capture log instead of deriving the list with global alignment before the runner exists. A dp-local runner now captures the plain decode bucket list and replays the graph sized to its local batch; every other runner is unchanged. Co-authored-by: charlotte12l --- .../sglang/srt/model_executor/model_runner.py | 12 ++++++- .../cuda_graph_setup.py | 6 ++-- .../runner/base_cuda_graph_runner.py | 13 +++++--- .../srt/model_executor/runner/base_runner.py | 6 +++- .../runner/decode_cuda_graph_runner.py | 24 +++++++++++--- python/sglang/srt/utils/common.py | 8 +++-- .../test_cuda_graph_setup.py | 33 +++++++++++++++---- .../test_model_runner_decode_rows.py | 6 +++- 8 files changed, 86 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 04f8524d0093..f892c3899860 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -847,6 +847,12 @@ def decode_num_tokens_per_req( dllm_config = DllmConfig.from_server_args(self.server_args) return dllm_config.block_size if dllm_config is not None else 1 + def decode_graph_gather_requirements(self) -> tuple[bool, bool]: + """MLP and attention gathers required by this runner's decode graphs.""" + return self._decode_cuda_graph_runner_cls().decode_graph_gather_requirements( + self + ) + def max_decode_logits_rows(self) -> int: """Rows the shared logits buffer needs.""" # Resolution can turn speculative_adaptive off, so the effective value @@ -866,7 +872,11 @@ def max_decode_logits_rows(self) -> int: 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) + capture_bs, _ = get_batch_sizes_to_capture( + self, + num_tokens_per_req, + gathered_buffer_required=any(self.decode_graph_gather_requirements()), + ) max_rows = max(max_rows, max(capture_bs) * num_tokens_per_req) return max_rows 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 ec77b9baa748..4bfd28f59ed7 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 @@ -711,12 +711,11 @@ def capture_decode_graph(*, model_runner: ModelRunner) -> GraphCapture: else: capture_name = f"{role} decode" num_tokens_per_req = 1 - capture_bs, _ = get_batch_sizes_to_capture(model_runner, num_tokens_per_req) decode_backend = get_exec().graph.cuda_graph_config.decode.backend logger.info( f"Capture {capture_name} {graph_backend[model_runner.device]} begin. " f"backend={decode_backend}, num_tokens_per_req={num_tokens_per_req}, " - f"bs={capture_bs}, avail mem={before_mem:.2f} GB" + f"avail mem={before_mem:.2f} GB" ) GraphRunnerCls = current_platform.get_graph_runner_cls() @@ -740,8 +739,11 @@ def capture_decode_graph(*, model_runner: ModelRunner) -> GraphCapture: after_mem = get_available_gpu_memory(model_runner.device, model_runner.gpu_id) memory_usage_gb = before_mem - after_mem capture_time = time.perf_counter() - tic + # The runner owns the bucket list: a dp-local draft aligns by its local + # batch, which a global re-derivation here would reject as empty. logger.info( f"Capture {capture_name} {graph_backend[model_runner.device]} end. " + f"bs={runner.capture_bs}, " f"elapsed={capture_time:.2f} s, " f"mem usage={memory_usage_gb:.2f} GB, avail mem={after_mem:.2f} GB." ) diff --git a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py index 8cbd23148469..ac5143ce88d8 100644 --- a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py @@ -28,8 +28,8 @@ get_flags, ) from sglang.srt.utils import ( + ceil_align, get_cuda_graph_batch_size_alignment, - get_cuda_graph_max_batch_size, ) if TYPE_CHECKING: @@ -62,7 +62,10 @@ def freeze_gc(enable_cudagraph_gc: bool): def get_batch_sizes_to_capture( - model_runner: ModelRunner, captured_req_width: int = 1 + model_runner: ModelRunner, + captured_req_width: int = 1, + *, + gathered_buffer_required: Optional[bool] = None, ) -> Tuple[List[int], List[int]]: """Build the (capture_bs, compile_bs) lists for the decode runner. @@ -73,7 +76,9 @@ def get_batch_sizes_to_capture( capture_bs = list(get_exec().graph.cuda_graph_config.decode.bs) num_max_requests = model_runner.req_to_token_pool.size - mul_base = get_cuda_graph_batch_size_alignment() + mul_base = get_cuda_graph_batch_size_alignment( + gathered_buffer_required=gathered_buffer_required + ) # TBO splits each request's rows across two micro-batches, so the # alignment constraint applies per request rather than per token row. alignment_width = captured_req_width @@ -81,7 +86,7 @@ def get_batch_sizes_to_capture( alignment_width = 1 # pad `num_max_requests` to avoid being filtered out - num_max_requests = get_cuda_graph_max_batch_size(num_max_requests) + num_max_requests = ceil_align(num_max_requests, mul_base) if max(capture_bs) > num_max_requests: # In some cases (e.g., with a small GPU or --max-running-requests), the #max-running-requests # is very small. We add more values here to make sure we capture the maximum bs. diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index e0fcc11148b5..f27ad22d9159 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -367,7 +367,11 @@ def _widest_decode_rows(self) -> Optional[int]: mr = self.model_runner tokens_per_req = mr.decode_num_tokens_per_req() - capture_bs, _ = get_batch_sizes_to_capture(mr, tokens_per_req) + capture_bs, _ = get_batch_sizes_to_capture( + mr, + tokens_per_req, + gathered_buffer_required=any(mr.decode_graph_gather_requirements()), + ) if not capture_bs: return None return max(capture_bs) * tokens_per_req 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 13e6cf1d958c..3413fc7114ce 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 @@ -245,10 +245,9 @@ def __init__( self.enable_torch_compile = get_flags().capture.enable_torch_compile self.disable_padding = get_exec().graph.disable_cuda_graph_padding self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder - self.require_mlp_tp_gather = ( - require_mlp_tp_gather() and not self._forward_is_dp_local(model_runner) + self.require_mlp_tp_gather, self.require_attn_tp_gather = ( + self.decode_graph_gather_requirements(model_runner) ) - self.require_attn_tp_gather = require_attn_tp_gather() # Composite predicates derive from the instance values so the dp-local # draft exemption above stays consistent (require_gathered_buffer == # mlp_tp_gather or attn_tp_gather; require_mlp_sync adds dp attention). @@ -330,7 +329,9 @@ def __init__( # --- bucket sizes --------------------------------------------- self.capture_bs, self.compile_bs = get_batch_sizes_to_capture( - model_runner, self.captured_req_width + model_runner, + self.captured_req_width, + gathered_buffer_required=self.require_gathered_buffer, ) if self.dllm_uses_input_embeds: max_requests = min( @@ -629,6 +630,17 @@ def _resolve_lora_variant(self, forward_batch: ForwardBatch): return "lora" return "nolora" + @classmethod + def decode_graph_gather_requirements(cls, model_runner) -> tuple[bool, bool]: + """(mlp_tp_gather, attn_tp_gather) needed by this runner's decode graphs. + + Startup sizing (logits rows, reduction workspaces) asks the runner class + so it aligns capture buckets exactly as the runner itself will.""" + return ( + require_mlp_tp_gather() and not cls._forward_is_dp_local(model_runner), + require_attn_tp_gather(), + ) + @staticmethod def _forward_is_dp_local(model_runner) -> bool: """The DSpark dense draft runs attn-TP-local (draft_tp_context): each @@ -1069,7 +1081,9 @@ def capture(self) -> None: if self.enable_torch_compile and not (get_flags().capture.enable_torch_compile): self.enable_torch_compile = False _, self.compile_bs = get_batch_sizes_to_capture( - self.model_runner, self.captured_req_width + self.model_runner, + self.captured_req_width, + gathered_buffer_required=self.require_gathered_buffer, ) profile_context = empty_context() # Holds the active torch profiler during capture so the backend can diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 0d22705f49ef..5d2d07456341 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -4221,11 +4221,15 @@ def require_mlp_sync(): return get_parallel().attn_dp_enabled or require_gathered_buffer() -def get_cuda_graph_batch_size_alignment() -> int: +def get_cuda_graph_batch_size_alignment( + *, gathered_buffer_required: Optional[bool] = None +) -> int: + if gathered_buffer_required is None: + gathered_buffer_required = require_gathered_buffer() alignment = 1 if get_exec().overlap.enable_two_batch_overlap: alignment *= 2 - if require_gathered_buffer(): + if gathered_buffer_required: alignment *= get_parallel().attn_tp_size # TODO: unverified on NVIDIA; drop the gate once validated on CUDA. if not is_hip() and alignment % get_parallel().attn_cp_size != 0: diff --git a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py b/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py index ac9aad7da824..22acfb824572 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py +++ b/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py @@ -67,27 +67,44 @@ def test_reuse_tables_pass_through_but_distinct_duplicates_raise(): ) -def test_model_runner_can_override_decode_graph_runner(monkeypatch): +def test_draft_startup_preserves_runner_local_capture_sizes(monkeypatch, caplog): + """Startup must not reject local draft buckets using global TP alignment.""" + from sglang.srt.model_executor.runner.base_cuda_graph_runner import ( + get_batch_sizes_to_capture, + ) from sglang.srt.runtime_context import get_context + from sglang.srt.utils import common + capture_bs = [1, 2, 3, 4, 8] # The capture decision reads the graph configuration and the MoE backends # out of the bags. override = get_context().override_server_args( - cuda_graph_config=SimpleNamespace(decode=SimpleNamespace(backend="default")), + cuda_graph_config=SimpleNamespace( + decode=SimpleNamespace(backend="default", bs=capture_bs) + ), + enable_torch_compile=False, + enable_two_batch_overlap=False, ) override.install() class CustomGraphRunner: def __init__(self, model_runner): self.model_runner = model_runner + self.capture_bs, _ = get_batch_sizes_to_capture( + model_runner, 7, gathered_buffer_required=False + ) class TestModelRunner: is_generation = True device = "cuda" gpu_id = 0 - is_draft_worker = False - spec_algorithm = SimpleNamespace(is_speculative=lambda: False) + is_draft_worker = True + spec_algorithm = SimpleNamespace(is_speculative=lambda: True) server_args = SimpleNamespace(model_impl="auto") + req_to_token_pool = SimpleNamespace(size=32) + + def decode_num_tokens_per_req(self): + return 7 def _decode_cuda_graph_runner_cls(self): return CustomGraphRunner @@ -95,18 +112,22 @@ def _decode_cuda_graph_runner_cls(self): model_runner = TestModelRunner() monkeypatch.setattr(cuda_graph_setup, "check_cuda_graph_backend", lambda *_: False) monkeypatch.setattr(cuda_graph_setup, "get_available_gpu_memory", lambda *_: 10.0) + monkeypatch.setattr(common, "require_gathered_buffer", lambda: True) monkeypatch.setattr( - cuda_graph_setup, "get_batch_sizes_to_capture", lambda *_: ([1], None) + common, "get_parallel", lambda: SimpleNamespace(attn_tp_size=16, attn_cp_size=1) ) monkeypatch.setattr( cuda_graph_setup.current_platform, "is_out_of_tree", lambda: False ) try: - capture = capture_decode_graph(model_runner=model_runner) + with caplog.at_level("INFO", logger=cuda_graph_setup.__name__): + capture = capture_decode_graph(model_runner=model_runner) assert isinstance(capture.runner, CustomGraphRunner) assert capture.runner.model_runner is model_runner + assert capture.runner.capture_bs == capture_bs + assert f"bs={capture_bs}" in caplog.text finally: override.restore() 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 index fe0668ba81b0..0f61ad18fa29 100644 --- a/test/registered/unit/model_executor/test_model_runner_decode_rows.py +++ b/test/registered/unit/model_executor/test_model_runner_decode_rows.py @@ -12,6 +12,9 @@ class _FakeModelRunner: max_decode_logits_rows = ModelRunner.max_decode_logits_rows + def decode_graph_gather_requirements(self): + return True, False + def __init__(self, *, initial_width: int, cuda_graph_bs: list[int]): self.initial_width = initial_width self.cuda_graph_bs = cuda_graph_bs @@ -20,7 +23,8 @@ 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): +def _alignment_8_capture_bs(runner, width, *, gathered_buffer_required): + assert gathered_buffer_required return ([bs for bs in runner.cuda_graph_bs if bs * width % 8 == 0], [])