diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py index ed388ad77..a38af66e4 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py +++ b/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." ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index f48b60b0b..1cff5c983 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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", - "musa": "cuda graph", - "cpu": "cpu graph", - "npu": "npu graph", + "cuda": "CUDA graph", + "musa": "CUDA graph", + "cpu": "CPU 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"mem usage={self.graph_mem_usage:.2f} GB. avail mem={after_mem:.2f} GB." + f"Capture {capture_name} {graph_backend[self.device]} end. " + 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"mem usage={mem_usage:.2f} GB. avail mem={after_mem:.2f} GB." + f"Capture {capture_name} CUDA graph end. " + 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): 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 49760cfc6..d71c93b9a 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 @@ -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 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 a95da5347..19eded792 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 @@ -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) 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 7f5fb616f..0dd66d48e 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 @@ -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 diff --git a/python/sglang/srt/model_executor/runner_backend_utils/__init__.py b/python/sglang/srt/model_executor/runner_backend_utils/__init__.py index 7e75ba203..bfa6ac7b7 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/__init__.py +++ b/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 -# paths (Full backend used by decode + EAGLE draft runners). The -# piecewise-specific variant lives in -# piecewise_cuda_graph.context_manager and points users at -# --disable-piecewise-cuda-graph, which doesn't apply here. +# Generic failure-message hint for decode-style CUDA graph capture paths +# (Full backend used by decode + EAGLE draft runners). 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" - "4. disable CUDA graph by --cuda-graph-backend-decode=disabled. (Not recommended. Huge performance loss)\n" + "3. disable decode CUDA graph by --cuda-graph-backend-decode=disabled. " + "(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" ) diff --git a/python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py b/python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py index a646556de..70c687297 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py +++ b/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" - "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" + "Piecewise CUDA graph failed.\n" + PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG ) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index c31628d11..a6f46ea19 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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): diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 3a1e97d8c..07f873b00 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/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):