diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index ea8c93963..f54e6f0a2 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -21,7 +21,7 @@ import inspect import logging import os from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, fields, is_dataclass from functools import partial from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union @@ -95,6 +95,13 @@ if TYPE_CHECKING: _has_foreach_copy = hasattr(torch, "_foreach_copy_") +# Opt-in (SGLANG_LOG_CG_BUFFERS=1): dump the per-buffer CUDA-graph memory +# breakdown at capture time. Startup-only, rank0; off by default. +_LOG_CG_BUFFERS = get_bool_env_var("SGLANG_LOG_CG_BUFFERS") +# Process-wide data_ptr -> first runner label, so the dump can flag buffers that +# share_buffers() coalesced into one physical allocation across the EAGLE families. +_cg_buffer_owner: Dict[int, str] = {} + def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None: """Call torch._foreach_copy_ grouped by (dst_dtype, src_dtype) pairs.""" @@ -501,6 +508,57 @@ def set_global_graph_memory_pool(val): global_graph_memory_pool = val +def _named_buffer_tensors(buffers) -> List[Tuple[str, torch.Tensor]]: + """Flatten a *InputBuffers dataclass into (name, tensor) pairs, mirroring + share_buffers()'s traversal (recurses into dict / dataclass fields).""" + out: List[Tuple[str, torch.Tensor]] = [] + for f in fields(buffers): + b = getattr(buffers, f.name) + if b is None: + continue + if is_dataclass(b): + b = vars(b) + if isinstance(b, dict): + for k, v in b.items(): + if torch.is_tensor(v): + out.append((f"{f.name}.{k}", v)) + elif torch.is_tensor(b): + out.append((f.name, b)) + return out + + +def log_input_buffer_sizes(label: str, buffers) -> None: + """One-time, rank0 dump of one runner's static CUDA-graph input/output + buffers: per-field shape/dtype/MB, flagging buffers that share_buffers() + aliased onto an earlier runner's allocation. Gated by SGLANG_LOG_CG_BUFFERS.""" + if not _LOG_CG_BUFFERS: + return + rows = _named_buffer_tensors(buffers) + rows.sort(key=lambda nt: -nt[1].numel() * nt[1].element_size()) + unique_mb = 0.0 + total_mb = 0.0 + lines = [] + for name, t in rows: + mb = t.numel() * t.element_size() / 2**20 + total_mb += mb + owner = _cg_buffer_owner.get(t.data_ptr()) + if owner is None: + _cg_buffer_owner[t.data_ptr()] = label + unique_mb += mb + tag = "" + else: + tag = f" " + lines.append( + f" {name:<40s} {str(tuple(t.shape)):<22s} " + f"{str(t.dtype).replace('torch.', ''):<10s} {mb:9.2f} MB{tag}" + ) + log_info_on_rank0( + logger, + f"[CG buffers] {label}: static input/output buffers " + f"(new={unique_mb:.2f} MB, incl-aliased={total_mb:.2f} MB)\n" + "\n".join(lines), + ) + + class CudaGraphRunner: """A CudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile.""" @@ -629,6 +687,7 @@ class CudaGraphRunner: ), ) self.buffers.share_buffers() + log_input_buffer_sizes(str(self.capture_forward_mode), self.buffers) self.tbo_plugin = TboCudaGraphRunnerPlugin() @@ -732,8 +791,18 @@ class CudaGraphRunner: torch.cuda.memory._record_memory_history() return profile_context + def _cuda_graph_mem_snapshot_path(self): + # Per-runner filename so the 3 EAGLE families don't overwrite one file. + label = str( + getattr(self, "capture_forward_mode", None) or type(self).__name__ + ) + for ch in ".:/ ": + label = label.replace(ch, "_") + return f"cuda_graph_mem_{label}.pickle" + def _post_process_after_profile(self, prof_context): - torch.cuda.memory._dump_snapshot(f"cuda_graph_runner_memory_usage.pickle") + snapshot_path = self._cuda_graph_mem_snapshot_path() + torch.cuda.memory._dump_snapshot(snapshot_path) torch.cuda.memory._record_memory_history(enabled=None) log_message = ( "Sorted by CUDA Time:\n" @@ -744,7 +813,7 @@ class CudaGraphRunner: + prof_context.key_averages(group_by_input_shape=True).table( sort_by="cpu_time_total", row_limit=10 ) - + "\n\nMemory Usage is saved to cuda_graph_runner_memory_usage.pickle\n" + + f"\n\nMemory Usage is saved to {snapshot_path}\n" ) logger.info(log_message) @@ -759,14 +828,20 @@ class CudaGraphRunner: self.model_runner.gpu_id, empty_cache=False, ) + is_rank0 = get_tensor_model_parallel_rank() == 0 + cg_label = str( + getattr(self, "capture_forward_mode", None) or type(self).__name__ + ) + capture_start_avail = avail_mem + prev_avail = avail_mem # Reverse the order to enable better memory sharing across cuda graphs. capture_range = ( tqdm.tqdm(list(reversed(self.capture_bs))) - if get_tensor_model_parallel_rank() == 0 + if is_rank0 else reversed(self.capture_bs) ) for i, bs in enumerate(capture_range): - if get_tensor_model_parallel_rank() == 0: + if is_rank0: avail_mem = get_available_gpu_memory( self.model_runner.device, self.model_runner.gpu_id, @@ -791,6 +866,33 @@ class CudaGraphRunner: self.graphs[key] = graph self.output_buffers[key] = output_buffers + if _LOG_CG_BUFFERS and is_rank0: + after = get_available_gpu_memory( + self.model_runner.device, + self.model_runner.gpu_id, + empty_cache=False, + ) + log_info_on_rank0( + logger, + f"[CG capture] {cg_label} bs={bs:>4d}: " + f"+{max(prev_avail - after, 0.0):6.3f} GB this shape " + f"(graph pool / attn workspace / outputs), avail={after:.3f} GB", + ) + prev_avail = after + + if _LOG_CG_BUFFERS and is_rank0: + final_avail = get_available_gpu_memory( + self.model_runner.device, + self.model_runner.gpu_id, + empty_cache=False, + ) + log_info_on_rank0( + logger, + f"[CG capture] {cg_label} TOTAL non-static-buffer growth: " + f"{max(capture_start_avail - final_avail, 0.0):.3f} GB " + f"across {len(self.capture_bs)} shapes", + ) + # Trigger CUDA graph capture for specific shapes. # Capture the large shapes first so that the smaller shapes # can reuse the memory pool allocated for the large shapes. diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 40e859b2d..f80a38a64 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -13,6 +13,7 @@ from sglang.srt.model_executor.cuda_graph_runner import ( DeepEPCudaGraphRunnerAdapter, get_batch_sizes_to_capture, get_global_graph_memory_pool, + log_input_buffer_sizes, model_capture_mode, set_global_graph_memory_pool, set_is_extend_in_batch, @@ -157,6 +158,7 @@ class EAGLEDraftCudaGraphRunner: global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu, ) self.buffers.share_buffers() + log_input_buffer_sizes("eagle_draft", self.buffers) # Capture try: diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 0bc5886a0..e6c3ee5cd 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -14,6 +14,7 @@ from sglang.srt.model_executor.cuda_graph_runner import ( LogitsProcessorOutput, get_batch_sizes_to_capture, get_global_graph_memory_pool, + log_input_buffer_sizes, model_capture_mode, set_global_graph_memory_pool, set_is_extend_in_batch, @@ -211,6 +212,7 @@ class EAGLEDraftExtendCudaGraphRunner: global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu, ) self.buffers.share_buffers() + log_input_buffer_sizes("eagle_draft_extend", self.buffers) # Capture try: