diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 2b2d4eeb8..eb225a43c 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -254,8 +254,8 @@ class GPUWorker: # Save output to file and return file path only if requested. Avoid the serialization # and deserialization overhead between scheduler_client and gpu_worker. - if req.save_output and req.return_file_paths_only and self.rank == 0: - if output_batch.output is not None: + if req.save_output and req.return_file_paths_only: + if self.rank == 0 and output_batch.output is not None: output_paths = save_outputs( output_batch.output, req.data_type, @@ -274,7 +274,15 @@ class GPUWorker: upscaling_scale=req.upscaling_scale, ) output_batch.output_file_paths = output_paths - output_batch.output = None + + # No rank needs to hold on to generated tensors once the file-path + # response has been materialized on rank 0 + output_batch.output = None + output_batch.audio = None + output_batch.audio_sample_rate = None + + if torch.cuda.is_initialized(): + torch.cuda.empty_cache() # TODO: extract to avoid duplication if req.perf_dump_path is not None or envs.SGLANG_DIFFUSION_STAGE_LOGGING: