[diffusion] optimization: stream mapped weights on a shared host/device pool (#38441)

Co-authored-by: Mick Qian <mickqian@radixark.ai>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

python/sglang/multimodal_gen/envs.py

@@ -23,6 +23,10 @@ if TYPE_CHECKING:
     SGLANG_DIFFUSION_TRACE_FUNCTION: int = 0
     SGLANG_DIFFUSION_DISABLE_EARLY_VAE_DECODER_CAST: bool = False
     SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE: bool = False
+    SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED: bool = False
+    SGLANG_DIFFUSION_DISABLE_MAPPED_DIRECT_READ: bool = False
+    SGLANG_DIFFUSION_DEBUG_HOST_MEMORY: bool = False
+    SGLANG_DIFFUSION_DEBUG_LAYERWISE_TIMING: bool = False
     SGLANG_DIFFUSION_DISABLE_LORA_MERGE_CACHE: bool = False
     SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD: str = "fork"
     SGLANG_DIFFUSION_TARGET_DEVICE: str = "cuda"
@@ -34,6 +38,10 @@ if TYPE_CHECKING:
     VERBOSE: bool = False
     SGLANG_DIFFUSION_SERVER_DEV_MODE: bool = False
     SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER: bool = False
+    SGLANG_DIFFUSION_HOST_SPILL_DIR: str = os.path.expanduser(
+        "~/.cache/sglang/diffusion/host_spill"
+    )
+    SGLANG_DIFFUSION_DISABLE_HOST_SPILL: bool = False
     SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB: float | None = None
     SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB: float | None = None
     SGLANG_DIFFUSION_STAGE_LOGGING: bool = False
@@ -248,6 +256,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER": _lazy_bool(
         "SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER"
     ),
+    # Where transformed weight copies (fused q/k/v, reordered rows) live as
+    # file mappings when host copies must stay reclaimable; reused across
+    # starts of the same checkpoint.
+    "SGLANG_DIFFUSION_HOST_SPILL_DIR": _lazy_str(
+        "SGLANG_DIFFUSION_HOST_SPILL_DIR",
+        os.path.expanduser("~/.cache/sglang/diffusion/host_spill"),
+    ),
+    "SGLANG_DIFFUSION_DISABLE_HOST_SPILL": _lazy_bool(
+        "SGLANG_DIFFUSION_DISABLE_HOST_SPILL"
+    ),
     # Test hook: make the host memory budget behave as if the machine had this
     # many GiB of RAM (available = this figure minus the process's own
     # anonymous memory). CI uses it to exercise the constrained placement
@@ -304,6 +322,28 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE": _lazy_bool(
         "SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE"
     ),
+    # Kill-switch: do not madvise(MADV_WILLNEED) mapped layers ahead of the
+    # courier; their pages arrive at fault-time readahead beats instead.
+    "SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED": _lazy_bool(
+        "SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED"
+    ),
+    # Kill-switch: on a shared host/device pool the courier reads mapped layers
+    # from their checkpoint files with O_DIRECT instead of through the page
+    # cache. This forces the mmap path.
+    "SGLANG_DIFFUSION_DISABLE_MAPPED_DIRECT_READ": _lazy_bool(
+        "SGLANG_DIFFUSION_DISABLE_MAPPED_DIRECT_READ"
+    ),
+    # Debug: after auto residency settles, log where this process's host memory
+    # sits -- per component and per kind (anonymous, mapped, pinned) -- next to
+    # the kernel's view of the process.
+    "SGLANG_DIFFUSION_DEBUG_HOST_MEMORY": _lazy_bool(
+        "SGLANG_DIFFUSION_DEBUG_HOST_MEMORY"
+    ),
+    # Debug: at the end of every layerwise stage, log where the courier and the
+    # compute thread spent their time (populate, memcpy, H2D, waits).
+    "SGLANG_DIFFUSION_DEBUG_LAYERWISE_TIMING": _lazy_bool(
+        "SGLANG_DIFFUSION_DEBUG_LAYERWISE_TIMING"
+    ),
     # Kill-switch: keep LoRA-merged weights in anonymous host memory instead
     # of the file-backed LoRA merge cache.
     "SGLANG_DIFFUSION_DISABLE_LORA_MERGE_CACHE": _lazy_bool(

python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py

@@ -4,7 +4,6 @@ import os
 
 import torch
 import torch.nn as nn
-from safetensors.torch import load_file as safetensors_load_file
 from safetensors.torch import safe_open
 from safetensors.torch import save_file as safetensors_save_file
 
@@ -170,6 +169,22 @@ def _decode_dtype_store_path(
     )
 
 
+def _load_safetensors_file(path: str) -> dict:
+    """The VAE checkpoint itself: read-only where host copies are redundant."""
+    from sglang.multimodal_gen.runtime.loader.utils import (
+        _load_safetensors_file as _load,
+    )
+
+    return _load(path)
+
+
+def _load_store(path: str) -> dict:
+    """Map the store read-only where host copies are redundant (see loader.utils)."""
+    from sglang.multimodal_gen.runtime.loader.utils import _load_safetensors_file
+
+    return _load_safetensors_file(path)
+
+
 def _assign_matching_store(vae, mapped: dict, dtype: torch.dtype) -> bool:
     """Adopt a decode-dtype store if it matches the module, else refuse."""
     state = vae.state_dict()
@@ -199,7 +214,7 @@ def _rehome_cast_weights_to_file(
     path = _decode_dtype_store_path(component_model_path, component_name, dtype)
     try:
         if os.path.exists(path):
-            mapped = safetensors_load_file(path)
+            mapped = _load_store(path)
             if mapped and _assign_matching_store(vae, mapped, dtype):
                 return len(mapped), True
             raise ValueError("existing decode-dtype store does not match the module")
@@ -215,7 +230,7 @@ def _rehome_cast_weights_to_file(
         tmp = f"{path}.tmp.{os.getpid()}"
         safetensors_save_file({k: v.contiguous() for k, v in cast_state.items()}, tmp)
         os.replace(tmp, path)
-        mapped = safetensors_load_file(path)
+        mapped = _load_store(path)
         if set(mapped) != set(cast_state):
             raise ValueError("decode-dtype store does not match the cast weights")
         vae.load_state_dict(mapped, strict=False, assign=True)
@@ -671,7 +686,7 @@ class VAELoader(WeightOverrideComponentLoader):
 
         loaded = {}
         for sf_path in safetensors_list:
-            loaded.update(safetensors_load_file(sf_path))
+            loaded.update(_load_safetensors_file(sf_path))
         _backfill_ltx2_audio_vae_latent_stats(loaded, component_type)
         num_deparameterized = adopt_plain_weight_norm_state(vae, loaded)
         target_state = vae.state_dict()

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -38,6 +38,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
     split_bitsandbytes_4bit_state,
 )
 from sglang.multimodal_gen.runtime.loader import rank_local_checkpoint
+from sglang.multimodal_gen.runtime.loader.host_spill import HostSpill
 from sglang.multimodal_gen.runtime.loader.utils import (
     finalize_loaded_model,
     get_param_names_mapping,
@@ -48,6 +49,9 @@ from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
 from sglang.multimodal_gen.runtime.loader.weight_utils import (
     safetensors_weights_iterator,
 )
+from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import (
+    host_copies_are_redundant,
+)
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.quantization_utils import (
@@ -315,6 +319,18 @@ def maybe_load_fsdp_model(
         # layerwise offload replaces block parameters with placeholders after
         # load, so compatible checkpoint tensors stay file-backed on CPU
         model._keep_checkpoint_mapping = True
+    host_spill = None
+    if (
+        weight_dir_list
+        and weights_iterator is None
+        and weight_load_plan.checkpoint_load_device.type == "cpu"
+        and host_copies_are_redundant()
+    ):
+        # Compatible tensors stay on the checkpoint mapping; the fused and
+        # reordered ones are materialized. On a shared pool their anonymous
+        # copies are never reclaimable (10.4 GiB for the H3 DiT's fused
+        # q/k/v), so they are materialized into file mappings instead.
+        host_spill = HostSpill.for_checkpoint(weight_dir_list)
     defer_cpu_placement = bool(
         component_starts_on_cpu
         and weight_load_plan.defer_cpu_placement
@@ -437,7 +453,10 @@ def maybe_load_fsdp_model(
             weight_load_plan.load_full_state_dict_on_device
         ),
         preconverted_state_dict=preconverted_state_dict,
+        host_spill=host_spill,
     )
+    if host_spill is not None:
+        host_spill.log_summary(type(model).__name__)
     if bnb_quant_states:
         attach_bitsandbytes_4bit_quant_states(
             dict(model.named_parameters()), bnb_quant_states
@@ -557,6 +576,7 @@ def load_model_from_full_model_state_dict(
         | None
     ) = None,
     allow_device_tensor_assignment: bool = False,
+    host_spill: HostSpill | None = None,
 ) -> _IncompatibleKeys:
     """
     Converting full state dict into a sharded state dict
@@ -589,7 +609,12 @@ def load_model_from_full_model_state_dict(
             full_sd_iterator,
             param_names_mapping,
             valid_target_names=set(meta_sd.keys()),
+            fused_tensor_factory=(None if host_spill is None else host_spill.tensor),
         )  # type: ignore
+        if host_spill is not None:
+            for name, tensor in custom_param_sd.items():
+                if isinstance(tensor, torch.Tensor):
+                    host_spill.seal(name, tensor.shape, tensor.dtype)
     else:
         custom_param_sd, reverse_param_names_mapping = preconverted_state_dict
 
@@ -734,11 +759,21 @@ def load_model_from_full_model_state_dict(
                 ):
                     sharded_tensor = full_tensor
                 else:
-                    sharded_tensor = torch.empty_like(
-                        meta_sharded_param,
-                        device=checkpoint_load_device,
-                        dtype=target_dtype,
+                    spilled = (
+                        None
+                        if host_spill is None or checkpoint_load_device.type != "cpu"
+                        else host_spill.tensor(
+                            target_param_name, meta_sharded_param.shape, target_dtype
+                        )
                     )
+                    if spilled is not None:
+                        sharded_tensor = spilled[0]
+                    else:
+                        sharded_tensor = torch.empty_like(
+                            meta_sharded_param,
+                            device=checkpoint_load_device,
+                            dtype=target_dtype,
+                        )
                     # Preserve requires_grad flag to avoid errors with non-floating dtypes
                     requires_grad = meta_sharded_param.requires_grad
                     temp_param = _make_param_like(actual_param, sharded_tensor)
@@ -759,6 +794,10 @@ def load_model_from_full_model_state_dict(
                             f"param_cls={type(actual_param).__name__}"
                         ) from exc
                     sharded_tensor = temp_param.data
+                    if host_spill is not None and spilled is not None:
+                        host_spill.seal(
+                            target_param_name, meta_sharded_param.shape, target_dtype
+                        )
             else:
                 # In cases where parts of the model aren't sharded, some parameters will be plain tensors
                 sharded_tensor = full_tensor

python/sglang/multimodal_gen/runtime/loader/host_spill.py

@@ -0,0 +1,158 @@
+"""File-backed homes for the host copies a checkpoint mapping cannot provide.
+
+A weight that is fused (q/k/v into one projection), sharded or otherwise
+transformed at load has no checkpoint bytes to stay mapped on, so the loader
+materializes it. Anonymous memory is the wrong home for that copy on a host
+that keeps everything else mapped: it is never reclaimable, and on a shared
+CPU/GPU pool it is memory the page cache and the device both lose. A shared
+file mapping under the cache directory holds the same bytes as page cache
+instead -- reclaimable, readable with O_DIRECT, and, once written, reusable
+by the next start of the same checkpoint.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import shutil
+from pathlib import Path
+from typing import Callable, Iterable
+
+import torch
+
+from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+# Below this a copy is not worth a file: the bookkeeping costs more than the
+# bytes it would return to the pool.
+MIN_SPILL_BYTES = 64 << 20
+# Free space to leave on the spill filesystem after a write.
+SPILL_DISK_RESERVE_BYTES = 2 << 30
+
+FusedTensorFactory = Callable[
+    [str, torch.Size, torch.dtype], "tuple[torch.Tensor, bool] | None"
+]
+
+
+def checkpoint_fingerprint(weight_dirs: Iterable[str]) -> str:
+    """Identity of a checkpoint on disk: each shard's path, size and mtime."""
+    digest = hashlib.sha1()
+    for weight_dir in sorted(str(d) for d in weight_dirs):
+        root = Path(weight_dir)
+        files = [root] if root.is_file() else sorted(root.glob("*.safetensors"))
+        for path in files:
+            try:
+                stat = path.stat()
+            except OSError:
+                continue
+            digest.update(
+                f"{path.resolve()}|{stat.st_size}|{stat.st_mtime_ns}\n".encode()
+            )
+    return digest.hexdigest()[:20]
+
+
+class HostSpill:
+    """Hands out file-backed tensors keyed by (parameter, dtype, shape).
+
+    A tensor comes back ``(tensor, filled)``: ``filled`` says an earlier run
+    wrote and sealed the same key, so the caller can skip producing it. A
+    caller that produces the bytes must ``seal`` the key afterwards; an
+    unsealed file is treated as garbage and rewritten.
+    """
+
+    def __init__(self, directory: str | os.PathLike[str], fingerprint: str):
+        self.directory = Path(directory) / fingerprint
+        self._disabled_reason: str | None = None
+        self.bytes_written = 0
+        self.bytes_reused = 0
+        self.count_written = 0
+        self.count_reused = 0
+        self._open: dict[str, str] = {}
+
+    @classmethod
+    def for_checkpoint(cls, weight_dirs: Iterable[str]) -> HostSpill | None:
+        if envs.SGLANG_DIFFUSION_DISABLE_HOST_SPILL:
+            return None
+        directory = os.path.expanduser(envs.SGLANG_DIFFUSION_HOST_SPILL_DIR)
+        return cls(directory, checkpoint_fingerprint(weight_dirs))
+
+    def _path(self, key: str) -> Path:
+        return self.directory / (hashlib.sha1(key.encode()).hexdigest() + ".bin")
+
+    def _disable(self, reason: str) -> None:
+        if self._disabled_reason is None:
+            self._disabled_reason = reason
+            logger.warning(
+                "Host spill disabled for this load: %s; transformed weights "
+                "fall back to anonymous memory.",
+                reason,
+            )
+
+    def tensor(
+        self, name: str, shape: torch.Size, dtype: torch.dtype
+    ) -> tuple[torch.Tensor, bool] | None:
+        """A file-backed tensor for ``name``, or None to use anonymous memory."""
+        if self._disabled_reason is not None:
+            return None
+        numel = 1
+        for dim in shape:
+            numel *= int(dim)
+        nbytes = numel * torch.empty((), dtype=dtype).element_size()
+        if nbytes < MIN_SPILL_BYTES:
+            return None
+        key = f"{name}|{dtype}|{tuple(int(d) for d in shape)}"
+        path = self._path(key)
+        sealed = path.with_suffix(".ok")
+        try:
+            self.directory.mkdir(parents=True, exist_ok=True)
+            filled = sealed.exists() and path.exists() and path.stat().st_size == nbytes
+            if not filled:
+                sealed.unlink(missing_ok=True)
+                free = shutil.disk_usage(self.directory).free
+                if free < nbytes + SPILL_DISK_RESERVE_BYTES:
+                    self._disable(
+                        f"{free / 2**30:.1f} GiB free under {self.directory}, "
+                        f"{nbytes / 2**30:.1f} GiB needed"
+                    )
+                    return None
+            storage = torch.from_file(str(path), shared=True, size=numel, dtype=dtype)
+        except (OSError, RuntimeError) as exc:
+            self._disable(f"{type(exc).__name__}: {exc}")
+            return None
+        tensor = storage.view(tuple(int(d) for d in shape))
+        if filled:
+            self.bytes_reused += nbytes
+            self.count_reused += 1
+        else:
+            self._open[key] = str(sealed)
+            self.bytes_written += nbytes
+            self.count_written += 1
+        return tensor, filled
+
+    def seal(self, name: str, shape: torch.Size, dtype: torch.dtype) -> None:
+        """Mark a key as completely written so the next start can reuse it."""
+        key = f"{name}|{dtype}|{tuple(int(d) for d in shape)}"
+        sealed = self._open.pop(key, None)
+        if sealed is None:
+            return
+        try:
+            with open(sealed, "w") as handle:
+                handle.write("ok\n")
+        except OSError as exc:
+            logger.debug("could not seal %s: %s", sealed, exc)
+
+    def log_summary(self, component: str) -> None:
+        if self.count_written == 0 and self.count_reused == 0:
+            return
+        logger.info(
+            "%s: %d transformed weights (%.2f GiB) live in file mappings under %s "
+            "(%d written, %d reused from an earlier start).",
+            component,
+            self.count_written + self.count_reused,
+            (self.bytes_written + self.bytes_reused) / 2**30,
+            self.directory,
+            self.count_written,
+            self.count_reused,
+        )

python/sglang/multimodal_gen/runtime/loader/readonly_safetensors.py

@@ -0,0 +1,95 @@
+"""Read-only mappings of safetensors files.
+
+safetensors maps a file for torch through ``UntypedStorage.from_file(shared=False)``:
+a private, writable mapping. On a shared CPU/GPU pool that permission costs
+memory: when the device copies from such a mapping the driver pins the pages
+with write intent, the kernel breaks copy-on-write, and every page copied
+becomes anonymous memory -- 1 GiB copied in 1 GiB of unreclaimable RAM, at a
+fraction of the bandwidth (0.1-3.8 GiB/s against 27 GiB/s on a GB10). A
+read-only mapping copies at full speed and leaves the pages as page cache.
+
+Frozen weights never need the write permission; a write into one here faults
+instead of silently copying the page, which is the invariant we want.
+"""
+
+from __future__ import annotations
+
+import json
+import mmap
+import os
+import struct
+import warnings
+from typing import Iterator
+
+import torch
+
+_DTYPES = {
+    "F64": torch.float64,
+    "F32": torch.float32,
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+    "I64": torch.int64,
+    "I32": torch.int32,
+    "I16": torch.int16,
+    "I8": torch.int8,
+    "U8": torch.uint8,
+    "BOOL": torch.bool,
+    "F8_E4M3": torch.float8_e4m3fn,
+    "F8_E5M2": torch.float8_e5m2,
+}
+
+# Mappings stay for the life of the process: the tensors handed out are views
+# into them, and the layerwise manager keeps such views as its host store.
+_MAPPINGS: dict[str, mmap.mmap] = {}
+
+
+def _mapping(path: str) -> mmap.mmap:
+    real = os.path.realpath(path)
+    mapped = _MAPPINGS.get(real)
+    if mapped is None:
+        fd = os.open(real, os.O_RDONLY)
+        try:
+            size = os.fstat(fd).st_size
+            mapped = mmap.mmap(fd, size, prot=mmap.PROT_READ, flags=mmap.MAP_PRIVATE)
+        finally:
+            os.close(fd)
+        _MAPPINGS[real] = mapped
+    return mapped
+
+
+def _header(mapped: mmap.mmap) -> tuple[dict, int]:
+    (n,) = struct.unpack("<Q", mapped[:8])
+    return json.loads(mapped[8 : 8 + n]), 8 + n
+
+
+def _tensor(mapped: mmap.mmap, base: int, meta: dict) -> torch.Tensor:
+    dtype = _DTYPES[meta["dtype"]]
+    start, end = meta["data_offsets"]
+    shape = tuple(meta["shape"])
+    if end == start:
+        return torch.empty(shape, dtype=dtype)
+    count = (end - start) // torch.empty((), dtype=dtype).element_size()
+    with warnings.catch_warnings():
+        # torch warns that the buffer is not writable; that is the point.
+        warnings.simplefilter("ignore")
+        flat = torch.frombuffer(mapped, dtype=dtype, count=count, offset=base + start)
+    return flat.view(shape)
+
+
+def safetensors_keys(path: str) -> list[str]:
+    header, _ = _header(_mapping(path))
+    return [name for name in header if name != "__metadata__"]
+
+
+def iter_safetensors_readonly(path: str) -> Iterator[tuple[str, torch.Tensor]]:
+    """(name, tensor) for every tensor in the file, as views of a read-only mapping."""
+    mapped = _mapping(path)
+    header, base = _header(mapped)
+    for name, meta in header.items():
+        if name == "__metadata__":
+            continue
+        yield name, _tensor(mapped, base, meta)
+
+
+def load_safetensors_readonly(path: str) -> dict[str, torch.Tensor]:
+    return dict(iter_safetensors_readonly(path))

python/sglang/multimodal_gen/runtime/loader/utils.py

@@ -181,10 +181,45 @@ def get_param_names_mapping(
     return mapping_fn
 
 
+def _fuse_tensors(
+    target_param_name: str,
+    tensors: list[torch.Tensor],
+    fused_tensor_factory: (
+        Callable[[str, torch.Size, torch.dtype], tuple[torch.Tensor, bool] | None]
+        | None
+    ),
+) -> torch.Tensor:
+    """Concatenate the pieces of one parameter along dim 0.
+
+    The factory, when given, provides the destination (a file mapping that
+    outlives anonymous memory) and says whether an earlier run already filled
+    it -- then the pieces are not even read.
+    """
+    if fused_tensor_factory is None or any(t.device.type != "cpu" for t in tensors):
+        return torch.cat(tensors, dim=0)
+    if (
+        len({tuple(t.shape[1:]) for t in tensors}) != 1
+        or len({t.dtype for t in tensors}) != 1
+    ):
+        return torch.cat(tensors, dim=0)
+    shape = torch.Size([sum(t.shape[0] for t in tensors), *tensors[0].shape[1:]])
+    provided = fused_tensor_factory(target_param_name, shape, tensors[0].dtype)
+    if provided is None:
+        return torch.cat(tensors, dim=0)
+    out, filled = provided
+    if not filled:
+        torch.cat(tensors, dim=0, out=out)
+    return out
+
+
 def hf_to_custom_state_dict(
     hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]],
     param_names_mapping: Callable[[str], tuple[str, Any, Any]],
     valid_target_names: set[str] | None = None,
+    fused_tensor_factory: (
+        Callable[[str, torch.Size, torch.dtype], tuple[torch.Tensor, bool] | None]
+        | None
+    ) = None,
     *,
     strict: bool = False,
 ) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]:
@@ -236,7 +271,9 @@ def hf_to_custom_state_dict(
                     to_merge_params[target_param_name][i]
                     for i in range(num_params_to_merge)
                 ]
-                full_tensor = torch.cat(sorted_tensors, dim=0)
+                full_tensor = _fuse_tensors(
+                    target_param_name, sorted_tensors, fused_tensor_factory
+                )
                 del to_merge_params[target_param_name]
             else:
                 continue
@@ -365,10 +402,18 @@ def keep_checkpoint_mapped(*, weight_bytes: int, component: str) -> bool:
     choice -- its pages are resident, where a mapping's first use pays a fault.
     """
     from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import (
+        host_copies_are_redundant,
         host_copies_would_not_fit,
         host_memory_available_bytes,
     )
 
+    if host_copies_are_redundant():
+        logger.info(
+            "%s stays on its checkpoint mapping: host and device share one "
+            "memory pool, so a copy would hold the same bytes twice.",
+            component,
+        )
+        return True
     if not host_copies_would_not_fit(weight_bytes):
         return False
     logger.info(
@@ -449,6 +494,27 @@ def _list_safetensors_files(
     return filter_duplicate_precision_variant_safetensors(found)
 
 
+def _load_safetensors_file(path: str) -> dict[str, torch.Tensor]:
+    """One safetensors file; a read-only mapping where host copies are redundant.
+
+    safetensors maps for torch through a private *writable* mapping, and on a
+    shared CPU/GPU pool a device copy from such a mapping copies every page it
+    touches into anonymous memory (the driver pins with write intent). A
+    read-only mapping copies at full speed and stays page cache.
+    """
+    from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import (
+        host_copies_are_redundant,
+    )
+
+    if host_copies_are_redundant():
+        from sglang.multimodal_gen.runtime.loader.readonly_safetensors import (
+            load_safetensors_readonly,
+        )
+
+        return load_safetensors_readonly(path)
+    return safetensors_load_file(path)
+
+
 def load_safetensors_state_dict(model_path: str) -> dict[str, torch.Tensor]:
     """Load one safetensors checkpoint, including an indexed sharded set."""
     index_path = _select_safetensors_index_file(model_path, _DEFAULT_SAFETENSORS_INDEX)
@@ -456,7 +522,7 @@ def load_safetensors_state_dict(model_path: str) -> dict[str, torch.Tensor]:
     if index_path is not None:
         state_dict: dict[str, torch.Tensor] = {}
         for path in safetensors_files:
-            state_dict.update(safetensors_load_file(path))
+            state_dict.update(_load_safetensors_file(path))
         return state_dict
 
     if not safetensors_files:
@@ -466,7 +532,7 @@ def load_safetensors_state_dict(model_path: str) -> dict[str, torch.Tensor]:
             f"Found {len(safetensors_files)} safetensors files in {model_path} "
             "and no index to disambiguate them."
         )
-    return safetensors_load_file(safetensors_files[0])
+    return _load_safetensors_file(safetensors_files[0])
 
 
 BYTES_PER_GB = 1024**3

python/sglang/multimodal_gen/runtime/loader/weight_readers/safetensors_mmap.py

@@ -4,6 +4,11 @@
 `safe_open` maps the file, so a CPU tensor it yields is a view into the
 checkpoint rather than a copy. Those pages are file-backed, which is what lets
 the kernel drop them under memory pressure even on a host with no swap.
+
+Where host copies are redundant (the device shares the host pool) the mapping
+is made read-only instead: safetensors maps writable, and a device copy from a
+writable private mapping there turns every page it touches into anonymous
+memory at a fraction of the bandwidth (see readonly_safetensors).
 """
 
 from typing import Callable, ClassVar, Iterator
@@ -12,6 +17,13 @@ import torch
 from safetensors.torch import safe_open
 from tqdm.auto import tqdm
 
+from sglang.multimodal_gen.runtime.loader.readonly_safetensors import (
+    iter_safetensors_readonly,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import (
+    host_copies_are_redundant,
+)
+
 _BAR_FORMAT = "{desc}: {percentage:.0f}%|{bar}| {n_fmt}/{total_fmt}"
 
 
@@ -40,6 +52,12 @@ class SafetensorsMmapReader:
             disable=not show_progress,
             bar_format=_BAR_FORMAT,
         ):
+            if device == "cpu" and host_copies_are_redundant():
+                for name, tensor in iter_safetensors_readonly(path):
+                    if key_filter is not None and not key_filter(name):
+                        continue
+                    yield name, tensor
+                continue
             with safe_open(path, framework="pt", device=device) as handle:
                 for name in handle.keys():  # noqa: SIM118
                     if key_filter is not None and not key_filter(name):

python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py

@@ -6,6 +6,7 @@ from typing import Mapping, MutableMapping, Protocol, Sequence
 import torch
 import torch.nn as nn
 
+from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
     COMPONENT_OFFLOAD,
     LAYERWISE_OFFLOAD,
@@ -650,6 +651,23 @@ class ComponentResidencyManager:
             )
             self._completed_warmup_phase_peaks = dict(self._warmup_phase_peaks)
         self._track_warmup_memory = False
+        if (
+            current_platform.device_shares_host_memory()
+            and torch.get_device_module().is_available()
+        ):
+            # One pool: every byte the caching allocator keeps reserved between
+            # requests is page cache the next request's streamed encoder cannot use.
+            torch.get_device_module().empty_cache()
+            if envs.SGLANG_DIFFUSION_DEBUG_HOST_MEMORY:
+                from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_breakdown import (
+                    log_host_memory_breakdown,
+                )
+
+                self._debug_requests_seen = getattr(self, "_debug_requests_seen", 0) + 1
+                log_host_memory_breakdown(
+                    self.placement_modules(),
+                    label=f"after request {self._debug_requests_seen}",
+                )
 
     def _begin_warmup_phase(
         self,

python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py

@@ -9,10 +9,20 @@ import torch.nn as nn
 from torch.distributed.fsdp import FSDPModule
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import (
+    shared_pool_available_bytes,
+)
 from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
     LayerwiseOffloadableModuleMixin,
 )
 from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+# Device growth between a component's stages on a shared pool: activations and
+# the allocator's reserve (9.4 GiB measured for H3 at 1344x768x124f) plus margin.
+SHARED_POOL_NEXT_STAGE_HEADROOM_BYTES = 12 * 1024**3
+
+logger = init_logger(__name__)
 
 if TYPE_CHECKING:
     from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
@@ -174,7 +184,15 @@ class ComponentOffloadStrategy(ComponentResidencyStrategy):
         self.wait_for_use(module, use, state)
         tensor = _module_reference_tensor(module)
         if tensor is not None and tensor.device.type != "cpu":
-            module.to("cpu", non_blocking=True)
+            # A non-blocking device->host move lands in pinned host memory the
+            # size of the component. On a shared pool that pins a second copy
+            # of the weights next to the device copy still being read from
+            # -- a 57 GiB DiT took 43 GiB of shared memory in under a minute
+            # and exhausted a GB10. Take the synchronous, pageable path there.
+            module.to(
+                "cpu",
+                non_blocking=not current_platform.device_shares_host_memory(),
+            )
         self._ready_events.pop(use.component_name, None)
 
     def finish_request(
@@ -231,6 +249,35 @@ class LayerwiseOffloadStrategy(ComponentResidencyStrategy):
             torch.mps.synchronize()
             module.restore_mps_cpu_non_layer_weights()
             torch.mps.empty_cache()
+        elif (
+            current_platform.is_cuda() and current_platform.device_shares_host_memory()
+        ):
+            # The stage's streamed layer windows are freed but still reserved
+            # by the caching allocator. On a shared pool that reserve is host
+            # memory the next stage's mapping needs as page cache; hand it back.
+            empty_cache = getattr(torch.get_device_module(), "empty_cache", None)
+            if empty_cache is not None:
+                empty_cache()
+            # And this component's own pages are now the least valuable in the
+            # cache until its next stage; say so before the next phase evicts.
+            # The room the cache will have for this component's next stream is
+            # what is available now less what the stages in between need.
+            room_bytes = max(
+                0, shared_pool_available_bytes() - SHARED_POOL_NEXT_STAGE_HEADROOM_BYTES
+            )
+            paged_out = 0
+            for manager in module.layerwise_offload_managers:
+                advise_cold = getattr(manager, "advise_mapped_pages_cold", None)
+                if advise_cold is not None:
+                    paged_out += int(advise_cold(room_bytes=room_bytes) or 0)
+            if paged_out:
+                logger.info(
+                    "Layerwise offload: paged out the first %.1f GiB of %s so the "
+                    "next request's stream fits the %.1f GiB the cache can give it.",
+                    paged_out / 1024**3,
+                    use.component_name,
+                    room_bytes / 1024**3,
+                )
 
     def finish_request(
         self,

python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_breakdown.py

@@ -0,0 +1,348 @@
+"""Where a worker's host memory sits, by component and by kind.
+
+Debug aid behind ``SGLANG_DIFFUSION_DEBUG_HOST_MEMORY``: on a shared
+host/device pool every anonymous byte the runtime keeps is a byte the page
+cache cannot hold, so the breakdown says what to cut.
+"""
+
+from __future__ import annotations
+
+import gc
+import logging
+from bisect import bisect_right
+from collections.abc import Mapping
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+GIB = 1024**3
+
+
+def _file_backed_ranges() -> list[tuple[int, int]]:
+    ranges: list[tuple[int, int]] = []
+    try:
+        with open("/proc/self/maps") as handle:
+            for line in handle:
+                fields = line.split()
+                if len(fields) < 6 or fields[5].startswith("["):
+                    continue
+                start, end = fields[0].split("-")
+                ranges.append((int(start, 16), int(end, 16)))
+    except OSError:
+        return []
+    ranges.sort()
+    return ranges
+
+
+def _kind(
+    tensor: torch.Tensor, starts: list[int], ranges: list[tuple[int, int]]
+) -> str:
+    if tensor.is_pinned():
+        return "pinned"
+    ptr = tensor.data_ptr()
+    index = bisect_right(starts, ptr) - 1
+    if index >= 0 and ranges[index][0] <= ptr < ranges[index][1]:
+        return "mapped"
+    return "anonymous"
+
+
+def _anon_vmas(min_bytes: int = 128 * 1024**2) -> list[tuple[int, int, int, str]]:
+    """(start, end, anonymous_bytes, vmflags) of anonymous mappings holding at least min_bytes."""
+    out: list[tuple[int, int, int, str]] = []
+    try:
+        start = end = 0
+        path = ""
+        anon = 0
+        flags = ""
+        with open("/proc/self/smaps") as handle:
+            for line in handle:
+                if line[0] in "0123456789abcdef" and "-" in line.split()[0]:
+                    if path in ("", "[anon]") and anon >= min_bytes:
+                        out.append((start, end, anon, flags))
+                    fields = line.split()
+                    start, end = (int(x, 16) for x in fields[0].split("-"))
+                    path = fields[5] if len(fields) >= 6 else ""
+                    anon = 0
+                    flags = ""
+                elif line.startswith("Anonymous:"):
+                    anon = int(line.split()[1]) * 1024
+                elif line.startswith("VmFlags:"):
+                    flags = line.split(":", 1)[1].strip()
+        if path in ("", "[anon]") and anon >= min_bytes:
+            out.append((start, end, anon, flags))
+    except OSError:
+        pass
+    return sorted(out, key=lambda item: -item[2])
+
+
+def _file_mapping_cow(top: int = 3) -> tuple[int, list[tuple[str, int]]]:
+    """Anonymous bytes inside file mappings: pages a write copied out of the file.
+
+    A private writable mapping (safetensors' from_file) stays "file-backed" by
+    pointer, but every page written to it becomes anonymous memory that no
+    longer drops under pressure. Returns (total, [(path, bytes)] for the largest).
+    """
+    per_path: dict[str, int] = {}
+    try:
+        path = ""
+        with open("/proc/self/smaps") as handle:
+            for line in handle:
+                if line[0] in "0123456789abcdef" and "-" in line.split()[0]:
+                    fields = line.split()
+                    path = fields[5] if len(fields) >= 6 else ""
+                    if path.startswith("[") or path.startswith("/dev/"):
+                        path = ""
+                elif path and line.startswith("Anonymous:"):
+                    anon = int(line.split()[1]) * 1024
+                    if anon:
+                        per_path[path] = per_path.get(path, 0) + anon
+    except OSError:
+        return 0, []
+    ranked = sorted(per_path.items(), key=lambda item: -item[1])
+    return sum(per_path.values()), ranked[:top]
+
+
+def _mallinfo() -> dict[str, float]:
+    try:
+        import ctypes
+        import ctypes.util
+
+        libc = ctypes.CDLL(ctypes.util.find_library("c"))
+
+        class MallInfo2(ctypes.Structure):
+            _fields_ = [
+                (name, ctypes.c_size_t)
+                for name in (
+                    "arena",
+                    "ordblks",
+                    "smblks",
+                    "hblks",
+                    "hblkhd",
+                    "usmblks",
+                    "fsmblks",
+                    "uordblks",
+                    "fordblks",
+                    "keepcost",
+                )
+            ]
+
+        libc.mallinfo2.restype = MallInfo2
+        info = libc.mallinfo2()
+        return {
+            "glibc_arena": info.arena / GIB,
+            "glibc_mmapped": info.hblkhd / GIB,
+            "glibc_in_use": info.uordblks / GIB,
+            "glibc_free": info.fordblks / GIB,
+        }
+    except Exception:
+        return {}
+
+
+def _sample_anon_vma(start: int, end: int) -> str:
+    """A few bytes from the start and the middle of a mapping, as hex, plus a
+    guess at what they hold (zeros, bf16-looking, other)."""
+    import ctypes
+
+    out = []
+    for offset in (0, (end - start) // 2 & ~4095):
+        try:
+            raw = ctypes.string_at(start + offset, 32)
+        except Exception:
+            out.append("unreadable")
+            continue
+        if not any(raw):
+            kind = "zeros"
+        else:
+            # bf16 weights: high bytes cluster around 0x3c-0x40 / 0xbc-0xc0
+            highs = raw[1::2]
+            kind = (
+                "bf16-like"
+                if sum(1 for b in highs if 0x38 <= (b & 0x7F) <= 0x42) >= 10
+                else "other"
+            )
+        out.append(f"@{offset:#x}:{raw[:16].hex()}({kind})")
+    return " ".join(out)
+
+
+def _tensor_ptrs_inside(start: int, end: int) -> tuple[int, float]:
+    count = 0
+    total = 0.0
+    for obj in gc.get_objects():
+        if isinstance(obj, torch.Tensor) and obj.device.type == "cpu":
+            try:
+                ptr = obj.data_ptr()
+                nbytes = obj.numel() * obj.element_size()
+            except Exception:
+                continue
+            if start <= ptr < end:
+                count += 1
+                total += nbytes / GIB
+    return count, total
+
+
+def _smaps_rollup() -> dict[str, float]:
+    totals: dict[str, float] = {}
+    try:
+        with open("/proc/self/smaps_rollup") as handle:
+            for line in handle:
+                key, _, rest = line.partition(":")
+                if key in (
+                    "Rss",
+                    "Anonymous",
+                    "Rss_File",
+                    "Rss_Shmem",
+                    "Private_Dirty",
+                ):
+                    totals[key] = int(rest.split()[0]) * 1024 / GIB
+    except OSError:
+        pass
+    return totals
+
+
+def host_memory_breakdown(modules: Mapping[str, object]) -> dict[str, dict[str, float]]:
+    """GiB of CPU tensor storage per component and kind; ``other`` covers
+    tensors no module owns (staging buffers, caches, activations kept alive)."""
+    ranges = _file_backed_ranges()
+    starts = [start for start, _ in ranges]
+    owners: dict[int, str] = {}
+    for name, module in modules.items():
+        if not isinstance(module, torch.nn.Module):
+            continue
+        for tensor in list(module.parameters()) + list(module.buffers()):
+            if tensor.device.type == "cpu":
+                owners[tensor.untyped_storage().data_ptr()] = name
+    seen: set[int] = set()
+    table: dict[str, dict[str, float]] = {}
+    for obj in gc.get_objects():
+        if not isinstance(obj, torch.Tensor) or obj.device.type != "cpu":
+            continue
+        try:
+            storage = obj.untyped_storage()
+            key = storage.data_ptr()
+            nbytes = storage.nbytes()
+        except Exception:
+            continue
+        if key in seen or nbytes == 0:
+            continue
+        seen.add(key)
+        owner = owners.get(key, "other")
+        kind = _kind(obj, starts, ranges)
+        table.setdefault(owner, {})
+        table[owner][kind] = table[owner].get(kind, 0.0) + nbytes / GIB
+    return table
+
+
+def log_host_memory_breakdown(modules: Mapping[str, object], *, label: str) -> None:
+    table = host_memory_breakdown(modules)
+    rollup = _smaps_rollup()
+    lines = [f"Host memory breakdown ({label}):"]
+    if rollup:
+        lines.append(
+            "  process: "
+            + " ".join(f"{key}={value:.2f}GiB" for key, value in sorted(rollup.items()))
+        )
+    for owner in sorted(table, key=lambda name: -sum(table[name].values())):
+        kinds = " ".join(
+            f"{kind}={value:.2f}GiB" for kind, value in sorted(table[owner].items())
+        )
+        lines.append(f"  {owner}: {kinds}")
+    cow_total, cow_top = _file_mapping_cow()
+    if cow_total:
+        lines.append(
+            f"  copy-on-write pages in file mappings: {cow_total / GIB:.2f}GiB; "
+            + ", ".join(
+                f"{'/'.join(path.rsplit('/', 3)[-3:])}={size / GIB:.2f}GiB"
+                for path, size in cow_top
+            )
+        )
+    device = torch.get_device_module()
+    if hasattr(device, "memory_allocated"):
+        lines.append(
+            f"  device: allocated={device.memory_allocated() / GIB:.2f}GiB "
+            f"reserved={device.memory_reserved() / GIB:.2f}GiB"
+        )
+    snapshot = getattr(device, "memory_snapshot", None)
+    if callable(snapshot):
+        try:
+            segments = snapshot()
+        except Exception:
+            segments = []
+        ranges = [
+            (
+                int(s.get("address", 0)),
+                int(s.get("address", 0)) + int(s.get("total_size", 0)),
+            )
+            for s in segments
+        ]
+        malloc = _mallinfo()
+        if malloc:
+            lines.append(
+                "  glibc: " + " ".join(f"{k}={v:.2f}GiB" for k, v in malloc.items())
+            )
+        host_stats = getattr(device, "host_memory_stats", None)
+        if callable(host_stats):
+            try:
+                hs = host_stats()
+                big = {
+                    k: v
+                    for k, v in hs.items()
+                    if isinstance(v, (int, float)) and v >= 64 * 1024**2
+                }
+                lines.append(
+                    "  pinned host allocator: "
+                    + (
+                        " ".join(
+                            f"{k}={v / GIB:.2f}GiB" for k, v in sorted(big.items())
+                        )
+                        or "no counter >= 64 MiB"
+                    )
+                )
+            except Exception as exc:
+                lines.append(f"  pinned host allocator: unavailable ({exc})")
+        for start, end, anon, flags in _anon_vmas()[:16]:
+            inside = any(a <= start < b for a, b in ranges)
+            count, held = _tensor_ptrs_inside(start, end)
+            lines.append(
+                f"  anon vma {start:#x}-{end:#x}: {anon / GIB:.2f}GiB "
+                f"{'INSIDE cuda segment' if inside else 'outside cuda segments'} [{flags}] "
+                f"live cpu tensors inside={count} ({held:.2f}GiB) sample: {_sample_anon_vma(start, end)}"
+            )
+        try:
+            config = torch.__config__.show()
+            lines.append(
+                "  torch build: "
+                + ", ".join(
+                    line.strip()
+                    for line in config.splitlines()
+                    if "MIMALLOC" in line or "ALLOC" in line.upper() and "USE_" in line
+                )[:300]
+            )
+        except Exception:
+            pass
+        big = sorted(
+            (
+                (int(segment.get("address", 0)), int(segment.get("total_size", 0)))
+                for segment in segments
+                if int(segment.get("total_size", 0)) >= 256 * 1024**2
+            ),
+            key=lambda item: -item[1],
+        )[:12]
+        lines.append(
+            "  device segments >= 256 MiB: "
+            + ", ".join(f"{address:#x}:{size / GIB:.2f}GiB" for address, size in big)
+        )
+    logger.info("\n".join(lines))
+
+
+def log_anon_vmas(label: str) -> None:
+    """Debug: the process's large anonymous mappings right now, for a timeline of
+    where they appear during startup."""
+    rollup = _smaps_rollup()
+    lines = [
+        f"Anonymous memory timeline ({label}): "
+        + " ".join(f"{k}={v:.2f}GiB" for k, v in sorted(rollup.items()))
+    ]
+    for start, end, anon, flags in _anon_vmas()[:8]:
+        lines.append(f"  anon vma {start:#x}-{end:#x}: {anon / GIB:.2f}GiB [{flags}]")
+    logger.info("\n".join(lines))

python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py

@@ -18,6 +18,7 @@ import os
 import psutil
 
 from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -103,13 +104,36 @@ def _cgroup_dirs(mount: str) -> list[str]:
     return dirs
 
 
-def cgroup_memory_limit_bytes() -> tuple[int, int] | None:
+# memory.stat keys for the page cache a cgroup is charged for: v2, then v1.
+_CGROUP_FILE_CACHE_KEYS = ("file", "cache")
+
+
+def _cgroup_file_cache_bytes(directory: str) -> int:
+    try:
+        with open(os.path.join(directory, "memory.stat")) as handle:
+            for line in handle:
+                key, _, value = line.partition(" ")
+                if key in _CGROUP_FILE_CACHE_KEYS:
+                    return int(value)
+    except (OSError, ValueError):
+        pass
+    return 0
+
+
+def cgroup_memory_limit_bytes(
+    *, exclude_file_cache: bool = False
+) -> tuple[int, int] | None:
     """This process's (cap, usage) under its cgroup, or None when uncapped.
 
     The tightest cap in the chain wins. A nested cgroup -- a systemd scope with
     MemoryMax, a container started with --cgroup-parent -- holds this process
     below whatever the mount root allows, and planning against the root would
     commit memory the process cannot have.
+
+    A cgroup is charged for the page cache it touches, so its usage grows by
+    the whole checkpoint the process maps. ``exclude_file_cache`` reports the
+    anonymous share alone, for callers that may spend cache the kernel would
+    reclaim under the cap anyway.
     """
     for mount, limit_name, usage_name in _CGROUP_MOUNTS:
         tightest = None
@@ -119,7 +143,10 @@ def cgroup_memory_limit_bytes() -> tuple[int, int] | None:
                 continue
             if tightest is not None and limit >= tightest[0]:
                 continue
-            tightest = (limit, _read_int(os.path.join(directory, usage_name)) or 0)
+            usage = _read_int(os.path.join(directory, usage_name)) or 0
+            if exclude_file_cache:
+                usage = max(0, usage - _cgroup_file_cache_bytes(directory))
+            tightest = (limit, usage)
         if tightest is not None:
             return tightest
     return None
@@ -155,6 +182,33 @@ def host_memory_available_bytes() -> int:
     return min(available, max(0, limit - usage))
 
 
+def shared_pool_available_bytes() -> int:
+    """Bytes a shared host/device pool can still give this process.
+
+    The device's own free figure is the kernel's MemFree, which leaves out the
+    page cache -- memory the kernel hands back on demand and a placement may
+    therefore spend. A cgroup cap is honoured on its anonymous share only, for
+    the same reason: the cache charged to the cgroup is reclaimed under the cap.
+    """
+    available = int(psutil.virtual_memory().available)
+    capped = cgroup_memory_limit_bytes(exclude_file_cache=True)
+    if capped is None:
+        return available
+    limit, anonymous = capped
+    return min(available, max(0, limit - anonymous))
+
+
+def host_copies_are_redundant() -> bool:
+    """Whether a host copy of a mapped weight buys nothing.
+
+    When host and device share one physical pool the device reads page-cache
+    pages directly, so a pinned or pageable copy holds the same bytes twice and
+    adds only pressure. The mapping is then the right home for every weight
+    that has one, whatever the free-memory reading says.
+    """
+    return current_platform.device_shares_host_memory()
+
+
 def host_copies_would_not_fit(weight_bytes: int) -> bool:
     """Whether copying `weight_bytes` into host memory would run the host out.
 
@@ -186,7 +240,13 @@ class HostPinBudget:
 
     def __init__(self, available_bytes: int | None = None) -> None:
         if available_bytes is None:
-            available_bytes = host_memory_available_bytes()
+            if host_copies_are_redundant():
+                # Nothing to pin for: on one pool the copy duplicates
+                # page-cache bytes the device can already read, and the mapped
+                # courier overlaps its transfers anyway.
+                available_bytes = 0
+            else:
+                available_bytes = host_memory_available_bytes()
         self.available_bytes = available_bytes
         self.reserve_bytes = max(
             int(available_bytes * HOST_RESERVE_FRACTION), MIN_HOST_RESERVE_BYTES

python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py

@@ -1,7 +1,13 @@
 import bisect
+import ctypes
+import ctypes.util
+import mmap
+import os
 import queue
 import re
+import sys
 import threading
+import time
 from collections.abc import Mapping, Sequence
 from contextlib import nullcontext
 from time import perf_counter
@@ -29,6 +35,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency
 from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import (
     HostPinBudget,
     describe_host_memory,
+    host_copies_are_redundant,
     host_copies_would_not_fit,
     host_memory_available_bytes,
     module_weight_bytes,
@@ -206,6 +213,354 @@ def _install_host_gather_hooks(
     module.register_forward_hook(_output_to_device)
 
 
+_MADV_WILLNEED = 3
+_libc = None
+if sys.platform == "linux":
+    try:
+        _libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
+    except OSError:
+        _libc = None
+_PAGE = os.sysconf("SC_PAGESIZE") if hasattr(os, "sysconf") else 4096
+
+
+_WILLNEED_MIN_AVAILABLE = 4 << 30  # bytes of MemAvailable required to advise
+
+
+def _willneed_headroom_ok(need_bytes: int) -> bool:
+    """Advised pages need somewhere to land, or the advice backfires.
+
+    Field-measured on a 32 GB host with MemAvailable at 0.29 GiB: pages read
+    ahead were evicted before the courier reached them, so every byte was
+    read twice and effective throughput fell to 0.67x of the unadvised run
+    (0.85 vs 1.27 GB/s). Only advise when the kernel has real headroom to
+    keep the window resident until it is consumed.
+    """
+    try:
+        with open("/proc/meminfo") as handle:
+            for line in handle:
+                if line.startswith("MemAvailable:"):
+                    available = int(line.split()[1]) * 1024
+                    return available >= max(_WILLNEED_MIN_AVAILABLE, 2 * need_bytes)
+    except (OSError, ValueError):
+        pass
+    return False
+
+
+def advise_willneed(tensors) -> int:
+    """Ask the kernel to read these mapped tensors' pages ahead, in bulk.
+
+    The default readahead pipeline feeds the drive in read_ahead_kb-sized
+    beats (128 KB), which holds a fast NVMe to a quarter of its sequential
+    throughput on a host too small to cache the checkpoint. MADV_WILLNEED
+    schedules the whole range at once, so the disk read for the next layer
+    runs at drive speed while the current layer computes. Best-effort and
+    Linux-only: on any failure the normal fault path still works — and on a
+    host with no free headroom the advice is withheld entirely, because a
+    window that cannot stay resident until consumed is read twice.
+    """
+    if _libc is None:
+        return 0
+    tensors = list(tensors)
+    need = 0
+    for tensor in tensors:
+        try:
+            need += tensor.untyped_storage().nbytes()
+        except Exception:
+            continue
+    if need == 0 or not _willneed_headroom_ok(need):
+        return 0
+    advised = 0
+    for tensor in tensors:
+        try:
+            storage = tensor.untyped_storage()
+            ptr = storage.data_ptr()
+            nbytes = storage.nbytes()
+        except Exception:
+            continue
+        if ptr == 0 or nbytes == 0:
+            continue
+        start = ptr & ~(_PAGE - 1)
+        length = (ptr + nbytes) - start
+        length = (length + _PAGE - 1) & ~(_PAGE - 1)
+        if (
+            _libc.madvise(
+                ctypes.c_void_p(start), ctypes.c_size_t(length), _MADV_WILLNEED
+            )
+            == 0
+        ):
+            advised += 1
+    return advised
+
+
+_MADV_COLD = 20  # Linux 5.4+: deactivate the pages; reclaimed first under pressure
+_MADV_PAGEOUT = 21  # Linux 5.4+: reclaim the pages now
+
+
+_MADV_POPULATE_READ = 22  # Linux 5.14+: fault the range in, in one sequential pass
+# Streamed layers faulted in concurrently on a cold pass. One sequential
+# stream gets ~1 GiB/s from a GB10's NVMe. Within one checkpoint shard six
+# streams measured 1.7 GiB/s and twelve 3.0 GiB/s, so the pool runs twelve.
+MAPPED_POPULATE_AHEAD = 12
+
+
+def populate_mapped_source(tensors) -> int:
+    """Fault a layer's mapped pages in before a parallel copy reads them.
+
+    A multi-threaded memcpy over an uncached mapping faults from many offsets
+    at once, which the kernel's readahead heuristics read as random access:
+    the drive is then fed 4 KiB at a time and a 1.1 GiB/s NVMe delivers a
+    fifth of that (a 125 s first denoise step on a GB10). One synchronous
+    MADV_POPULATE_READ per tensor keeps the read sequential and full-speed;
+    on an older kernel it fails with EINVAL and the WILLNEED path remains.
+    """
+    if _libc is None:
+        return 0
+    populated = 0
+    for tensor in tensors:
+        try:
+            ptr = tensor.data_ptr()
+            nbytes = tensor.numel() * tensor.element_size()
+        except Exception:
+            continue
+        if ptr == 0 or nbytes == 0:
+            continue
+        start = ptr & ~(_PAGE - 1)
+        length = (ptr + nbytes) - start
+        length = (length + _PAGE - 1) & ~(_PAGE - 1)
+        if (
+            _libc.madvise(
+                ctypes.c_void_p(start), ctypes.c_size_t(length), _MADV_POPULATE_READ
+            )
+            == 0
+        ):
+            populated += 1
+        else:
+            advise_willneed([tensor])
+    return populated
+
+
+def _advise_mapped_source_cold(tensor: torch.Tensor, *, reclaim: bool = False) -> None:
+    """Tell the kernel a mapped tensor's file pages are cold once a copy holds them.
+
+    MADV_COLD deactivates the pages without dropping them: under pressure the
+    kernel reclaims them ahead of anything hot, and otherwise keeps them. A
+    just-copied 45 GiB encoder otherwise looks like the hottest data on the
+    box and the kernel swaps idle anonymous memory instead -- measured on a
+    GB10 as 15 GiB of swap traffic and a wedged host.
+
+    ``reclaim`` asks for MADV_PAGEOUT instead: the pages go now. That is for
+    a permanent materialization on a shared pool, where the device grows by
+    the same bytes the source pages hold and the driver satisfies device
+    allocations from free memory without waiting for cache reclaim -- a 57 GiB
+    DiT copied with MemFree near zero ended in NVRM out-of-memory twice.
+    """
+    if not sys.platform.startswith("linux") or tensor.device.type != "cpu":
+        return
+    if _libc is None:
+        return
+    page = mmap.PAGESIZE
+    start = tensor.data_ptr()
+    end = start + tensor.numel() * tensor.element_size()
+    start -= start % page
+    if end <= start:
+        return
+    _libc.madvise(
+        ctypes.c_void_p(start),
+        ctypes.c_size_t(end - start),
+        _MADV_PAGEOUT if reclaim else _MADV_COLD,
+    )
+
+
+def _shared_pool_hosting(
+    totals: Dict[int, int], mapped: Dict[int, int]
+) -> Dict[int, str]:
+    """Hosting when host and device draw from one pool.
+
+    A mapped layer stays mapped: the device reads page-cache pages directly, so
+    a pinned or pageable copy would hold the same bytes twice. Only a layer with
+    no mapping at all -- an anonymous fused weight -- keeps a pageable copy.
+    """
+    return {
+        layer_idx: "mapped" if mapped.get(layer_idx, 0) > 0 else "pageable"
+        for layer_idx in totals
+    }
+
+
+_DIRECT_ALIGN = 4096
+_DIRECT_CHUNK = 64 << 20
+# Components smaller than this keep the page-cache path: their pages fit the
+# cache next to a resident DiT, and a component re-streamed many times per
+# request (the H3 video VAE: 36 layers, 4.5 GiB, ~200 passes per decode)
+# must not go to the drive on every pass.
+MAPPED_DIRECT_READ_MIN_BYTES = 8 * 1024**3
+
+
+class _DirectReader:
+    """Read a mapped tensor's bytes from its checkpoint file with O_DIRECT.
+
+    Measured on a GB10: the NVMe delivers 9.9 GiB/s to an O_DIRECT reader,
+    but 1.1 GiB/s (one stream) to 3.7 GiB/s (twelve) through the page cache,
+    whose page allocation and reclaim are the wall on a shared pool. Reading
+    straight into the courier's pinned slot skips the cache entirely: no
+    pages to populate, evict or reclaim, and no cache footprint at all.
+    """
+
+    def __init__(self) -> None:
+        self._ranges: List[tuple[int, int, int, str]] = []
+        self._starts: List[int] = []
+        self._fds: Dict[str, int] = {}
+        self._located: Dict[int, Optional[tuple[str, int, int]]] = {}
+        self.refresh()
+
+    def refresh(self) -> None:
+        ranges = []
+        try:
+            with open("/proc/self/maps") as handle:
+                for line in handle:
+                    fields = line.split()
+                    if len(fields) < 6 or fields[5].startswith("["):
+                        continue
+                    start, end = (int(x, 16) for x in fields[0].split("-"))
+                    ranges.append((start, end, int(fields[2], 16), fields[5]))
+        except OSError:
+            ranges = []
+        ranges.sort()
+        self._ranges = ranges
+        self._starts = [r[0] for r in ranges]
+        self._located.clear()
+
+    def locate(self, tensor: torch.Tensor) -> Optional[tuple[str, int, int]]:
+        """(path, file offset, nbytes) of the tensor's bytes, or None if unmapped."""
+        ptr = tensor.data_ptr()
+        nbytes = tensor.numel() * tensor.element_size()
+        cached = self._located.get(ptr)
+        if cached is not None or ptr in self._located:
+            return cached
+        found = None
+        for attempt in range(2):
+            index = bisect.bisect_right(self._starts, ptr) - 1
+            if index >= 0:
+                start, end, file_offset, path = self._ranges[index]
+                if start <= ptr and ptr + nbytes <= end:
+                    found = (path, file_offset + (ptr - start), nbytes)
+                    break
+            if attempt == 0:
+                self.refresh()
+        self._located[ptr] = found
+        return found
+
+    def fd(self, path: str) -> int:
+        fd = self._fds.get(path)
+        if fd is None:
+            fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECT"))
+            self._fds[path] = fd
+        return fd
+
+    def read_into(
+        self, buffer: memoryview, path: str, aligned_offset: int, span: int
+    ) -> None:
+        """Fill buffer[:span] with the file bytes at aligned_offset (both 4 KiB aligned)."""
+        fd = self.fd(path)
+        pos = 0
+        while pos < span:
+            want = min(_DIRECT_CHUNK, span - pos)
+            got = os.preadv(fd, [buffer[pos : pos + want]], aligned_offset + pos)
+            if got <= 0:
+                # past the end of the file: the tail of the last aligned span
+                break
+            pos += got
+
+    def close(self) -> None:
+        for fd in self._fds.values():
+            try:
+                os.close(fd)
+            except OSError:
+                pass
+        self._fds.clear()
+
+
+def _aligned_span(file_offset: int, nbytes: int) -> tuple[int, int]:
+    """(aligned start, span) covering [file_offset, file_offset + nbytes) at 4 KiB granularity."""
+    start = file_offset & ~(_DIRECT_ALIGN - 1)
+    end = (file_offset + nbytes + _DIRECT_ALIGN - 1) & ~(_DIRECT_ALIGN - 1)
+    return start, end - start
+
+
+class _MappedPopulator:
+    """Fault upcoming mapped layers in from several threads at once.
+
+    On the pass that may find pages cold, the manager hands the next few
+    streamed layers to this pool so the drive stays saturated while the
+    courier stages and the model computes the current one. A layer is
+    populated once per request; ``reset`` starts the next request over.
+    """
+
+    def __init__(self, workers: int = MAPPED_POPULATE_AHEAD) -> None:
+        self._tasks: queue.Queue[Optional[tuple[int, list]]] = queue.Queue()
+        self._lock = threading.Condition()
+        self._pending: Set[int] = set()
+        self._done: Set[int] = set()
+        self.stats = {"populate_s": 0.0, "bytes": 0, "layers": 0}
+        self._threads = [
+            threading.Thread(target=self._run, name=f"mapped-populate-{i}", daemon=True)
+            for i in range(max(1, workers))
+        ]
+        for thread in self._threads:
+            thread.start()
+
+    def submit(self, layer_idx: int, tensors) -> bool:
+        tensors = list(tensors)
+        if not tensors:
+            return False
+        with self._lock:
+            if layer_idx in self._pending or layer_idx in self._done:
+                return False
+            self._pending.add(layer_idx)
+        self._tasks.put((layer_idx, tensors))
+        return True
+
+    def wait(self, layer_idx: int) -> bool:
+        """Block until a submitted layer is populated; False if it never was."""
+        with self._lock:
+            if layer_idx not in self._pending and layer_idx not in self._done:
+                return False
+            while layer_idx in self._pending:
+                self._lock.wait()
+            return layer_idx in self._done
+
+    def reset(self) -> None:
+        with self._lock:
+            self._done.clear()
+
+    def close(self) -> None:
+        for _ in self._threads:
+            self._tasks.put(None)
+        for thread in self._threads:
+            thread.join(timeout=5.0)
+
+    def _run(self) -> None:
+        while True:
+            task = self._tasks.get()
+            if task is None:
+                return
+            layer_idx, tensors = task
+            started = time.perf_counter()
+            try:
+                populate_mapped_source(tensors)
+            except Exception:  # advisory only; the courier reads the pages anyway
+                pass
+            with self._lock:
+                self.stats["populate_s"] += time.perf_counter() - started
+                self.stats["bytes"] += sum(
+                    t.numel() * t.element_size() for t in tensors
+                )
+                self.stats["layers"] += 1
+            with self._lock:
+                self._pending.discard(layer_idx)
+                self._done.add(layer_idx)
+                self._lock.notify_all()
+
+
 class MappedLayerCourier:
     """Ships a mapped layer's weights to the device off the compute thread.
 
@@ -232,24 +587,69 @@ class MappedLayerCourier:
         weight_metadata: Dict[int, Dict[str, Dict[str, Any]]],
         device: torch.device,
         pin_slots: bool,
+        cold_source: Optional[Callable[[int], bool]] = None,
+        populate_source: Optional[Callable[[int], bool]] = None,
+        await_populated: Optional[Callable[[int], bool]] = None,
+        direct_copy: bool = False,
+        direct_read: bool = False,
     ) -> None:
         self._mapped_cpu_weights = mapped_cpu_weights
+        # Read each layer's bytes from the checkpoint file with O_DIRECT into the
+        # pinned slot instead of copying them out of the page cache.
+        self.direct_read = bool(direct_read) and hasattr(os, "O_DIRECT")
+        self._reader: Optional[_DirectReader] = (
+            _DirectReader() if self.direct_read else None
+        )
+        self._slot_views: Dict[int, Any] = {}
+        # On a shared host/device pool the device reads host pages directly, so
+        # a layer goes mapping -> device in one copy and the pinned staging
+        # slots (2.7 GiB on the H3 encoder + VAE) are not allocated at all.
+        self._direct_copy = direct_copy
+        self.stats = {
+            "layers": 0,
+            "bytes": 0,
+            "slot_sync_s": 0.0,
+            "populate_wait_s": 0.0,
+            "populate_s": 0.0,
+            "memcpy_s": 0.0,
+            "h2d_issue_s": 0.0,
+            "direct_read_s": 0.0,
+            "direct_read_bytes": 0,
+        }
+        # Blocks until a populator thread has faulted the layer in; True if one
+        # did, so the courier does not fault the same range a second time.
+        self._await_populated = await_populated
         self._weight_metadata = weight_metadata
         self._device = device
+        # Whether a layer's file pages may go cold once its copy is staged:
+        # true for resident layers, never for layers re-read every step.
+        self._cold_source = cold_source
+        # Whether to fault the layer in sequentially before the staging copy:
+        # true on a request's first pass over the layers, when pages may be cold.
+        self._populate_source = populate_source
         slot_bytes = max(
             (
-                sum(t.numel() * t.element_size() for t in weights.values())
+                sum(
+                    t.numel() * t.element_size()
+                    + (2 * _DIRECT_ALIGN if self.direct_read else 0)
+                    for t in weights.values()
+                )
                 for weights in mapped_cpu_weights.values()
                 if weights
             ),
             default=0,
         )
+        slot_bytes = (slot_bytes + _DIRECT_ALIGN - 1) & ~(_DIRECT_ALIGN - 1)
         if slot_bytes <= 0:
             raise ValueError("no mapped weights to ship")
-        self._slots = [
-            torch.empty(slot_bytes, dtype=torch.uint8, pin_memory=pin_slots)
-            for _ in range(self._NUM_SLOTS)
-        ]
+        self._slots = (
+            []
+            if direct_copy
+            else [
+                torch.empty(slot_bytes, dtype=torch.uint8, pin_memory=pin_slots)
+                for _ in range(self._NUM_SLOTS)
+            ]
+        )
         self._slot_events: List[Optional[Any]] = [None] * self._NUM_SLOTS
         self._stream = torch.get_device_module().Stream()
         self._tasks: queue.Queue[Optional[int]] = queue.Queue()
@@ -292,6 +692,8 @@ class MappedLayerCourier:
     def close(self) -> None:
         self._tasks.put(None)
         self._thread.join(timeout=5.0)
+        if self._reader is not None:
+            self._reader.close()
 
     def _run(self) -> None:
         slot_turn = 0
@@ -313,37 +715,129 @@ class MappedLayerCourier:
                 return
 
     def _ship(self, layer_idx: int, slot_turn: int):
-        slot = self._slots[slot_turn]
-        previous = self._slot_events[slot_turn]
-        if previous is not None:
-            # the previous transfer through this slot must land before reuse
-            previous.synchronize()
+        stats = self.stats
         tensors: Dict[str, torch.Tensor] = {}
-        offset = 0
+        if not self._direct_copy:
+            slot = self._slots[slot_turn]
+            previous = self._slot_events[slot_turn]
+            if previous is not None:
+                # the previous transfer through this slot must land before reuse
+                started = time.perf_counter()
+                previous.synchronize()
+                stats["slot_sync_s"] += time.perf_counter() - started
+        if (
+            not self.direct_read
+            and self._populate_source is not None
+            and self._populate_source(layer_idx)
+        ):
+            started = time.perf_counter()
+            if self._await_populated is not None and self._await_populated(layer_idx):
+                stats["populate_wait_s"] += time.perf_counter() - started
+            else:
+                populate_mapped_source(self._mapped_cpu_weights[layer_idx].values())
+                stats["populate_s"] += time.perf_counter() - started
+        layer_bytes = 0
         with torch.inference_mode(False), torch.no_grad():
-            staged = []
-            for name, cpu_tensor in self._mapped_cpu_weights[layer_idx].items():
-                width = cpu_tensor.element_size()
-                if offset % width:
-                    offset += width - (offset % width)
-                start = offset // width
-                window = slot.view(cpu_tensor.dtype)[
-                    start : start + cpu_tensor.numel()
-                ].view(cpu_tensor.shape)
-                window.copy_(cpu_tensor)
-                offset += cpu_tensor.numel() * width
-                staged.append((name, window))
             event = torch.get_device_module().Event()
-            with torch.get_device_module().stream(self._stream):
-                for name, window in staged:
-                    meta = self._weight_metadata[layer_idx][name]
-                    gpu_tensor = torch.empty(
-                        meta["shape"], dtype=meta["dtype"], device=self._device
+            if self._direct_copy:
+                started = time.perf_counter()
+                with torch.get_device_module().stream(self._stream):
+                    for name, cpu_tensor in self._mapped_cpu_weights[layer_idx].items():
+                        meta = self._weight_metadata[layer_idx][name]
+                        gpu_tensor = torch.empty(
+                            meta["shape"], dtype=meta["dtype"], device=self._device
+                        )
+                        gpu_tensor.copy_(cpu_tensor, non_blocking=True)
+                        layer_bytes += cpu_tensor.numel() * cpu_tensor.element_size()
+                        if self._cold_source is not None and self._cold_source(
+                            layer_idx
+                        ):
+                            _advise_mapped_source_cold(cpu_tensor, reclaim=True)
+                        tensors[name] = gpu_tensor
+                    event.record(self._stream)
+                stats["h2d_issue_s"] += time.perf_counter() - started
+            else:
+                offset = 0
+                staged = []
+                started = time.perf_counter()
+                slot_bytes_view = None
+                if self.direct_read:
+                    slot_bytes_view = self._slot_views.get(slot_turn)
+                    if slot_bytes_view is None:
+                        slot_bytes_view = memoryview(slot.numpy())
+                        self._slot_views[slot_turn] = slot_bytes_view
+                for name, cpu_tensor in self._mapped_cpu_weights[layer_idx].items():
+                    width = cpu_tensor.element_size()
+                    nbytes = cpu_tensor.numel() * width
+                    located = (
+                        self._reader.locate(cpu_tensor) if self.direct_read else None
                     )
-                    gpu_tensor.copy_(window, non_blocking=True)
-                    tensors[name] = gpu_tensor
-                event.record(self._stream)
-        self._slot_events[slot_turn] = event
+                    if located is not None:
+                        path, file_offset, _ = located
+                        aligned_start, span = _aligned_span(file_offset, nbytes)
+                        skew = file_offset - aligned_start
+                        if offset % _DIRECT_ALIGN:
+                            offset += _DIRECT_ALIGN - (offset % _DIRECT_ALIGN)
+                        if skew % width == 0 and offset + span <= slot.numel():
+                            read_started = time.perf_counter()
+                            try:
+                                self._reader.read_into(
+                                    slot_bytes_view[offset : offset + span],
+                                    path,
+                                    aligned_start,
+                                    span,
+                                )
+                            except OSError as exc:
+                                logger.warning(
+                                    "Layerwise offload: O_DIRECT read failed (%s); "
+                                    "mapped layers return to the page-cache copy.",
+                                    exc,
+                                )
+                                self.direct_read = False
+                                located = None
+                            else:
+                                stats["direct_read_s"] += (
+                                    time.perf_counter() - read_started
+                                )
+                                stats["direct_read_bytes"] += span
+                                window = (
+                                    slot[offset + skew : offset + skew + nbytes]
+                                    .view(cpu_tensor.dtype)
+                                    .view(cpu_tensor.shape)
+                                )
+                                offset += span
+                        else:
+                            located = None
+                    if located is None:
+                        if offset % width:
+                            offset += width - (offset % width)
+                        start = offset // width
+                        window = slot.view(cpu_tensor.dtype)[
+                            start : start + cpu_tensor.numel()
+                        ].view(cpu_tensor.shape)
+                        window.copy_(cpu_tensor)
+                        if self._cold_source is not None and self._cold_source(
+                            layer_idx
+                        ):
+                            _advise_mapped_source_cold(cpu_tensor, reclaim=True)
+                        offset += nbytes
+                    layer_bytes += nbytes
+                    staged.append((name, window))
+                stats["memcpy_s"] += time.perf_counter() - started
+                started = time.perf_counter()
+                with torch.get_device_module().stream(self._stream):
+                    for name, window in staged:
+                        meta = self._weight_metadata[layer_idx][name]
+                        gpu_tensor = torch.empty(
+                            meta["shape"], dtype=meta["dtype"], device=self._device
+                        )
+                        gpu_tensor.copy_(window, non_blocking=True)
+                        tensors[name] = gpu_tensor
+                    event.record(self._stream)
+                stats["h2d_issue_s"] += time.perf_counter() - started
+                self._slot_events[slot_turn] = event
+        stats["layers"] += 1
+        stats["bytes"] += layer_bytes
         return event, tensors
 
 
@@ -414,6 +908,14 @@ class LayerwiseOffloadManager:
         # Armed on the first denoise forward, so that the load-time prefetch below
         # does not pin the whole resident set before the DiT is the active component.
         self._residency_active = False
+        # True while load_all_layers materializes every layer for a resident
+        # placement; every mapped source is then read exactly once.
+        self._materializing_all = False
+        # True from a request's start until its last layer has run once: the
+        # pass in which a mapped layer's pages may not be in the page cache.
+        self._first_pass = True
+        self._mapped_populator: Optional[_MappedPopulator] = None
+        self._debug_collect_wait_s = 0.0
         # True once _initialize builds the CPU buffers; unlike `enabled` it
         # never flips back, so disable_offload/enable_offload can toggle
         # `enabled` without losing track of which managers can be re-armed.
@@ -614,6 +1116,19 @@ class LayerwiseOffloadManager:
         buys a whole layer's worth of per-step overlap.
         """
         totals, mapped = self._layer_byte_totals(layer_groups)
+        if host_copies_are_redundant():
+            hosting = _shared_pool_hosting(totals, mapped)
+            logger.info(
+                "Layerwise offload: %s keeps %d of %d layers on the checkpoint "
+                "mapping (host and device share one memory pool, so a pinned or "
+                "pageable copy would hold the same bytes twice); %d layers "
+                "without a mapping stay pageable.",
+                self._pin_component_name,
+                sum(1 for where in hosting.values() if where == "mapped"),
+                len(totals),
+                sum(1 for where in hosting.values() if where == "pageable"),
+            )
+            return hosting
         pinned_bytes = 0
         hosting: Dict[int, str] = {}
         pin_order: List[int] = []
@@ -956,6 +1471,134 @@ class LayerwiseOffloadManager:
             if layer_idx not in retain:
                 self.release_layer(layer_idx)
 
+    def advise_mapped_pages_cold(self, *, room_bytes: Optional[int] = None) -> int:
+        """Hand this stage's mapped pages back in the order that keeps the next
+        request fast. Returns the bytes paged out.
+
+        With host and device in one pool the page cache cannot always hold the
+        whole request cycle. Plain LRU then evicts, at each phase boundary,
+        exactly the pages the next phase needs: a cyclic scan just larger
+        than the cache misses everywhere (measured on a GB10: ~100 GiB re-read
+        per request). Deactivating a component's pages (MADV_COLD) once its
+        stage is over fixes the cross-component case.
+
+        When the component's own stream is larger than the room the cache will
+        have for it (``room_bytes``), the same pathology happens inside the
+        stage: faulting layer 40 in evicts layer 0, which the next request
+        reads first. Measured on a GB10 with a 45 GiB encoder and ~30 GiB of
+        room, every page fault reclaimed synchronously and the stage took
+        38 s at 1.2 GiB/s. So the head of the stream -- as many layers as do
+        not fit -- is paged out now, deterministically, and the tail is left
+        cold: the next request reads the head from disk at full parallel
+        speed into free pages and hits the cache for everything else.
+        """
+        order = [
+            idx for idx in self._streamed_order if self._mapped_cpu_weights.get(idx)
+        ]
+        order += [
+            idx
+            for idx in self._mapped_cpu_weights
+            if idx not in set(order) and self._mapped_cpu_weights.get(idx)
+        ]
+        layer_bytes = {
+            idx: sum(
+                t.numel() * t.element_size()
+                for t in self._mapped_cpu_weights[idx].values()
+            )
+            for idx in order
+        }
+        total = sum(layer_bytes.values())
+        excess = 0
+        if room_bytes is not None and total > room_bytes:
+            excess = total - room_bytes
+        paged_out = 0
+        for idx in order:
+            reclaim = paged_out < excess
+            for tensor in self._mapped_cpu_weights[idx].values():
+                _advise_mapped_source_cold(tensor, reclaim=reclaim)
+            if reclaim:
+                paged_out += layer_bytes[idx]
+        return paged_out
+
+    def _ensure_mapped_populator(self) -> Optional[_MappedPopulator]:
+        if self._mapped_populator is None and _libc is not None:
+            self._mapped_populator = _MappedPopulator()
+        return self._mapped_populator
+
+    def _await_mapped_populated(self, layer_idx: int) -> bool:
+        populator = self._mapped_populator
+        return populator is not None and populator.wait(layer_idx)
+
+    def _populate_ahead(self, layer_idx: int) -> None:
+        """On a cold pass, fault the next streamed layers in from parallel threads."""
+        populator = self._ensure_mapped_populator()
+        if populator is None:
+            return
+        for ahead in self._next_streamed(after=layer_idx, count=MAPPED_POPULATE_AHEAD):
+            if ahead in self._gpu_layers or ahead in self._courier_inflight:
+                continue
+            populator.submit(ahead, self._mapped_cpu_weights.get(ahead, {}).values())
+
+    def _log_debug_timing(self) -> None:
+        """Debug: where this stage's layer traffic spent its time."""
+        if not envs.SGLANG_DIFFUSION_DEBUG_LAYERWISE_TIMING:
+            return
+        courier = self._mapped_courier
+        populator = self._mapped_populator
+        if courier is None or not courier.stats["layers"]:
+            self._debug_collect_wait_s = 0.0
+            return
+        cs = courier.stats
+        ps = (
+            populator.stats
+            if populator is not None
+            else {"populate_s": 0.0, "bytes": 0, "layers": 0}
+        )
+        logger.info(
+            "Layerwise timing %s: layers=%d bytes=%.2fGiB direct=%s | courier: slot_sync=%.2fs "
+            "populate_wait=%.2fs populate=%.2fs memcpy=%.2fs (direct_read=%.2fs %.2fGiB) "
+            "h2d_issue=%.2fs | populator: "
+            "layers=%d bytes=%.2fGiB busy=%.2fs | compute thread collect wait=%.2fs",
+            self.layers_attr_str,
+            cs["layers"],
+            cs["bytes"] / (1024**3),
+            courier._direct_copy,
+            cs["slot_sync_s"],
+            cs["populate_wait_s"],
+            cs["populate_s"],
+            cs["memcpy_s"],
+            cs["direct_read_s"],
+            cs["direct_read_bytes"] / (1024**3),
+            cs["h2d_issue_s"],
+            ps["layers"],
+            ps["bytes"] / (1024**3),
+            ps["populate_s"],
+            self._debug_collect_wait_s,
+        )
+        for key in cs:
+            cs[key] = 0.0 if isinstance(cs[key], float) else 0
+        if populator is not None:
+            with populator._lock:
+                for key in ps:
+                    ps[key] = 0.0 if isinstance(ps[key], float) else 0
+        self._debug_collect_wait_s = 0.0
+
+    def _mapped_source_may_be_cold(self, layer_idx: int) -> bool:
+        """Whether this copy is the first read of the layer in this request."""
+        return self._first_pass or self._materializing_all
+
+    def _mapped_source_is_cold(self, layer_idx: int) -> bool:
+        """Whether a mapped layer's file pages may go once its device copy lands.
+
+        Only while every layer is being materialized for a permanent resident
+        placement: those pages are not read again until a demotion. A
+        stage-scoped resident set is re-armed from the same pages on the next
+        request, and a streamed layer is re-read every step; marking either
+        cold made the kernel evict exactly what the next request needed
+        (measured on a GB10 as a 60 s first denoise step re-reading 23 GiB).
+        """
+        return self._materializing_all
+
     @torch.compiler.disable
     def _activate_residency(self) -> None:
         """Arm the resident set on the first denoise forward. The pinning itself is
@@ -1021,6 +1664,21 @@ class LayerwiseOffloadManager:
             if courier is not None and courier.submit(layer_idx):
                 self._courier_inflight.add(layer_idx)
                 ship_mapped = True
+                if (
+                    not envs.SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED
+                    and not courier.direct_read
+                ):
+                    # Schedule the disk read for this layer's pages now, in
+                    # one bulk request, so it overlaps the previous layer's
+                    # compute instead of trickling in at fault-time beats.
+                    advise_willneed(self._mapped_cpu_weights[layer_idx].values())
+                    if self._first_pass:
+                        # On the pass that may find pages cold, keep the drive
+                        # saturated several layers ahead from parallel threads:
+                        # one sequential stream idles the NVMe at ~1 GiB/s
+                        # while a layer is staged and computed; twelve streams
+                        # within one shard measured ~3 GiB/s on a GB10.
+                        self._populate_ahead(layer_idx)
 
         # create gpu buffer and load from CPU buffer
         gpu_buffers: Dict[torch.dtype, torch.Tensor] = {}
@@ -1050,10 +1708,16 @@ class LayerwiseOffloadManager:
                     # on the compute thread rather than ahead of it, and a page
                     # the kernel has reclaimed is faulted back in here.
                     cpu_tensor = self._mapped_cpu_weights[layer_idx][name]
+                    if not envs.SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED:
+                        # A blocking read on this thread: fault the range in
+                        # sequentially so a cold cache fills at drive speed.
+                        populate_mapped_source([cpu_tensor])
                     gpu_tensor = torch.empty(
                         meta["shape"], dtype=meta["dtype"], device=self.device
                     )
                     gpu_tensor.copy_(cpu_tensor, non_blocking=False)
+                    if self._mapped_source_is_cold(layer_idx):
+                        _advise_mapped_source_cold(cpu_tensor, reclaim=True)
                     target.data = self._wrap_for_target(target, gpu_tensor)
                     continue
 
@@ -1107,6 +1771,19 @@ class LayerwiseOffloadManager:
                 weight_metadata=self._weight_metadata,
                 device=self.device,
                 pin_slots=current_platform.is_cuda(),
+                # Measured on a GB10: a copy_ straight from file-backed pages ran
+                # at 0.67 GiB/s (the driver stages pageable sources page by
+                # page) and the process's anonymous memory grew past 100 GiB;
+                # the pinned slots stay even on a shared pool.
+                direct_copy=False,
+                direct_read=(
+                    host_copies_are_redundant()
+                    and not envs.SGLANG_DIFFUSION_DISABLE_MAPPED_DIRECT_READ
+                    and self._mapped_bytes >= MAPPED_DIRECT_READ_MIN_BYTES
+                ),
+                cold_source=self._mapped_source_is_cold,
+                populate_source=self._mapped_source_may_be_cold,
+                await_populated=self._await_mapped_populated,
             )
             logger.info(
                 "Layerwise offload: %s ships mapped layers through a courier "
@@ -1128,6 +1805,7 @@ class LayerwiseOffloadManager:
     def _collect_mapped_layer(self, layer_idx: int) -> None:
         """Bind a shipped layer's tensors on the compute thread."""
         courier = self._mapped_courier
+        started = time.perf_counter()
         try:
             event, tensors = courier.collect(layer_idx)
         except BaseException as exc:
@@ -1143,6 +1821,10 @@ class LayerwiseOffloadManager:
             self._courier_inflight.discard(layer_idx)
             self.prefetch_layer(layer_idx, non_blocking=False)
             return
+        # debug counter; a manager built without __init__ (tests) has none yet
+        self._debug_collect_wait_s = getattr(self, "_debug_collect_wait_s", 0.0) + (
+            time.perf_counter() - started
+        )
         compute_stream = torch.get_device_module().current_stream()
         compute_stream.wait_event(event)
         with torch.inference_mode(False), torch.no_grad():
@@ -1195,6 +1877,9 @@ class LayerwiseOffloadManager:
     def release_all(self) -> None:
         """Release every layer, including the resident ones: this ends the
         denoise stage that the resident set is scoped to."""
+        self._log_debug_timing()
+        if self._mapped_populator is not None:
+            self._mapped_populator.reset()
         if not self.enabled or self.device is None:
             return
         if self.copy_stream is not None:
@@ -1207,6 +1892,9 @@ class LayerwiseOffloadManager:
 
         for layer_idx in list(self._gpu_layers):
             self.release_layer(layer_idx, force=True)
+        # The next use starts a new request; its first pass over the layers may
+        # find their pages evicted and is the one worth faulting in sequentially.
+        self._first_pass = True
 
     @torch.compiler.disable
     def load_all_layers(self) -> None:
@@ -1216,9 +1904,54 @@ class LayerwiseOffloadManager:
         if self.copy_stream is not None:
             torch.get_device_module().current_stream().wait_stream(self.copy_stream)
 
-        for layer_idx in range(self.num_layers):
-            if layer_idx not in self._gpu_layers:
-                self.prefetch_layer(layer_idx, non_blocking=False)
+        self._materializing_all = True
+        try:
+            for layer_idx in range(self.num_layers):
+                if layer_idx not in self._gpu_layers:
+                    # Anonymous host stores can fill the copy stream without a
+                    # per-layer host wait. Checkpoint mappings still use the
+                    # synchronous path: the mapped courier has a bounded slot
+                    # ring intended to overlap one forward, not materialize a
+                    # whole model at once.
+                    self.prefetch_layer(
+                        layer_idx,
+                        non_blocking=not bool(self._mapped_cpu_weights.get(layer_idx)),
+                    )
+        finally:
+            self._materializing_all = False
+        if self.copy_stream is not None:
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
+
+    def release_host_stores(self) -> None:
+        """Drop rollback stores after a resident placement is validated.
+
+        The real device tensors must already be materialized and the manager
+        disabled. Repacking pinned stores as pageable here would copy the full
+        checkpoint for data that will never be streamed again.
+        """
+        if self.enabled:
+            raise RuntimeError("cannot release host stores while offload is enabled")
+        if self._mapped_courier is not None:
+            self._mapped_courier.close()
+            self._mapped_courier = None
+        if self._mapped_populator is not None:
+            self._mapped_populator.close()
+            self._mapped_populator = None
+        if self._courier_inflight:
+            raise RuntimeError(
+                "cannot release host stores with mapped copies in flight"
+            )
+
+        self._pin_budget.release(self.pinned_host_weight_bytes())
+        self._consolidated_cpu_weights.clear()
+        self._strided_cpu_weights.clear()
+        self._mapped_cpu_weights.clear()
+        self._mps_cpu_weights.clear()
+        self._weight_metadata.clear()
+        self._layer_hosting.clear()
+        self._prefetch_events.clear()
+        self._mapped_bytes = 0
+        self._configured = False
 
     @torch.compiler.disable
     def sync_layer_to_cpu(self, layer_idx: int) -> None:
@@ -1842,6 +2575,12 @@ class LayerwiseOffloadableModuleMixin:
         policies = ", ".join(sorted({manager.residency_policy for manager in managers}))
         total_layers = sum(manager.num_layers for manager in managers)
         resident_layers = sum(manager.resident_layers for manager in managers)
+        if envs.SGLANG_DIFFUSION_DEBUG_HOST_MEMORY:
+            from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_breakdown import (
+                log_anon_vmas,
+            )
+
+            log_anon_vmas(f"layerwise offload ready for {component_name}")
         logger.info(
             "Layerwise offload ready for %s in %.2fs: groups=%d, layers=%d, "
             "prefetch/group=%s, resident=%d/%d, policy=%s",

python/sglang/multimodal_gen/runtime/platforms/cuda.py

@@ -45,6 +45,12 @@ pynvml = import_pynvml()  # type: ignore[no-untyped-call]
 torch.backends.cuda.enable_cudnn_sdp(False)
 
 
+@lru_cache(maxsize=None)
+def _device_is_integrated(device_index: int) -> bool:
+    # A static device property, asked on every planner cost evaluation.
+    return bool(torch.cuda.get_device_properties(device_index).is_integrated)
+
+
 def device_id_to_physical_device_id(device_id: int) -> int:
     if "CUDA_VISIBLE_DEVICES" in os.environ:
         device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
@@ -628,6 +634,15 @@ class CudaPlatformBase(Platform):
 
         return free_gpu_memory / (1 << 30)
 
+    @classmethod
+    def device_shares_host_memory(cls) -> bool:
+        if not torch.cuda.is_available():
+            return False
+        try:
+            return _device_is_integrated(torch.cuda.current_device())
+        except (RuntimeError, AssertionError):
+            return False
+
     @classmethod
     def _resolve_default_attn_backend(cls) -> AttentionBackendEnum:
         if cls.is_sm120():

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -434,6 +434,16 @@ class Platform:
         """Whether automatic DiT layerwise offload is enabled on this platform."""
         return True
 
+    @classmethod
+    def device_shares_host_memory(cls) -> bool:
+        """Whether the accelerator draws from the same physical pool as the host.
+
+        On such a part (DGX Spark's GB10, Jetson) a device allocation is host
+        memory the kernel no longer has, and a host copy of a mapped weight is
+        a second copy of bytes the page cache already holds.
+        """
+        return False
+
     @classmethod
     def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module:
         """Apply platform-specific optimizations to VAE after loading."""

python/sglang/multimodal_gen/runtime/server_args/auto_tune.py

@@ -4,6 +4,7 @@ ServerArgsAutoTuner tunes the ServerArgs based on the desired performance mode
 
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING
 
 from sglang.multimodal_gen import envs
@@ -88,6 +89,28 @@ IMAGE_GEN_KEEP_RESIDENT_MIN_AVAILABLE_GB = 45.0
 DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB = 120.0
 
 
+# torch's CPU allocator (mimalloc since 2.13) keeps freed pages in its arenas
+# and backs them with transparent huge pages, so the fused-weight copies the
+# loader frees once a component is promoted stayed resident: 18.2 GiB of
+# anonymous memory on a GB10 after the DiT went resident, 5.0 GiB with these.
+# Read at process start, so they are set for the workers to inherit; a torch
+# without mimalloc ignores them.
+SHARED_POOL_CPU_ALLOCATOR_DEFAULTS = {
+    "MIMALLOC_PURGE_DELAY": "0",
+    "MIMALLOC_ALLOW_LARGE_OS_PAGES": "0",
+}
+
+
+def apply_shared_pool_cpu_allocator_defaults(environ) -> list[str]:
+    """Set the CPU allocator defaults not already chosen; return the names set."""
+    applied = []
+    for name, value in SHARED_POOL_CPU_ALLOCATOR_DEFAULTS.items():
+        if name not in environ:
+            environ[name] = value
+            applied.append(name)
+    return applied
+
+
 class ServerArgsAutoTuner:
     """Auto-tunes the server-arg for the given performance-mode, based on practical deployment experience with different model architectures"""
 
@@ -456,6 +479,104 @@ class ServerArgsAutoTuner:
             args.text_encoder_cpu_offload = False
         if args.image_encoder_cpu_offload is None:
             args.image_encoder_cpu_offload = False
+        if (
+            args.pin_cpu_memory
+            and not args.is_arg_explicitly_set("pin_cpu_memory")
+            and current_platform.device_shares_host_memory()
+        ):
+            # The device reads host pages directly on a shared pool, so a
+            # pinned copy of a mapped weight is the same bytes held twice.
+            args.pin_cpu_memory = False
+            logger.info(
+                "Host and device share one memory pool: pinned host weight "
+                "copies are disabled (pass --pin-cpu-memory true to override)."
+            )
+        if (
+            current_platform.device_shares_host_memory()
+            and "PYTORCH_CUDA_ALLOC_CONF" not in os.environ
+        ):
+            # Every byte the caching allocator keeps reserved is a byte the
+            # page cache -- the home of every mapped weight here -- cannot
+            # hold. Measured on a GB10: ~30 GiB of reserved-but-idle segments
+            # forced the encoder and the DiT to take turns being re-read from
+            # disk. Expandable segments let the reserve follow the live peak.
+            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+            logger.info(
+                "Host and device share one memory pool: PYTORCH_CUDA_ALLOC_CONF="
+                "expandable_segments:True so the allocator's reserve does not "
+                "crowd out the page cache."
+            )
+        if current_platform.device_shares_host_memory():
+            applied = apply_shared_pool_cpu_allocator_defaults(os.environ)
+            if applied:
+                logger.info(
+                    "Host and device share one memory pool: %s so the CPU "
+                    "allocator returns freed weight copies to the pool.",
+                    " ".join(f"{name}={os.environ[name]}" for name in applied),
+                )
+        if current_platform.device_shares_host_memory():
+            try:
+                import psutil
+
+                swap_total = psutil.swap_memory().total
+            except Exception:
+                swap_total = 0
+            try:
+                # A cgroup with swap disabled (memory.swap.max = 0) protects
+                # this process whatever the host has mounted.
+                with open("/sys/fs/cgroup/memory.swap.max") as handle:
+                    if handle.read().strip() == "0":
+                        swap_total = 0
+            except OSError:
+                pass
+            try:
+                with open("/sys/fs/cgroup/memory.max") as handle:
+                    uncapped = handle.read().strip() == "max"
+            except OSError:
+                uncapped = True
+            if uncapped:
+                # The driver takes device memory from free pages and does not
+                # wait for the kernel to reclaim page cache: with the cache
+                # full and MemFree near zero, device growth fails outright
+                # (NVRM out-of-memory on a GB10, three runs). A cgroup limit a
+                # little under physical memory makes the kernel reclaim this
+                # process's cache ahead of its own allocations.
+                logger.warning(
+                    "Host and device share one memory pool and this process has "
+                    "no cgroup memory limit: device allocations may fail while "
+                    "the page cache holds the free memory. Run with a limit a few "
+                    "GiB under physical memory (for example docker --memory)."
+                )
+            if swap_total > 0:
+                # Under page-cache pressure the kernel prefers swapping idle
+                # anonymous memory -- here the DiT's fused weight copies --
+                # over dropping cache, and every denoise step then swaps them
+                # back in. Measured on a GB10: 128 s first steps and a 54 s
+                # text encoder with 143 GiB of swap enabled.
+                logger.warning(
+                    "Host and device share one memory pool and swap is enabled "
+                    "(%.0f GiB): the kernel may swap out weight copies under "
+                    "page-cache pressure. Run with swap off for this process "
+                    "(container --memory-swap equal to --memory, or "
+                    "vm.swappiness=0).",
+                    swap_total / 1024**3,
+                )
+            if args.dit_cpu_offload or args.text_encoder_cpu_offload:
+                # Whole-component offload holds a component twice while it
+                # moves: the device copy plus a host copy the size of the
+                # component. On a shared pool both come out of the same
+                # memory. Measured on a GB10: a 57 GiB DiT moving back to
+                # the host at the end of a denoise stage exhausted the pool.
+                logger.warning(
+                    "Host and device share one memory pool and whole-component "
+                    "CPU offload is enabled (dit_cpu_offload=%s, "
+                    "text_encoder_cpu_offload=%s): moving a component holds it "
+                    "twice while it moves. Prefer layerwise offload, where "
+                    "residency is armed layer by layer from the checkpoint "
+                    "mapping.",
+                    bool(args.dit_cpu_offload),
+                    bool(args.text_encoder_cpu_offload),
+                )
 
     def _normalize_performance_mode(self) -> str:
         args = self.server_args

python/sglang/multimodal_gen/test/unit/test_host_spill.py

@@ -0,0 +1,104 @@
+import os
+import sys
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.runtime.loader import host_spill as host_spill_module
+from sglang.multimodal_gen.runtime.loader.host_spill import (
+    HostSpill,
+    checkpoint_fingerprint,
+)
+from sglang.multimodal_gen.runtime.loader.utils import (
+    MappedRegions,
+    hf_to_custom_state_dict,
+)
+
+
+@pytest.fixture(autouse=True)
+def _small_spill_threshold(monkeypatch):
+    monkeypatch.setattr(host_spill_module, "MIN_SPILL_BYTES", 0)
+    monkeypatch.setattr(host_spill_module, "SPILL_DISK_RESERVE_BYTES", 0)
+
+
+def test_spilled_tensor_is_file_backed_and_reused_after_sealing(tmp_path):
+    spill = HostSpill(tmp_path, "ckpt")
+    shape = torch.Size([4, 8])
+    tensor, filled = spill.tensor("blocks.0.qkv", shape, torch.bfloat16)
+    assert not filled
+    tensor.copy_(torch.arange(32, dtype=torch.bfloat16).view(4, 8))
+    if sys.platform == "linux":
+        assert MappedRegions().holds(tensor)
+    # unsealed: the next start must not trust it
+    again, filled = HostSpill(tmp_path, "ckpt").tensor(
+        "blocks.0.qkv", shape, torch.bfloat16
+    )
+    assert not filled
+    spill.seal("blocks.0.qkv", shape, torch.bfloat16)
+    reused, filled = HostSpill(tmp_path, "ckpt").tensor(
+        "blocks.0.qkv", shape, torch.bfloat16
+    )
+    assert filled
+    assert torch.equal(reused.float(), torch.arange(32, dtype=torch.float32).view(4, 8))
+    # a different dtype or shape is a different file
+    other, filled = HostSpill(tmp_path, "ckpt").tensor(
+        "blocks.0.qkv", shape, torch.float16
+    )
+    assert not filled
+    del again, other
+
+
+def test_spill_disables_itself_when_the_disk_is_full(tmp_path, monkeypatch):
+    spill = HostSpill(tmp_path, "ckpt")
+
+    class _Usage:
+        free = 0
+
+    monkeypatch.setattr(host_spill_module.shutil, "disk_usage", lambda _p: _Usage())
+    monkeypatch.setattr(host_spill_module, "SPILL_DISK_RESERVE_BYTES", 1 << 30)
+    assert spill.tensor("w", torch.Size([2, 2]), torch.float32) is None
+    assert spill.tensor("w2", torch.Size([2, 2]), torch.float32) is None
+    assert spill.count_written == 0
+
+
+@pytest.mark.parametrize("strict", [False, True])
+def test_fused_weights_are_concatenated_into_the_provided_tensor(tmp_path, strict):
+    spill = HostSpill(tmp_path, "ckpt")
+    q = torch.full((2, 3), 1.0)
+    k = torch.full((2, 3), 2.0)
+    v = torch.full((2, 3), 3.0)
+
+    def mapping(name):
+        prefix, _, which = name.rpartition(".")
+        return f"{prefix}.qkv", {"q": 0, "k": 1, "v": 2}[which], 3
+
+    weights = [("blocks.0.q", q), ("blocks.0.k", k), ("blocks.0.v", v)]
+    merged, _ = hf_to_custom_state_dict(
+        iter(weights), mapping, fused_tensor_factory=spill.tensor, strict=strict
+    )
+    fused = merged["blocks.0.qkv"]
+    assert torch.equal(fused, torch.cat([q, k, v], dim=0))
+    if sys.platform == "linux":
+        assert MappedRegions().holds(fused)
+    assert spill.count_written == 1
+    spill.seal("blocks.0.qkv", fused.shape, fused.dtype)
+
+    # the next start reuses the sealed file without reading the pieces
+    reused_spill = HostSpill(tmp_path, "ckpt")
+    merged_again, _ = hf_to_custom_state_dict(
+        iter([(n, torch.zeros_like(t)) for n, t in weights]),
+        mapping,
+        fused_tensor_factory=reused_spill.tensor,
+        strict=strict,
+    )
+    assert torch.equal(merged_again["blocks.0.qkv"], torch.cat([q, k, v], dim=0))
+    assert reused_spill.count_reused == 1
+
+
+def test_fingerprint_changes_with_the_checkpoint_files(tmp_path):
+    shard = tmp_path / "model-00001-of-00002.safetensors"
+    shard.write_bytes(b"a" * 16)
+    before = checkpoint_fingerprint([str(tmp_path)])
+    shard.write_bytes(b"b" * 32)
+    os.utime(shard, ns=(1, 1))
+    assert checkpoint_fingerprint([str(tmp_path)]) != before

python/sglang/multimodal_gen/test/unit/test_mapped_willneed.py

@@ -0,0 +1,93 @@
+"""advise_willneed hands the kernel page-aligned ranges and never raises.
+
+What matters: the madvise call receives a page-aligned start and a length
+that covers the tensor's storage, odd offsets round outward, and failures
+(no libc, bad pointers) degrade to zero advice instead of an exception.
+"""
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.runtime.managers.memory_managers import (
+    layerwise_offload as lo,
+)
+
+
+class _RecordingLibc:
+    def __init__(self, ret=0):
+        self.calls = []
+        self.ret = ret
+
+    def madvise(self, addr, length, advice):
+        self.calls.append((addr.value, length.value, advice))
+        return self.ret
+
+
+@pytest.fixture()
+def libc(monkeypatch):
+    fake = _RecordingLibc()
+    monkeypatch.setattr(lo, "_libc", fake)
+    monkeypatch.setattr(lo, "_willneed_headroom_ok", lambda need: True)
+    return fake
+
+
+def test_ranges_are_page_aligned_and_cover_the_storage(libc):
+    t = torch.zeros(1024, dtype=torch.float32)
+    advised = lo.advise_willneed([t])
+
+    assert advised == 1
+    ((addr, length, advice),) = libc.calls
+    page = lo._PAGE
+    assert advice == lo._MADV_WILLNEED
+    assert addr % page == 0
+    ptr = t.untyped_storage().data_ptr()
+    nbytes = t.untyped_storage().nbytes()
+    assert addr <= ptr
+    assert addr + length >= ptr + nbytes
+    assert length % page == 0
+
+
+def test_a_failing_madvise_counts_nothing(monkeypatch):
+    monkeypatch.setattr(lo, "_libc", _RecordingLibc(ret=-1))
+    assert lo.advise_willneed([torch.zeros(16)]) == 0
+
+
+def test_no_libc_is_a_quiet_noop(monkeypatch):
+    monkeypatch.setattr(lo, "_libc", None)
+    assert lo.advise_willneed([torch.zeros(16)]) == 0
+
+
+def test_empty_and_broken_tensors_are_skipped(libc):
+    class Broken:
+        def untyped_storage(self):
+            raise RuntimeError("no storage")
+
+    assert lo.advise_willneed([Broken(), torch.empty(0)]) == 0
+    assert libc.calls == []
+
+
+def test_no_headroom_withholds_the_advice(monkeypatch):
+    fake = _RecordingLibc()
+    monkeypatch.setattr(lo, "_libc", fake)
+    monkeypatch.setattr(lo, "_willneed_headroom_ok", lambda need: False)
+    assert lo.advise_willneed([torch.zeros(1024)]) == 0
+    assert fake.calls == []
+
+
+def test_headroom_reads_memavailable(monkeypatch, tmp_path):
+    meminfo = tmp_path / "meminfo"
+
+    real_open = open
+
+    def fake_open(path, *a, **k):
+        if path == "/proc/meminfo":
+            return real_open(meminfo, *a, **k)
+        return real_open(path, *a, **k)
+
+    monkeypatch.setattr("builtins.open", fake_open)
+
+    meminfo.write_text("MemTotal: 32 kB\nMemAvailable: 16777216 kB\n")  # 16 GiB
+    assert lo._willneed_headroom_ok(1 << 30)
+
+    meminfo.write_text("MemTotal: 32 kB\nMemAvailable: 524288 kB\n")  # 0.5 GiB
+    assert not lo._willneed_headroom_ok(1 << 30)

python/sglang/multimodal_gen/test/unit/test_readonly_safetensors.py

@@ -0,0 +1,82 @@
+import sys
+
+import pytest
+import torch
+from safetensors.torch import load_file, save_file
+
+from sglang.multimodal_gen.runtime.loader.readonly_safetensors import (
+    iter_safetensors_readonly,
+    load_safetensors_readonly,
+    safetensors_keys,
+)
+
+
+def _write(tmp_path):
+    tensors = {
+        "a.weight": torch.randn(64, 32, dtype=torch.bfloat16),
+        "b.bias": torch.arange(17, dtype=torch.float32),
+        "c.empty": torch.empty(0, 4, dtype=torch.float16),
+        "d.flag": torch.tensor([True, False]),
+        "e.int": torch.arange(6, dtype=torch.int64).view(2, 3),
+    }
+    path = tmp_path / "model.safetensors"
+    save_file(tensors, str(path))
+    return path, tensors
+
+
+def test_readonly_load_matches_safetensors(tmp_path):
+    path, tensors = _write(tmp_path)
+    ours = load_safetensors_readonly(str(path))
+    theirs = load_file(str(path))
+    assert set(ours) == set(theirs) == set(tensors)
+    for name in tensors:
+        assert ours[name].dtype == theirs[name].dtype
+        assert ours[name].shape == theirs[name].shape
+        assert torch.equal(ours[name], theirs[name])
+    assert safetensors_keys(str(path)) == list(theirs)
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="/proc/self/maps")
+def test_readonly_mapping_has_no_write_permission(tmp_path):
+    path, _ = _write(tmp_path)
+    tensor = dict(iter_safetensors_readonly(str(path)))["a.weight"]
+    ptr = tensor.data_ptr()
+    perms = None
+    for line in open("/proc/self/maps"):
+        fields = line.split()
+        low, high = (int(x, 16) for x in fields[0].split("-"))
+        if low <= ptr < high:
+            perms = fields[1]
+            break
+    assert perms is not None and perms.startswith("r--"), perms
+
+
+def test_mmap_reader_maps_read_only_where_host_copies_are_redundant(
+    tmp_path, monkeypatch
+):
+    from sglang.multimodal_gen.runtime.loader.weight_readers import safetensors_mmap
+
+    path, tensors = _write(tmp_path)
+    monkeypatch.setattr(safetensors_mmap, "host_copies_are_redundant", lambda: True)
+    reader = safetensors_mmap.SafetensorsMmapReader()
+    got = dict(
+        reader.iter_weights(
+            [str(path)],
+            device="cpu",
+            to_cpu=True,
+            key_filter=lambda name: name != "b.bias",
+            show_progress=False,
+        )
+    )
+    assert set(got) == set(tensors) - {"b.bias"}
+    assert torch.equal(got["a.weight"], tensors["a.weight"])
+    if sys.platform == "linux":
+        ptr = got["a.weight"].data_ptr()
+        perms = next(
+            line.split()[1]
+            for line in open("/proc/self/maps")
+            if int(line.split()[0].split("-")[0], 16)
+            <= ptr
+            < int(line.split()[0].split("-")[1], 16)
+        )
+        assert perms.startswith("r--"), perms

python/sglang/multimodal_gen/test/unit/test_vae_loader.py

@@ -314,7 +314,7 @@ class TestDirectGPUVAEState(unittest.TestCase):
                 "optimize_vae",
                 side_effect=lambda vae: vae,
             ),
-            patch.object(vae_loader, "safetensors_load_file") as legacy_load,
+            patch("safetensors.torch.load_file") as legacy_load,
         ):
             safetensors_save_file(
                 {"proj.weight": expected_weight, "scale": expected_scale},