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

[diffusion] fix: preserve tensor stride when offloading rollout weights to pinned host memory (#32420)

Browse Source
This commit is contained in:
Kangrui Du
2026-07-27 12:05:13 -07:00
committed by GitHub
parent 1da062f018
commit 8a311d1c88
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py
+10 -1
View File
@@ -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. # Async D2H into pinned host memory; caller synchronizes once after the batch.
for t in list(module.parameters()) + list(module.buffers()): for t in list(module.parameters()) + list(module.buffers()):
if t.device.type == "cuda": 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) pin.copy_(t.data, non_blocking=True)
t.data = pin t.data = pin