chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion]: clear file-path-only outputs on all ranks to prevent TP GPU memory skew (#20353)

Browse Source
This commit is contained in:
Ratish P
2026-03-12 14:59:09 +05:30
committed by GitHub
parent 5f1bfb0d28
commit 4e5ca92249
1 changed files with 11 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
View File

@@ -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: