From ab9de886c57bfa4160cf8f41c4d97c23df7e5af0 Mon Sep 17 00:00:00 2001 From: Ratish P <114130421+Ratish1@users.noreply.github.com> Date: Sun, 8 Mar 2026 16:56:17 +0530 Subject: [PATCH] [diffusion] reduce LayerwiseOffloadManager reserved GPU memory (#20042) --- .../runtime/utils/layerwise_offload.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py index b6b782a01..05a3d676d 100644 --- a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py @@ -66,6 +66,7 @@ class LayerwiseOffloadManager: self._named_parameters: Dict[str, torch.nn.Parameter] = {} self._named_buffers: Dict[str, torch.Tensor] = {} + self._offload_placeholders: Dict[torch.dtype, torch.Tensor] = {} # Store forward hooks for removal self._forward_hooks: List[Any] = [] @@ -80,6 +81,13 @@ class LayerwiseOffloadManager: except Exception: return None + def _get_shared_empty_tensor(self, dtype: torch.dtype) -> torch.Tensor: + placeholder = self._offload_placeholders.get(dtype) + if placeholder is None: + placeholder = torch.empty((1,), device=self.device, dtype=dtype) + self._offload_placeholders[dtype] = placeholder + return placeholder + @torch.compiler.disable def _initialize(self) -> None: if not self.enabled: @@ -126,7 +134,7 @@ class LayerwiseOffloadManager: "shape": weight.shape, } - weight.data = torch.empty((1,), device=self.device, dtype=dtype) + weight.data = self._get_shared_empty_tensor(dtype) current_offset += numel @@ -212,15 +220,13 @@ class LayerwiseOffloadManager: # clear prefetch event, since it's useless and needs to be reset self._prefetch_events.pop(layer_idx, None) - if layer_idx <= 0: - return - if layer_idx not in self._gpu_layers: return for name, meta in self._weight_metadata.get(layer_idx, {}).items(): target = self.get_target_with_name(name) - target.data = torch.empty((1,), device=self.device, dtype=meta["dtype"]) + # Wraparound prefetch will reload the layer when it is needed again + target.data = self._get_shared_empty_tensor(meta["dtype"]) self._gpu_layers.discard(layer_idx)