From cc54d8e8d0b3c66bb59a1a281db9c52caf2d815a Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 5 May 2026 22:30:43 +0800 Subject: [PATCH] [diffusion] chore: clean CUDA cache only at explicit release points (#24397) --- .../runtime/managers/component_manager.py | 28 +++++++++++++++++++ .../runtime/managers/gpu_worker.py | 3 ++ .../runtime/utils/perf_logger.py | 8 +++++- .../test/server/perf_baselines.json | 12 ++++---- .../test/server/test_server_utils.py | 7 ++++- 5 files changed, 50 insertions(+), 8 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/component_manager.py index b43a16ead..ae35a6449 100644 --- a/python/sglang/multimodal_gen/runtime/managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/component_manager.py @@ -453,7 +453,9 @@ class ComponentResidencyManager: return strategy = self.strategy_for(use.component_name, module) self._trace("finish", use, strategy, module) + was_on_cuda = self._module_on_cuda(module) strategy.finish_use(module, use, self.state) + self._empty_cache_after_large_release(use, strategy, module, was_on_cuda) def finish_request(self) -> None: if not self.enabled and not self._uses_seen and self._active_use is None: @@ -492,7 +494,11 @@ class ComponentResidencyManager: else: action = "request_resident" if preferred else "request_finish" self._trace(action, use, strategy, module) + was_on_cuda = self._module_on_cuda(module) strategy.finish_request(module, use, self.state, preferred=preferred) + self._empty_cache_after_large_release( + use, strategy, module, was_on_cuda + ) self._trace("request_end") def stage_name(self, stage: ComponentResidencyStage) -> str: @@ -652,6 +658,28 @@ class ComponentResidencyManager: buffer = next(module.buffers(), None) return buffer.device.type if buffer is not None else None + def _module_on_cuda(self, module: nn.Module | None) -> bool: + return self._module_device(module) == "cuda" + + def _empty_cache_after_large_release( + self, + use: ComponentUse, + strategy: ComponentResidencyStrategy, + module: nn.Module, + was_on_cuda: bool, + ) -> None: + """explicitly empty cache after potential release of large component""" + if not use.memory_intensive: + return + released_cuda_storage = was_on_cuda and not self._module_on_cuda(module) + released_layerwise_storage = isinstance(strategy, LayerwiseOffloadStrategy) + if not (released_cuda_storage or released_layerwise_storage): + return + if not torch.get_device_module().is_available(): + return + torch.get_device_module().empty_cache() + self._trace("empty_cache", use, strategy, module, detail="after_release") + _GLOBAL_COMPONENT_RESIDENCY_MANAGER: ComponentResidencyManager | None = None diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 1a224d4f5..02438fd07 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -396,6 +396,9 @@ class GPUWorker: output_batch = OutputBatch() output_batch.error = f"Error executing {error_context}: {e}" self._record_output_peak_memory(output_batch) + # clean cache if OOM + if torch.cuda.is_initialized(): + torch.cuda.empty_cache() return output_batch def _record_output_peak_memory(self, output_batch: OutputBatch) -> None: diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 9b1eb68cf..0de2b363d 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py @@ -209,7 +209,13 @@ class StageProfiler: if self.log_stage_start_end: msg = f"[{self.stage_name}] started..." if self.logger.isEnabledFor(logging.DEBUG): - msg += f" ({round(current_platform.get_available_gpu_memory(), 2)} GB left)" + # This debug-only memory log runs at every stage boundary in CI. + # Keep it observational; cache cleanup is handled at explicit + # failure and component-release points. + available_memory = current_platform.get_available_gpu_memory( + empty_cache=False + ) + msg += f" ({round(available_memory, 2)} GB left)" self.logger.info(msg) if (self.log_timing and self.metrics) or self.log_stage_start_end: diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 7d5a7e96a..ffcd75394 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -755,8 +755,8 @@ }, "joyai_image_edit_ti2i": { "stages_ms": { - "InputValidationStage": 32.2, - "ImageEncodingStage": 948.69, + "InputValidationStage": 29.6, + "ImageEncodingStage": 740.87, "ImageVAEEncodingStage": 70.47, "LatentPreparationStage": 0.17, "TimestepPreparationStage": 20.66, @@ -1963,12 +1963,12 @@ "qwen_image_edit_2511_ti2i": { "stages_ms": { "DecodingStage": 19.39, - "InputValidationStage": 55.64, + "InputValidationStage": 54.97, "DenoisingStage": 22253.55, "ImageEncodingStage": 733.67, - "LatentPreparationStage": 0.19, + "LatentPreparationStage": 0.13, "TimestepPreparationStage": 12.36, - "ImageVAEEncodingStage": 96.46 + "ImageVAEEncodingStage": 84.65 }, "denoise_step_ms": { "0": 416.33, @@ -2012,7 +2012,7 @@ "38": 557.72, "39": 559.24 }, - "expected_e2e_ms": 23525.28, + "expected_e2e_ms": 23405.08, "expected_avg_denoise_ms": 556.18, "expected_median_denoise_ms": 560.34, "estimated_full_test_time_s": 143.7 diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index 8b28b9cf0..aedaaf1e6 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -619,12 +619,17 @@ class PerformanceValidator: if stage == "DenoisingStage" else self.tolerances.non_denoise_stage ) + if stage.endswith("DecodingStage"): + tolerance = max(tolerance, 0.9) + min_abs_tolerance_ms = 250.0 + else: + min_abs_tolerance_ms = 120.0 self._assert_le( f"Stage '{stage}'", actual, expected, tolerance, - min_abs_tolerance_ms=120.0, # relax absolute tolerance for non-denoising stages + min_abs_tolerance_ms=min_abs_tolerance_ms, )