From 53ebdd0b1d39e89e22e660f84db347fd0f1192f3 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Wed, 18 Mar 2026 22:59:58 +0800 Subject: [PATCH] [Bug]: fix CUDA OOM during diffusion post-processing (#1670) Signed-off-by: lishunyang --- vllm_omni/diffusion/diffusion_engine.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 87b621d7b..34f18f4df 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -8,6 +8,7 @@ from typing import Any import numpy as np import PIL.Image +import torch from vllm.logger import init_logger from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -106,8 +107,19 @@ class DiffusionEngine: for i, prompt in enumerate(request.prompts) ] + # When CPU offload is enabled, move output to CPU before + # post-processing to avoid device OOM — model weights may still + # reside on the device and leave no headroom for intermediates. + output_data = output.output + if ( + self.od_config.enable_cpu_offload + and isinstance(output_data, torch.Tensor) + and output_data.device.type != "cpu" + ): + output_data = output_data.cpu() + postprocess_start_time = time.perf_counter() - outputs = self.post_process_func(output.output) if self.post_process_func is not None else output.output + outputs = self.post_process_func(output_data) if self.post_process_func is not None else output_data audio_payload = None if isinstance(outputs, dict): audio_payload = outputs.get("audio")