From ab9750fb355edaa792b351fd3183a3141d9a6bd7 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 11 Sep 2026 22:04:40 +0800 Subject: [PATCH] [diffusion] optimization: pin layerwise host stores in place at their exact size (#39021) Co-authored-by: Claude Opus 5 --- .../memory_managers/layerwise_offload.py | 119 ++++++++++++++++-- .../test/unit/test_layerwise_offload.py | 37 ++++++ 2 files changed, 146 insertions(+), 10 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py index 44c2a0e66..ee96fe374 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py @@ -1,6 +1,7 @@ import bisect import ctypes import ctypes.util +import math import mmap import os import queue @@ -8,6 +9,7 @@ import re import sys import threading import time +import weakref from collections.abc import Iterator, Mapping, Sequence from contextlib import nullcontext from time import perf_counter @@ -395,6 +397,92 @@ _DIRECT_CHUNK = 64 << 20 MAPPED_DIRECT_READ_MIN_BYTES = 8 * 1024**3 +def _pin_in_place(tensor: torch.Tensor) -> bool: + """Page-lock an allocated CPU tensor in place, at exactly its size. + + torch's pinned pool rounds each block up to a power of two (a 1.20 GiB + layer store costs 2 GiB), which the pin budget never charged for. + """ + if not torch.cuda.is_available(): + return False + storage = tensor.untyped_storage() + nbytes = storage.nbytes() + if nbytes == 0: + return False + try: + cudart = torch.cuda.cudart() + if int(cudart.cudaHostRegister(storage.data_ptr(), nbytes, 0)) != 0: + return False + except Exception: + return False + weakref.finalize(storage, _unpin_in_place, storage.data_ptr()) + return True + + +def _unpin_in_place(data_ptr: int) -> None: + try: + torch.cuda.cudart().cudaHostUnregister(data_ptr) + except Exception: + pass + + +def _shared_storage(nbytes: int) -> Optional[torch.UntypedStorage]: + """Anonymous shared memory (memfd) of exactly `nbytes`, or None off Linux. + + Shared rather than private anonymous memory so the locked pages stay on + the shmem side of the process's accounting, where cudaHostAlloc kept + them: the host-anon figures (and the forced-host-view budget that + subtracts them) keep meaning "pageable copies". + """ + if not hasattr(os, "memfd_create"): + return None + try: + fd = os.memfd_create("sglang-pinned-store", 0) + except OSError: + return None + try: + os.ftruncate(fd, nbytes) + return torch.UntypedStorage.from_file( + f"/proc/self/fd/{fd}", shared=True, nbytes=nbytes + ) + except Exception: + return None + finally: + os.close(fd) + + +# Below this, torch's own pinned pool: its power-of-two rounding costs little on +# small blocks, and two small blocks can share a page, which cudaHostRegister +# refuses to lock twice. +_REGISTER_MIN_BYTES = 64 << 20 + + +def _pinned_empty(*size: int, dtype: torch.dtype, stride=None) -> torch.Tensor: + """A pinned CPU tensor of exactly this size; torch's pinned pool as fallback.""" + shape = size[0] if stride is not None else size + if stride is None: + storage_numel = math.prod(shape) + else: + storage_numel = 1 + sum((d - 1) * st for d, st in zip(shape, stride) if d > 0) + storage_numel = storage_numel if math.prod(shape) else 0 + nbytes = storage_numel * dtype.itemsize + if nbytes >= _REGISTER_MIN_BYTES: + storage = _shared_storage(nbytes) + tensor = torch.empty(0, dtype=dtype) + if storage is not None: + tensor.set_(storage, 0, shape, stride or tensor.stride()) + elif stride is None: + tensor = torch.empty(*size, dtype=dtype) + else: + tensor = torch.empty_strided(size=shape, stride=stride, dtype=dtype) + if _pin_in_place(tensor): + return tensor + del tensor, storage + if stride is None: + return torch.empty(*size, dtype=dtype, pin_memory=True) + return torch.empty_strided(size=shape, stride=stride, dtype=dtype, pin_memory=True) + + class _DirectReader: """Read a mapped tensor's bytes from its checkpoint file with O_DIRECT. @@ -646,7 +734,11 @@ class MappedLayerCourier: [] if direct_copy else [ - torch.empty(slot_bytes, dtype=torch.uint8, pin_memory=pin_slots) + ( + _pinned_empty(slot_bytes, dtype=torch.uint8) + if pin_slots + else torch.empty(slot_bytes, dtype=torch.uint8) + ) for _ in range(self._NUM_SLOTS) ] ) @@ -1304,12 +1396,18 @@ class LayerwiseOffloadManager: # Preserve non-contiguous layouts such as the transposed FP8 # weight views expected by CUTLASS kernels. - cpu_tensor = torch.empty_strided( - size=local_weight.shape, - stride=local_weight.stride(), - dtype=dtype, - pin_memory=pin_this_layer, - ) + if pin_this_layer: + cpu_tensor = _pinned_empty( + local_weight.shape, + dtype=dtype, + stride=local_weight.stride(), + ) + else: + cpu_tensor = torch.empty_strided( + size=local_weight.shape, + stride=local_weight.stride(), + dtype=dtype, + ) if pin_this_layer: yield cpu_tensor.untyped_storage() cpu_tensor.copy_(local_weight) @@ -1341,9 +1439,10 @@ class LayerwiseOffloadManager: total_numel = current_offset # create concatenated CPU buffer (in pinned memory) - cpu_buffer = torch.empty( - total_numel, dtype=dtype, pin_memory=pin_this_layer - ) + if pin_this_layer: + cpu_buffer = _pinned_empty(total_numel, dtype=dtype) + else: + cpu_buffer = torch.empty(total_numel, dtype=dtype) if pin_this_layer: yield cpu_buffer.untyped_storage() diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index ad35ddae1..b3f3006fc 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -2266,3 +2266,40 @@ def test_mixed_scm_and_dbcache_step_schedule(monkeypatch, step_kinds): on_gpu = idx in manager._gpu_layers assert _layer_weight_ok(model.blocks[idx]) is on_gpu, (kind, idx) manager.prepare_for_next_req(non_blocking=False) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_large_pinned_stores_are_registered_in_place_at_exact_size(monkeypatch): + pooled = [] + empty = torch.empty + + def record_pool_use(*args, **kwargs): + if kwargs.get("pin_memory"): + pooled.append(kwargs) + return empty(*args, **kwargs) + + monkeypatch.setattr(torch, "empty", record_pool_use) + nbytes = layerwise_offload_mod._REGISTER_MIN_BYTES + 4096 + tensor = layerwise_offload_mod._pinned_empty(nbytes, dtype=torch.uint8) + # locked where it was allocated: pinned, no pool block, no rounding + assert tensor.is_pinned() + assert tensor.untyped_storage().nbytes() == nbytes + assert pooled == [] + del tensor + gc.collect() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_small_pinned_stores_keep_using_the_pool(monkeypatch): + pooled = [] + empty = torch.empty + + def record_pool_use(*args, **kwargs): + if kwargs.get("pin_memory"): + pooled.append(kwargs) + return empty(*args, **kwargs) + + monkeypatch.setattr(torch, "empty", record_pool_use) + tensor = layerwise_offload_mod._pinned_empty(1024, dtype=torch.float32) + assert tensor.is_pinned() + assert len(pooled) == 1