From 145bd54f1b941aef4d383faf225a17ee4766c6fd Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Sat, 10 Jan 2026 03:29:13 -0800 Subject: [PATCH] Piecewise Cuda Graph Memory Usage (#15927) Signed-off-by: Oasis-Git --- .../srt/compilation/cuda_piecewise_backend.py | 9 ++---- python/sglang/srt/server_args.py | 31 +++++++++++-------- 2 files changed, 21 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/compilation/cuda_piecewise_backend.py b/python/sglang/srt/compilation/cuda_piecewise_backend.py index 1955fd3d5..8ca1e6a43 100644 --- a/python/sglang/srt/compilation/cuda_piecewise_backend.py +++ b/python/sglang/srt/compilation/cuda_piecewise_backend.py @@ -140,11 +140,6 @@ class CUDAPiecewiseBackend: if self.is_last_graph and not self.to_be_compiled_sizes: self.check_for_ending_compilation() - # Skip CUDA graphs if this entry doesn't use them OR - # if we're supposed to skip them globally - # skip_cuda_graphs = get_forward_context().skip_cuda_graphs - # if not entry.use_cudagraph or skip_cuda_graphs: - # return entry.runnable(*args) if is_in_pcg_torch_compile(): return entry.runnable(*args) @@ -172,7 +167,9 @@ class CUDAPiecewiseBackend: stack.enter_context(patch("torch.cuda.empty_cache", lambda: None)) # mind-exploding: carefully manage the reference and memory. stream = get_pcg_capture_stream() - assert stream is not None, "PCG capture stream is not set" + assert ( + stream is not None + ), "PCG capture stream is not set, please check if runtime recompilation happened" with torch.cuda.graph(cudagraph, pool=self.graph_pool, stream=stream): # `output` is managed by pytorch's cudagraph pool output = entry.runnable(*args) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 8eeec5ea3..f046c2775 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -950,10 +950,13 @@ class ServerArgs: self.cuda_graph_max_bs = max(self.cuda_graph_bs) if self.piecewise_cuda_graph_max_tokens is None: - # Capture piecewise cuda graph tokens up to the chunked prefill size. Two benefits: - # 1. cuda graph acceleration for all prefill lengths. - # 2. do not need more temporary memory for activations. Less fragmentation. - self.piecewise_cuda_graph_max_tokens = self.chunked_prefill_size + # Refer to pr #15927, by default we set the piecewise cuda graph max tokens to the chunked prefill size by default. + # For MLA backend, the introduction of piecewise cuda graph will influence the kernel dispatch difference compared to the original mode. + # To avoid the performance regression, we set the max tokens to 2048 by default. + if not self.use_mla_backend(): + self.piecewise_cuda_graph_max_tokens = self.chunked_prefill_size + else: + self.piecewise_cuda_graph_max_tokens = 2048 if self.piecewise_cuda_graph_tokens is None: self.piecewise_cuda_graph_tokens = ( @@ -983,6 +986,11 @@ class ServerArgs: if self.cuda_graph_max_bs > 300: reserved_mem += self.cuda_graph_max_bs * self.dp_size * 1.5 + # For piecewise cuda graphs + if self.enable_piecewise_cuda_graph: + # Only calculate the memory overhead for Non-Torch Memory use since the Torch Memory can be reused with Cuda Graph Capture + reserved_mem += len(self.piecewise_cuda_graph_tokens) * 8 + if gpu_mem is not None and gpu_mem > 60 * 1024: reserved_mem = max(reserved_mem, 10 * 1024) @@ -994,10 +1002,6 @@ class ServerArgs: # eagle draft models and cuda graphs reserved_mem += 2 * 1024 - # For piecewise cuda graphs - if self.enable_piecewise_cuda_graph: - reserved_mem += self.piecewise_cuda_graph_max_tokens // 8 - self.mem_fraction_static = ( round((gpu_mem - reserved_mem) / gpu_mem, 3) if gpu_mem is not None @@ -1049,8 +1053,9 @@ class ServerArgs: list(range(4, 33, 4)) + list(range(48, 257, 16)) + list(range(288, 513, 32)) - + list(range(640, 4096 + 1, 128)) - + list(range(4352, self.piecewise_cuda_graph_max_tokens + 1, 256)) + + list(range(640, 1024 + 1, 64)) + + list(range(1280, 4096 + 1, 256)) + + list(range(4608, self.piecewise_cuda_graph_max_tokens + 1, 512)) ) capture_sizes = [ @@ -4177,9 +4182,9 @@ class ServerArgs: ) parser.add_argument( "--piecewise-cuda-graph-tokens", - type=json_list_type, - default=ServerArgs.piecewise_cuda_graph_tokens, - help="Set the list of tokens when using piecewise cuda graph.", + type=int, + nargs="+", + help="Set the list of token lengths for piecewise cuda graph capture.", ) parser.add_argument( "--piecewise-cuda-graph-compiler",