[diffusion] logging: improve peak vram logging (#18865)

python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py

@@ -21,8 +21,9 @@ from sglang.multimodal_gen.runtime.entrypoints.cli.utils import (
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.perf_logger import (
+    MemorySnapshot,
     PerformanceLogger,
-    RequestTimings,
+    RequestMetrics,
 )
 from sglang.multimodal_gen.utils import FlexibleArgumentParser
 
@@ -72,11 +73,22 @@ def maybe_dump_performance(args: argparse.Namespace, server_args, prompt: str, r
     if not (args.perf_dump_path and timings_dict):
         return
 
-    timings = RequestTimings(request_id=timings_dict.get("request_id"))
+    timings = RequestMetrics(request_id=timings_dict.get("request_id"))
     timings.stages = timings_dict.get("stages", {})
     timings.steps = timings_dict.get("steps", [])
     timings.total_duration_ms = timings_dict.get("total_duration_ms", 0)
 
+    # restore memory snapshots from serialized dict
+    memory_snapshots_dict = timings_dict.get("memory_snapshots", {})
+    for checkpoint_name, snapshot_dict in memory_snapshots_dict.items():
+        snapshot = MemorySnapshot(
+            allocated_mb=snapshot_dict.get("allocated_mb", 0.0),
+            reserved_mb=snapshot_dict.get("reserved_mb", 0.0),
+            peak_allocated_mb=snapshot_dict.get("peak_allocated_mb", 0.0),
+            peak_reserved_mb=snapshot_dict.get("peak_reserved_mb", 0.0),
+        )
+        timings.memory_snapshots[checkpoint_name] = snapshot
+
     PerformanceLogger.dump_benchmark_report(
         file_path=args.perf_dump_path,
         timings=timings,

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -45,7 +45,10 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import (
     globally_suppress_loggers,
     init_logger,
 )
-from sglang.multimodal_gen.runtime.utils.perf_logger import PerformanceLogger
+from sglang.multimodal_gen.runtime.utils.perf_logger import (
+    PerformanceLogger,
+    capture_memory_snapshot,
+)
 
 logger = init_logger(__name__)
 
@@ -146,11 +149,20 @@ class GPUWorker:
         )
 
     def do_mem_analysis(self, output_batch: OutputBatch):
-        peak_memory_bytes = torch.cuda.max_memory_allocated()
+        final_snapshot = capture_memory_snapshot()
-        output_batch.peak_memory_mb = peak_memory_bytes / (1024**2)
+        if output_batch.timings:
-        peak_memory_gb = peak_memory_bytes / (1024**3)
+            output_batch.timings.record_memory_snapshot("mem_analysis", final_snapshot)
+
+        # for details on max_memory_reserved: https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.max_memory_reserved.html
+        peak_reserved_bytes = torch.cuda.max_memory_reserved()
+        peak_allocated_bytes = torch.cuda.max_memory_allocated()
+
+        output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
+        peak_reserved_gb = peak_reserved_bytes / (1024**3)
+        peak_allocated_gb = peak_allocated_bytes / (1024**3)
+
         remaining_gpu_mem_gb = (
-            current_platform.get_device_total_memory() / (1024**3) - peak_memory_gb
+            current_platform.get_device_total_memory() / (1024**3) - peak_reserved_gb
         )
         can_stay_resident = self.get_can_stay_resident_components(remaining_gpu_mem_gb)
         suggested_args = set()
@@ -173,8 +185,13 @@ class GPUWorker:
         suggested_args_str = (
             ", ".join(sorted(suggested_args)) if suggested_args else "None"
         )
+
+        pool_overhead_gb = peak_reserved_gb - peak_allocated_gb
+
         logger.info(
-            f"Peak GPU memory: {peak_memory_gb:.2f} GB, "
+            f"Peak GPU memory: {peak_reserved_gb:.2f} GB, "
+            f"Peak allocated: {peak_allocated_gb:.2f} GB, "
+            f"Memory pool overhead: {pool_overhead_gb:.2f} GB ({pool_overhead_gb/peak_reserved_gb*100:.1f}%), "
             f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. "
             f"Components that could stay resident (based on the last request workload): {can_stay_resident}. "
             f"Related offload server args to disable: {suggested_args_str}"
@@ -193,6 +210,11 @@ class GPUWorker:
 
             start_time = time.monotonic()
 
+            # capture memory baseline before forward
+            if self.rank == 0 and req.timings:
+                baseline_snapshot = capture_memory_snapshot()
+                req.timings.record_memory_snapshot("before_forward", baseline_snapshot)
+
             req.log(server_args=self.server_args)
             result = self.pipeline.forward(req, self.server_args)
 
@@ -210,6 +232,13 @@ class GPUWorker:
             else:
                 output_batch = result
 
+            # capture memory after forward (peak)
+            if self.rank == 0 and output_batch.timings:
+                peak_snapshot = capture_memory_snapshot()
+                output_batch.timings.record_memory_snapshot(
+                    "after_forward", peak_snapshot
+                )
+
             if self.rank == 0 and not req.suppress_logs:
                 self.do_mem_analysis(output_batch)
 

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -26,7 +26,7 @@ from sglang.multimodal_gen.configs.sample.teacache import (
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-from sglang.multimodal_gen.runtime.utils.perf_logger import RequestTimings
+from sglang.multimodal_gen.runtime.utils.perf_logger import RequestMetrics
 from sglang.multimodal_gen.utils import align_to
 
 logger = init_logger(__name__)
@@ -152,7 +152,7 @@ class Req:
     VSA_sparsity: float = 0.0
 
     # stage logging
-    timings: Optional["RequestTimings"] = None
+    timings: Optional["RequestMetrics"] = None
 
     # results
     output: torch.Tensor | None = None
@@ -267,7 +267,7 @@ class Req:
         if self.guidance_scale_2 is None:
             self.guidance_scale_2 = self.guidance_scale
 
-        self.timings = RequestTimings(request_id=self.request_id)
+        self.timings = RequestMetrics(request_id=self.request_id)
 
         if self.is_warmup:
             self.set_as_warmup()
@@ -329,8 +329,8 @@ class OutputBatch:
     error: str | None = None
     output_file_paths: list[str] | None = None
 
-    # logged timings info, directly from Req.timings
+    # logged metrics info, directly from Req.timings
-    timings: Optional["RequestTimings"] = None
+    timings: Optional["RequestMetrics"] = None
 
     # For ComfyUI integration: noise prediction from denoising stage
     noise_pred: torch.Tensor | None = None

python/sglang/multimodal_gen/runtime/utils/perf_logger.py

@@ -25,14 +25,32 @@ logger = init_logger(__name__)
 
 
 @dataclasses.dataclass
-class RequestTimings:
+class MemorySnapshot:
-    """A lightweight data class to store performance timings for a single request."""
+    allocated_mb: float  # current allocated memory
+    reserved_mb: float  # current reserved memory (actual VRAM)
+    peak_allocated_mb: float  # peak allocated since last reset
+    peak_reserved_mb: float  # peak reserved since last reset
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "allocated_mb": round(self.allocated_mb, 2),
+            "reserved_mb": round(self.reserved_mb, 2),
+            "peak_allocated_mb": round(self.peak_allocated_mb, 2),
+            "peak_reserved_mb": round(self.peak_reserved_mb, 2),
+        }
+
+
+@dataclasses.dataclass
+class RequestMetrics:
+    """Performance metrics for a single request, including timings and memory snapshots."""
 
     def __init__(self, request_id: str):
         self.request_id = request_id
         self.stages: Dict[str, float] = {}
         self.steps: list[float] = []
         self.total_duration_ms: float = 0.0
+        # memory tracking: {checkpoint_name: MemorySnapshot}
+        self.memory_snapshots: Dict[str, MemorySnapshot] = {}
 
     @property
     def total_duration_s(self) -> float:
@@ -47,13 +65,20 @@ class RequestTimings:
         assert index == len(self.steps)
         self.steps.append(duration_s * 1000)
 
+    def record_memory_snapshot(self, checkpoint_name: str, snapshot: MemorySnapshot):
+        self.memory_snapshots[checkpoint_name] = snapshot
+
     def to_dict(self) -> Dict[str, Any]:
-        """Serializes the timing data to a dictionary."""
+        """Serializes the metrics data to a dictionary."""
         return {
             "request_id": self.request_id,
             "stages": self.stages,
             "steps": self.steps,
             "total_duration_ms": self.total_duration_ms,
+            "memory_snapshots": {
+                name: snapshot.to_dict()
+                for name, snapshot in self.memory_snapshots.items()
+            },
         }
 
 
@@ -90,6 +115,28 @@ def get_git_commit_hash() -> str:
         return "N/A"
 
 
+def capture_memory_snapshot() -> MemorySnapshot:
+    if not torch.cuda.is_available():
+        return MemorySnapshot(
+            allocated_mb=0.0,
+            reserved_mb=0.0,
+            peak_allocated_mb=0.0,
+            peak_reserved_mb=0.0,
+        )
+
+    allocated = torch.cuda.memory_allocated()
+    reserved = torch.cuda.memory_reserved()
+    peak_allocated = torch.cuda.max_memory_allocated()
+    peak_reserved = torch.cuda.max_memory_reserved()
+
+    return MemorySnapshot(
+        allocated_mb=allocated / (1024**2),
+        reserved_mb=reserved / (1024**2),
+        peak_allocated_mb=peak_allocated / (1024**2),
+        peak_reserved_mb=peak_reserved / (1024**2),
+    )
+
+
 @dataclasses.dataclass
 class RequestPerfRecord:
     request_id: str
@@ -101,6 +148,7 @@ class RequestPerfRecord:
     stages: list[dict]
     steps: list[float]
     total_duration_ms: float
+    memory_snapshots: dict[str, dict] = dataclasses.field(default_factory=dict)
 
     def __init__(
         self,
@@ -110,6 +158,7 @@ class RequestPerfRecord:
         stages,
         steps,
         total_duration_ms,
+        memory_snapshots=None,
         timestamp=None,
     ):
         self.request_id = request_id
@@ -123,20 +172,22 @@ class RequestPerfRecord:
         self.stages = stages
         self.steps = steps
         self.total_duration_ms = total_duration_ms
+        self.memory_snapshots = memory_snapshots or {}
 
 
 class StageProfiler:
     """
-    A unified context manager, records timing information (usually of a single Stage or a step) into a provided RequestTimings object (usually from a Req).
+    A unified context manager, records performance metrics (usually of a single Stage or a step) into a provided RequestMetrics object (usually from a Req).
     """
 
     def __init__(
         self,
         stage_name: str,
         logger: _SGLDiffusionLogger,
-        timings: Optional["RequestTimings"],
+        timings: Optional["RequestMetrics"],
         log_stage_start_end: bool = False,
         perf_dump_path_provided: bool = False,
+        capture_memory: bool = False,
     ):
         self.stage_name = stage_name
         self.timings = timings
@@ -144,6 +195,7 @@ class StageProfiler:
         self.start_time = 0.0
         self.log_timing = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING
         self.log_stage_start_end = log_stage_start_end
+        self.capture_memory = capture_memory
 
     def __enter__(self):
         if self.log_stage_start_end:
@@ -194,6 +246,13 @@ class StageProfiler:
             else:
                 self.timings.record_stage(self.stage_name, execution_time_s)
 
+            # capture memory snapshot after stage if requested
+            if self.capture_memory and torch.cuda.is_available():
+                snapshot = capture_memory_snapshot()
+                self.timings.record_memory_snapshot(
+                    f"after_{self.stage_name}", snapshot
+                )
+
         return False
 
 
@@ -203,14 +262,14 @@ class PerformanceLogger:
 
     Serves both as a runtime logger (stream to file) and a dump utility.
 
-    Notice that ""RequestTimings"" stores the performance metrics of a single request
+    Notice that RequestMetrics stores the performance metrics of a single request
     """
 
     @classmethod
     def dump_benchmark_report(
         cls,
         file_path: str,
-        timings: "RequestTimings",
+        timings: "RequestMetrics",
         meta: Optional[Dict[str, Any]] = None,
         tag: str = "benchmark_dump",
     ):
@@ -228,6 +287,11 @@ class PerformanceLogger:
             for idx, duration_ms in enumerate(timings.steps)
         ]
 
+        memory_checkpoints = {
+            name: snapshot.to_dict()
+            for name, snapshot in timings.memory_snapshots.items()
+        }
+
         report = {
             "timestamp": datetime.now(UTC).isoformat(),
             "request_id": timings.request_id,
@@ -236,6 +300,7 @@ class PerformanceLogger:
             "total_duration_ms": timings.total_duration_ms,
             "steps": formatted_steps,
             "denoise_steps_ms": denoise_steps_ms,
+            "memory_checkpoints": memory_checkpoints,
             "meta": meta or {},
         }
 
@@ -251,7 +316,7 @@ class PerformanceLogger:
     @classmethod
     def log_request_summary(
         cls,
-        timings: "RequestTimings",
+        timings: "RequestMetrics",
         tag: str = "total_inference_time",
     ):
         """logs the stage metrics and total duration for a completed request
@@ -264,6 +329,11 @@ class PerformanceLogger:
             for name, duration_ms in timings.stages.items()
         ]
 
+        memory_checkpoints = {
+            name: snapshot.to_dict()
+            for name, snapshot in timings.memory_snapshots.items()
+        }
+
         record = RequestPerfRecord(
             timings.request_id,
             commit_hash=get_git_commit_hash(),
@@ -271,6 +341,7 @@ class PerformanceLogger:
             stages=formatted_stages,
             steps=timings.steps,
             total_duration_ms=timings.total_duration_ms,
+            memory_snapshots=memory_checkpoints,
         )
 
         try: