From 8a311d1c889244ab1f857d7df79de7e5f0a6891c Mon Sep 17 00:00:00 2001
From: Kangrui Du <89372739+Rockdu@users.noreply.github.com>
Date: Mon, 27 Jul 2026 12:05:13 -0700
Subject: [PATCH] [diffusion] fix: preserve tensor stride when offloading
 rollout weights to pinned host memory (#32420)

---
 .../memory_managers/memory_occupation_controller.py   | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py
index 4a3d55d76..612445b8e 100644
--- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py
+++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py
@@ -20,7 +20,16 @@ def _module_to_pinned_cpu(module: torch.nn.Module) -> None:
     # Async D2H into pinned host memory; caller synchronizes once after the batch.
     for t in list(module.parameters()) + list(module.buffers()):
         if t.device.type == "cuda":
-            pin = torch.empty(t.shape, dtype=t.dtype, device="cpu", pin_memory=True)
+            # Mirror stride/layout like srt/utils/offloader.py: torch.empty() would force
+            # contiguous and silently drop channels_last_3d VAE weights.
+            pin = torch.empty_strided(
+                size=t.size(),
+                stride=t.stride(),
+                dtype=t.dtype,
+                layout=t.layout,
+                device="cpu",
+                pin_memory=True,
+            )
             pin.copy_(t.data, non_blocking=True)
             t.data = pin