[cuda_graph] Size dp-local decode graphs by the runner's own gather requirement

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
 

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."
     )

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,15 +76,17 @@ 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
     if get_exec().overlap.enable_two_batch_overlap:
         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.

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

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

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:

test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py

@@ -67,46 +67,67 @@ 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
 
     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()
 

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], [])