diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py b/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py index fb6554f4b..0c8b2c0c2 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 0b6995ced..bd92efbd2 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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() - output_batch.peak_memory_mb = peak_memory_bytes / (1024**2) - peak_memory_gb = peak_memory_bytes / (1024**3) + final_snapshot = capture_memory_snapshot() + if output_batch.timings: + 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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 8c5ff72c7..34c41a6fa 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/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 - timings: Optional["RequestTimings"] = None + # logged metrics info, directly from Req.timings + timings: Optional["RequestMetrics"] = None # For ComfyUI integration: noise prediction from denoising stage noise_pred: torch.Tensor | None = None diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 7c60d9352..badfe1420 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py @@ -25,14 +25,32 @@ logger = init_logger(__name__) @dataclasses.dataclass -class RequestTimings: - """A lightweight data class to store performance timings for a single request.""" +class MemorySnapshot: + 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: