Clean up CUDA graph capture logs (#28937)

python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py

@@ -118,7 +118,8 @@ class MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(
                 tic = time.perf_counter()
                 before_mem = get_available_gpu_memory(self.device, self.gpu_id)
                 logger.info(
-                    f"Capture draft extend cuda graph begin (step {step}). This can take up to several minutes. avail mem={before_mem:.2f} GB"
+                    f"Capture draft extend CUDA graph begin. step={step}, "
+                    f"avail mem={before_mem:.2f} GB"
                 )
 
                 self.runners[step].init_buffers_and_capture(
@@ -133,5 +134,8 @@ class MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(
 
                 after_mem = get_available_gpu_memory(self.device, self.gpu_id)
                 logger.info(
-                    f"Capture draft extend cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB."
+                    "Capture draft extend CUDA graph end. "
+                    f"step={step}, elapsed={time.perf_counter() - tic:.2f} s, "
+                    f"mem usage={(before_mem - after_mem):.2f} GB, "
+                    f"avail mem={after_mem:.2f} GB."
                 )

python/sglang/srt/model_executor/model_runner.py

@@ -166,6 +166,7 @@ from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
 from sglang.srt.model_executor.runner import (
     EagerRunner,
     PrefillCudaGraphRunner,
+    get_batch_sizes_to_capture,
 )
 from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
 from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
@@ -2538,15 +2539,32 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         graph_backend = defaultdict(
             lambda: f"{current_platform.device_name} graph",
             {
-                "cuda": "cuda graph",
+                "cuda": "CUDA graph",
-                "musa": "cuda graph",
+                "musa": "CUDA graph",
-                "cpu": "cpu graph",
+                "cpu": "CPU graph",
-                "npu": "npu graph",
+                "npu": "NPU graph",
             },
         )
+        role = "draft" if self.is_draft_worker else "target"
+        if self.spec_algorithm.is_speculative():
+            capture_name = f"{role} verify"
+            num_tokens_per_bs = (
+                self.spec_algorithm.get_num_tokens_per_bs_for_target_verify(
+                    self.server_args.speculative_num_draft_tokens,
+                    self.is_draft_worker,
+                )
+            )
+        else:
+            capture_name = f"{role} decode"
+            num_tokens_per_bs = 1
+        capture_bs, _ = get_batch_sizes_to_capture(self, num_tokens_per_bs)
+        decode_backend = self.server_args.cuda_graph_config.decode.backend
         logger.info(
-            f"Capture {graph_backend[self.device]} begin. This can take up to several minutes. avail mem={before_mem:.2f} GB"
+            f"Capture {capture_name} {graph_backend[self.device]} begin. "
+            f"backend={decode_backend}, num_tokens_per_bs={num_tokens_per_bs}, "
+            f"bs={capture_bs}, avail mem={before_mem:.2f} GB"
         )
+
         if current_platform.is_out_of_tree():
             GraphRunnerCls = current_platform.get_graph_runner_cls()
             self.decode_cuda_graph_runner = GraphRunnerCls(self)
@@ -2567,12 +2585,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         after_mem = get_available_gpu_memory(self.device, self.gpu_id)
         self.graph_mem_usage = before_mem - after_mem
         logger.info(
-            f"Capture {graph_backend[self.device]} end. Time elapsed: {time.perf_counter() - tic:.2f} s. "
+            f"Capture {capture_name} {graph_backend[self.device]} end. "
-            f"mem usage={self.graph_mem_usage:.2f} GB. avail mem={after_mem:.2f} GB."
+            f"elapsed={time.perf_counter() - tic:.2f} s, "
+            f"mem usage={self.graph_mem_usage:.2f} GB, avail mem={after_mem:.2f} GB."
         )
 
     def init_prefill_cuda_graph(self, force_for_draft_worker: bool = False):
-        """Initialize piecewise CUDA graph runner."""
+        """Initialize prefill CUDA graph runner."""
         self.prefill_cuda_graph_runner = None
 
         if check_cuda_graph_backend(Phase.PREFILL, Backend.DISABLED):
@@ -2613,17 +2632,17 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             self.prefill_cuda_graph_runner = self.eager_runner
             return
 
-        # Disable piecewise CUDA graph for non-language models
+        # Disable prefill CUDA graph for non-language models
         if not hasattr(self.model, "model"):
             logger.warning(
-                "Disable piecewise CUDA graph because the model is not a language model"
+                "Disable prefill CUDA graph because the model is not a language model"
             )
             return
 
-        # Disable piecewise CUDA graph for non capture size
+        # Disable prefill CUDA graph for non capture size
         if not self.server_args.cuda_graph_config.prefill.bs:
             logger.warning(
-                "Disable piecewise CUDA graph because the capture size is not set"
+                "Disable prefill CUDA graph because the capture size is not set"
             )
             return
 
@@ -2638,7 +2657,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             layer_model = language_model
         else:
             logger.warning(
-                "Disable piecewise CUDA graph because the model does not have a 'layers' attribute"
+                "Disable prefill CUDA graph because the model does not have a 'layers' attribute"
             )
             return
 
@@ -2710,14 +2729,20 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             # TODO(yuwei): support Non-Standard GQA
             log_info_on_rank0(
                 logger,
-                "Disable piecewise CUDA graph because some layers do not apply Standard GQA",
+                "Disable prefill CUDA graph because some layers do not apply Standard GQA",
             )
             return
 
         tic = time.perf_counter()
         before_mem = get_available_gpu_memory(self.device, self.gpu_id)
+        prefill_backend = self.server_args.cuda_graph_config.prefill.backend
+        role = "draft" if self.is_draft_worker else "target"
+        capture_name = f"{role} prefill"
+        capture_num_tokens = sorted(self.server_args.cuda_graph_config.prefill.bs)
         logger.info(
-            f"Capture piecewise CUDA graph begin. avail mem={before_mem:.2f} GB"
+            f"Capture {capture_name} CUDA graph begin. "
+            f"backend={prefill_backend}, num_tokens={capture_num_tokens}, "
+            f"avail mem={before_mem:.2f} GB"
         )
 
         self.prefill_cuda_graph_runner = PrefillCudaGraphRunner(self)
@@ -2725,8 +2750,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         after_mem = get_available_gpu_memory(self.device, self.gpu_id)
         mem_usage = before_mem - after_mem
         logger.info(
-            f"Capture piecewise CUDA graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. "
+            f"Capture {capture_name} CUDA graph end. "
-            f"mem usage={mem_usage:.2f} GB. avail mem={after_mem:.2f} GB."
+            f"elapsed={time.perf_counter() - tic:.2f} s, "
+            f"mem usage={mem_usage:.2f} GB, avail mem={after_mem:.2f} GB."
         )
 
     def init_threads_binding(self):

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

@@ -65,7 +65,7 @@ def get_batch_sizes_to_capture(
     """
 
     server_args = model_runner.server_args
-    capture_bs = server_args.cuda_graph_config.decode.bs
+    capture_bs = list(server_args.cuda_graph_config.decode.bs)
     num_max_requests = model_runner.req_to_token_pool.size
 
     mul_base = 1

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

@@ -93,7 +93,6 @@ from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_strea
 from sglang.srt.utils import (
     empty_context,
     get_available_gpu_memory,
-    log_info_on_rank0,
     require_attn_tp_gather,
     require_gathered_buffer,
     require_mlp_sync,
@@ -251,7 +250,6 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(
             model_runner, self.num_tokens_per_bs
         )
-        log_info_on_rank0(logger, f"Capture cuda graph bs {self.capture_bs}")
         if KTRANSFORMERS_AVAILABLE:
             KTMoEWrapper.set_capture_batch_sizes(self.capture_bs)
 

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

@@ -66,6 +66,9 @@ from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend impor
 from sglang.srt.model_executor.runner_backend.utils import (
     resolve_prefill_backend,
 )
+from sglang.srt.model_executor.runner_backend_utils import (
+    PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     set_tc_piecewise_forward_context,
 )
@@ -77,7 +80,6 @@ from sglang.srt.utils import (
     get_bool_env_var,
     is_hip,
     is_npu,
-    log_info_on_rank0,
     require_attn_tp_gather,
     require_mlp_tp_gather,
 )
@@ -135,10 +137,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         )
         self.max_bs = model_runner.req_to_token_pool.size
 
-        log_info_on_rank0(
-            logger, f"Capture cuda graph num tokens {self.capture_num_tokens}"
-        )
-
         self.capture_forward_mode = ForwardMode.EXTEND
         self.capture_hidden_mode = CaptureHiddenMode.NULL
         # If returning hidden states is enabled, or if speculative prefill
@@ -222,7 +220,15 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         self._prefill_static_buffers: Optional[Dict[str, torch.Tensor]] = None
         self.static_draft_hidden_states: Optional[torch.Tensor] = None
         self.layer_model = None
-        self.backend = resolve_prefill_backend(self)
+        try:
+            self.backend = resolve_prefill_backend(self)
+        except RuntimeError as e:
+            if _prefill_backend_name == Backend.TC_PIECEWISE:
+                raise Exception(
+                    f"Capture prefill CUDA graph failed: {e}\n"
+                    f"{PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG}"
+                )
+            raise
         if isinstance(self.backend, BreakableCudaGraphBackend):
             with torch.device(self.device):
                 self._prefill_static_buffers = {
@@ -304,7 +310,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # --- capture --------------------------------------------------
         self.device_module.synchronize()
         self.model_runner.tp_group.barrier()
-        self.capture()
+        try:
+            self.capture()
+        except RuntimeError as e:
+            raise Exception(
+                f"Capture prefill CUDA graph failed: {e}\n"
+                f"{PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG}"
+            )
 
         self.raw_num_tokens = 0
 

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

@@ -9,16 +9,22 @@ Subpackages:
 Backends in cuda_graph_backend/ import from here. Runners do not.
 """
 
-# Generic failure-message hint for non-piecewise CUDA graph capture
+# Generic failure-message hint for decode-style CUDA graph capture paths
-# paths (Full backend used by decode + EAGLE draft runners). The
+# (Full backend used by decode + EAGLE draft runners).
-# piecewise-specific variant lives in
-# piecewise_cuda_graph.context_manager and points users at
-# --disable-piecewise-cuda-graph, which doesn't apply here.
 CUDA_GRAPH_CAPTURE_FAILED_MSG = (
     "Possible solutions:\n"
     "1. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n"
     "2. set --cuda-graph-max-bs-decode to a smaller value (e.g., 16)\n"
-    "3. disable torch compile by not using --enable-torch-compile\n"
+    "3. disable decode CUDA graph by --cuda-graph-backend-decode=disabled. "
-    "4. disable CUDA graph by --cuda-graph-backend-decode=disabled. (Not recommended. Huge performance loss)\n"
+    "(Not recommended. Huge performance loss)\n"
+    "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
+)
+
+PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
+    "Possible solutions:\n"
+    "1. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n"
+    "2. set --cuda-graph-max-bs-prefill to a smaller value (e.g., 2048)\n"
+    "3. disable prefill CUDA graph by --cuda-graph-backend-prefill=disabled. "
+    "(Not recommended. Performance loss)\n"
     "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
 )

python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py

@@ -21,13 +21,13 @@ This module deliberately does **not** own torch.compile-specific state
 
 from __future__ import annotations
 
-import logging
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, List, Optional
 
-logger = logging.getLogger(__name__)
+from sglang.srt.model_executor.runner_backend_utils import (
-
+    PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -45,21 +45,11 @@ def is_in_tc_piecewise_cuda_graph() -> bool:
 def enable_tc_piecewise_cuda_graph():
     """Mark the enclosed scope as "we are inside a piecewise CUDA graph
     capture/replay". Sets _in_tc_piecewise_cuda_graph true for the duration.
-
-    Errors during capture surface a hint that lets users disable the
-    feature while filing a bug.
     """
     global _in_tc_piecewise_cuda_graph
     _in_tc_piecewise_cuda_graph = True
     try:
         yield
-    except Exception as e:
-        logger.error(
-            "Piecewise CUDA Graph failed with error: %s\n%s",
-            e,
-            TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
-        )
-        raise
     finally:
         _in_tc_piecewise_cuda_graph = False
 
@@ -112,7 +102,5 @@ def set_tc_piecewise_forward_context(
 
 
 TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
-    "Piecewise CUDA Graph capture failed.\n"
+    "Piecewise CUDA graph failed.\n" + PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG
-    "To work around this error, add --cuda-graph-backend-prefill=disabled to your launch command.\n"
-    "Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose"
 )

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -42,7 +42,10 @@ from sglang.srt.model_executor.cuda_graph_config import (
 )
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardBatch
 from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
-from sglang.srt.model_executor.runner import DecodeCudaGraphRunner
+from sglang.srt.model_executor.runner import (
+    DecodeCudaGraphRunner,
+    get_batch_sizes_to_capture,
+)
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.adaptive_runtime_state import (
     AdaptiveController,
@@ -382,12 +385,16 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             "musa": EAGLEDraftCudaGraphRunner,
         }
         # Capture draft
+        decode_backend = self.server_args.cuda_graph_config.decode.backend
+        capture_bs, _ = get_batch_sizes_to_capture(self.draft_runner)
         if self.speculative_num_steps > 1:
             tic = time.perf_counter()
             before_mem = get_available_gpu_memory(self.device, self.gpu_id)
             log_info_on_rank0(
                 logger,
-                f"Capture draft cuda graph begin. This can take up to several minutes. avail mem={before_mem:.2f} GB",
+                f"Capture draft decode CUDA graph begin. backend={decode_backend}, "
+                f"num_tokens_per_bs={self.topk}, bs={capture_bs}, "
+                f"avail mem={before_mem:.2f} GB",
             )
             self.cuda_graph_runner = Device2DraftCudaGraphRunner[
                 self.target_worker.device
@@ -395,7 +402,10 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             after_mem = get_available_gpu_memory(self.device, self.gpu_id)
             log_info_on_rank0(
                 logger,
-                f"Capture draft cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB.",
+                "Capture draft decode CUDA graph end. "
+                f"elapsed={time.perf_counter() - tic:.2f} s, "
+                f"mem usage={(before_mem - after_mem):.2f} GB, "
+                f"avail mem={after_mem:.2f} GB.",
             )
 
         Device2ExtendCudaGraphRunner = {
@@ -438,7 +448,9 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             before_mem = get_available_gpu_memory(self.device, self.gpu_id)
             log_info_on_rank0(
                 logger,
-                f"Capture draft extend cuda graph begin. This can take up to several minutes. avail mem={before_mem:.2f} GB",
+                f"Capture draft extend CUDA graph begin. backend={decode_backend}, "
+                f"num_tokens_per_bs={self.speculative_num_draft_tokens}, "
+                f"bs={capture_bs}, avail mem={before_mem:.2f} GB",
             )
             self.cuda_graph_runner_for_draft_extend = Device2ExtendCudaGraphRunner[
                 self.target_worker.device
@@ -448,7 +460,10 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             after_mem = get_available_gpu_memory(self.device, self.gpu_id)
             log_info_on_rank0(
                 logger,
-                f"Capture draft extend cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB.",
+                "Capture draft extend CUDA graph end. "
+                f"elapsed={time.perf_counter() - tic:.2f} s, "
+                f"mem usage={(before_mem - after_mem):.2f} GB, "
+                f"avail mem={after_mem:.2f} GB.",
             )
 
     def draft(self, batch: ScheduleBatch):

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -701,7 +701,8 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
                 tic = time.perf_counter()
                 before_mem = get_available_gpu_memory(self.device, self.gpu_id)
                 logger.info(
-                    f"Capture draft extend cuda graph begin (step {step}). This can take up to several minutes. avail mem={before_mem:.2f} GB"
+                    f"Capture draft extend CUDA graph begin. step={step}, "
+                    f"avail mem={before_mem:.2f} GB"
                 )
 
                 self.runners[step].init_buffers_and_capture(
@@ -716,7 +717,10 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
 
                 after_mem = get_available_gpu_memory(self.device, self.gpu_id)
                 logger.info(
-                    f"Capture draft extend cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB."
+                    "Capture draft extend CUDA graph end. "
+                    f"step={step}, elapsed={time.perf_counter() - tic:.2f} s, "
+                    f"mem usage={(before_mem - after_mem):.2f} GB, "
+                    f"avail mem={after_mem:.2f} GB."
                 )
 
     def reset_buffers(self, forward_batch, batch_result):