minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Clean up CUDA graph capture logs (#28937)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-06-22 06:15:26 -07:00
committed by GitHub
parent ead39d38fc
commit ad9723af03
9 changed files with 112 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py
+6 -2
View File
@@ -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
+43 -17
View File
@@ -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):
python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py
+1 -1
View File
@@ -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
-2
View File
@@ -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
+19 -7
View File
@@ -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
+13 -7
View File
@@ -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"
)
python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py
+4 -16
View File
@@ -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
)
python/sglang/srt/speculative/eagle_worker_v2.py
+20 -5
View File
@@ -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
+6 -2
View File
@@ -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):