[Diffusion] Support peak memory record in offline generate and serving (#15610)

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

@@ -97,6 +97,9 @@ class GPUWorker:
         req = batch[0]
         output_batch = None
         try:
+            if self.rank == 0:
+                torch.cuda.reset_peak_memory_stats()
+
             start_time = time.monotonic()
             timings = RequestTimings(request_id=req.request_id)
             req.timings = timings
@@ -104,6 +107,10 @@ class GPUWorker:
             output_batch = self.pipeline.forward(req, self.server_args)
             duration_ms = (time.monotonic() - start_time) * 1000
 
+            if self.rank == 0:
+                peak_memory_bytes = torch.cuda.max_memory_allocated()
+                output_batch.peak_memory_mb = peak_memory_bytes / (1024**2)
+
             if output_batch.timings:
                 output_batch.timings.total_duration_ms = duration_ms
                 PerformanceLogger.log_request_summary(timings=output_batch.timings)