[diffusion] chore: clean CUDA cache only at explicit release points (#24397)

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
 

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:

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:

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

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,
             )