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

python/sglang/multimodal_gen/benchmarks/bench_serving.py

@@ -63,6 +63,7 @@ class RequestFuncOutput:
     error: str = ""
     start_time: float = 0.0
     response_body: Dict[str, Any] = field(default_factory=dict)
+    peak_memory_mb: float = 0.0
 
 
 class BaseDataset(ABC):
@@ -371,6 +372,8 @@ async def async_request_image_sglang(
                     resp_json = await response.json()
                     output.response_body = resp_json
                     output.success = True
+                    if "peak_memory_mb" in resp_json:
+                        output.peak_memory_mb = resp_json["peak_memory_mb"]
                 else:
                     output.error = f"HTTP {response.status}: {await response.text()}"
                     output.success = False
@@ -398,6 +401,8 @@ async def async_request_image_sglang(
                     resp_json = await response.json()
                     output.response_body = resp_json
                     output.success = True
+                    if "peak_memory_mb" in resp_json:
+                        output.peak_memory_mb = resp_json["peak_memory_mb"]
                 else:
                     output.error = f"HTTP {response.status}: {await response.text()}"
                     output.success = False
@@ -406,6 +411,7 @@ async def async_request_image_sglang(
             output.success = False
 
     output.latency = time.perf_counter() - output.start_time
+
     if pbar:
         pbar.update(1)
     return output
@@ -537,6 +543,8 @@ async def async_request_video_sglang(
                     if status == "completed":
                         output.success = True
                         output.response_body = status_data
+                        if "peak_memory_mb" in status_data:
+                            output.peak_memory_mb = status_data["peak_memory_mb"]
                         break
                     elif status == "failed":
                         output.success = False
@@ -557,6 +565,7 @@ async def async_request_video_sglang(
             break
 
     output.latency = time.perf_counter() - output.start_time
+
     if pbar:
         pbar.update(1)
     return output
@@ -568,6 +577,7 @@ def calculate_metrics(outputs: List[RequestFuncOutput], total_duration: float):
 
     num_success = len(success_outputs)
     latencies = [o.latency for o in success_outputs]
+    peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0]
 
     metrics = {
         "duration": total_duration,
@@ -578,6 +588,9 @@ def calculate_metrics(outputs: List[RequestFuncOutput], total_duration: float):
         "latency_median": np.median(latencies) if latencies else 0,
         "latency_p99": np.percentile(latencies, 99) if latencies else 0,
         "latency_p50": np.percentile(latencies, 50) if latencies else 0,
+        "peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
+        "peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
+        "peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
     }
 
     return metrics
@@ -719,6 +732,24 @@ async def benchmark(args):
     print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
     print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
 
+    if metrics["peak_memory_mb_max"] > 0:
+        print(f"{'-' * 50}")
+        print(
+            "{:<40} {:<15.2f}".format(
+                "Peak Memory Max (MB):", metrics["peak_memory_mb_max"]
+            )
+        )
+        print(
+            "{:<40} {:<15.2f}".format(
+                "Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"]
+            )
+        )
+        print(
+            "{:<40} {:<15.2f}".format(
+                "Peak Memory Median (MB):", metrics["peak_memory_mb_median"]
+            )
+        )
+
     print("\n" + "=" * 60)
 
     if args.output_file:

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -209,6 +209,7 @@ class DiffGenerator:
 
         results = []
         total_start_time = time.perf_counter()
+
         # 2. send requests to scheduler, one at a time
         # TODO: send batch when supported
         for request_idx, req in enumerate(requests):
@@ -245,6 +246,7 @@ class DiffGenerator:
                             "prompts": req.prompt,
                             "size": (req.height, req.width, req.num_frames),
                             "generation_time": timer.duration,
+                            "peak_memory_mb": output_batch.peak_memory_mb,
                             "timings": (
                                 output_batch.timings.to_dict()
                                 if output_batch.timings
@@ -262,6 +264,16 @@ class DiffGenerator:
         total_gen_time = time.perf_counter() - total_start_time
         log_batch_completion(logger, len(results), total_gen_time)
 
+        if results:
+            peak_memories = [r.get("peak_memory_mb", 0) for r in results]
+            if peak_memories:
+                max_peak_memory = max(peak_memories)
+                avg_peak_memory = sum(peak_memories) / len(peak_memories)
+                logger.info(
+                    f"Memory usage - Max peak: {max_peak_memory:.2f} MB, "
+                    f"Avg peak: {avg_peak_memory:.2f} MB"
+                )
+
         if len(results) == 0:
             return None
         else:

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -122,7 +122,9 @@ async def generations(
     )
 
     # Run synchronously for images and save to disk
-    save_file_path = await process_generation_batch(async_scheduler_client, batch)
+    save_file_path, result = await process_generation_batch(
+        async_scheduler_client, batch
+    )
 
     await IMAGE_STORE.upsert(
         request_id,
@@ -137,14 +139,17 @@ async def generations(
     if resp_format == "b64_json":
         with open(save_file_path, "rb") as f:
             b64 = base64.b64encode(f.read()).decode("utf-8")
-        return ImageResponse(
+        response_kwargs = {
-            data=[
+            "data": [
                 ImageResponseData(
                     b64_json=b64,
                     revised_prompt=request.prompt,
                 )
             ]
-        )
+        }
+        if result.peak_memory_mb and result.peak_memory_mb > 0:
+            response_kwargs["peak_memory_mb"] = result.peak_memory_mb
+        return ImageResponse(**response_kwargs)
     else:
         # Return error, not supported
         raise HTTPException(
@@ -219,7 +224,9 @@ async def edits(
     )
     batch = _build_req_from_sampling(sampling)
 
-    save_file_path = await process_generation_batch(async_scheduler_client, batch)
+    save_file_path, result = await process_generation_batch(
+        async_scheduler_client, batch
+    )
 
     await IMAGE_STORE.upsert(
         request_id,
@@ -236,12 +243,18 @@ async def edits(
     if (response_format or "b64_json").lower() == "b64_json":
         with open(save_file_path, "rb") as f:
             b64 = base64.b64encode(f.read()).decode("utf-8")
-        return ImageResponse(
+        response_kwargs = {
-            data=[ImageResponseData(b64_json=b64, revised_prompt=prompt)]
+            "data": [ImageResponseData(b64_json=b64, revised_prompt=prompt)]
-        )
+        }
+        if result.peak_memory_mb and result.peak_memory_mb > 0:
+            response_kwargs["peak_memory_mb"] = result.peak_memory_mb
+        return ImageResponse(**response_kwargs)
     else:
         url = f"/v1/images/{request_id}/content"
-        return ImageResponse(data=[ImageResponseData(url=url, revised_prompt=prompt)])
+        response_kwargs = {"data": [ImageResponseData(url=url, revised_prompt=prompt)]}
+        if result.peak_memory_mb and result.peak_memory_mb > 0:
+            response_kwargs["peak_memory_mb"] = result.peak_memory_mb
+        return ImageResponse(**response_kwargs)
 
 
 @router.get("/{image_id}/content")

python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py

@@ -14,6 +14,7 @@ class ImageResponseData(BaseModel):
 class ImageResponse(BaseModel):
     created: int = Field(default_factory=lambda: int(time.time()))
     data: List[ImageResponseData]
+    peak_memory_mb: Optional[float] = None
 
 
 class ImageGenerationsRequest(BaseModel):
@@ -50,6 +51,7 @@ class VideoResponse(BaseModel):
     completed_at: Optional[int] = None
     expires_at: Optional[int] = None
     error: Optional[Dict[str, Any]] = None
+    peak_memory_mb: Optional[float] = None
 
 
 class VideoGenerationsRequest(BaseModel):

python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py

@@ -182,7 +182,10 @@ async def process_generation_batch(
     total_time = time.perf_counter() - total_start_time
     log_batch_completion(logger, 1, total_time)
 
-    return save_file_path
+    if result.peak_memory_mb and result.peak_memory_mb > 0:
+        logger.info(f"Peak memory usage: {result.peak_memory_mb:.2f} MB")
+
+    return save_file_path, result
 
 
 def merge_image_input_list(*inputs: Union[List, Any, None]) -> List:

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -118,11 +118,15 @@ async def _dispatch_job_async(job_id: str, batch: Req) -> None:
     from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 
     try:
-        await process_generation_batch(async_scheduler_client, batch)
+        _, result = await process_generation_batch(async_scheduler_client, batch)
-        await VIDEO_STORE.update_fields(
+        update_fields = {
-            job_id,
+            "status": "completed",
-            {"status": "completed", "progress": 100, "completed_at": int(time.time())},
+            "progress": 100,
-        )
+            "completed_at": int(time.time()),
+        }
+        if result.peak_memory_mb and result.peak_memory_mb > 0:
+            update_fields["peak_memory_mb"] = result.peak_memory_mb
+        await VIDEO_STORE.update_fields(job_id, update_fields)
     except Exception as e:
         logger.error(f"{e}")
         await VIDEO_STORE.update_fields(

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)

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

@@ -281,3 +281,4 @@ class OutputBatch:
 
     # logged timings info, directly from Req.timings
     timings: Optional["RequestTimings"] = None
+    peak_memory_mb: float = 0.0