CUDA graph: opt-in per-buffer + per-shape decode CG memory logging

Attribute decode CUDA-graph GPU memory at capture time, gated by
SGLANG_LOG_CG_BUFFERS=1 (off by default, startup-only, rank0):
  1. Per-runner static input/output buffer table (name/shape/dtype/MB)
     across the 3 EAGLE families (target-verify, draft, draft-extend),
     flagging buffers that share_buffers() aliased onto an earlier family.
  2. Per-shape graph-pool growth deltas in the capture loop + a per-family
     TOTAL, isolating the shared activation pool from per-shape outputs.
Also fix --enable-profile-cuda-graph dumping all families to one
overwritten cuda_graph_runner_memory_usage.pickle: now writes a unique
cuda_graph_mem_<family>.pickle per runner.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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"   <alias of {owner}>"
+        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.

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:

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: