From d77f3fccbfddb7c1450cdd64a72d4b5b4f724edc Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Mon, 22 Dec 2025 21:21:21 +0800 Subject: [PATCH] [Diffusion] Support peak memory record in offline generate and serving (#15610) --- .../benchmarks/bench_serving.py | 31 +++++++++++++++++++ .../entrypoints/diffusion_generator.py | 12 +++++++ .../runtime/entrypoints/openai/image_api.py | 31 +++++++++++++------ .../runtime/entrypoints/openai/protocol.py | 2 ++ .../runtime/entrypoints/openai/utils.py | 5 ++- .../runtime/entrypoints/openai/video_api.py | 14 ++++++--- .../runtime/managers/gpu_worker.py | 7 +++++ .../runtime/pipelines_core/schedule_batch.py | 1 + 8 files changed, 88 insertions(+), 15 deletions(-) diff --git a/python/sglang/multimodal_gen/benchmarks/bench_serving.py b/python/sglang/multimodal_gen/benchmarks/bench_serving.py index 8960c1391..65106183e 100644 --- a/python/sglang/multimodal_gen/benchmarks/bench_serving.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index c42e1515d..993e8b6aa 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index ab3af9988..8a2ace597 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/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( - data=[ + response_kwargs = { + "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( - data=[ImageResponseData(b64_json=b64, revised_prompt=prompt)] - ) + response_kwargs = { + "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") diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py index 610f00415..8b894928f 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py +++ b/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): diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index 08d658300..0b61e44e3 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index d467d0d45..a75c41fc3 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/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) - await VIDEO_STORE.update_fields( - job_id, - {"status": "completed", "progress": 100, "completed_at": int(time.time())}, - ) + _, result = await process_generation_batch(async_scheduler_client, batch) + update_fields = { + "status": "completed", + "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( diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index a3c414040..8358b35b3 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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) 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 df878ace4..ea4a62de8 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/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