[diffusion] cache LoRA-merged weights in files the page cache can hold

python/sglang/multimodal_gen/envs.py

@@ -23,6 +23,7 @@
     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_LORA_MERGE_CACHE: bool = False
     SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD: str = "fork"
     SGLANG_DIFFUSION_TARGET_DEVICE: str = "cuda"
     SGLANG_DIFFUSION_PLATFORM_OVERRIDE: str = ""
@@ -278,6 +279,11 @@ def _getter():
     "SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE": _lazy_bool(
         "SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE"
     ),
+    # 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(
+        "SGLANG_DIFFUSION_DISABLE_LORA_MERGE_CACHE"
+    ),
     # ================== cache-dit Env Vars ==================
     # Enable cache-dit acceleration for DiT inference
     # CUDA-IPC transport for 2-rank Ulysses all-to-all (NVLink same-node)

python/sglang/multimodal_gen/runtime/layers/lora/linear.py

@@ -74,14 +74,24 @@ def __init__(
         base_layer: nn.Module,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
+        snapshot_base: bool = True,
     ):
         super().__init__()
         self.base_layer: nn.Module = base_layer
 
         self.merged: bool = False
         # Immutable base-weight snapshot; `to("cpu")` may alias CPU storage.
-        # Use `clone()` so merge updates cannot mutate this backup tensor.
-        self.cpu_weight = base_layer.weight.detach().to("cpu").clone()
+        # Use `clone()` so in-place merge updates cannot mutate this backup.
+        # With snapshot_base=False the snapshot is a zero-copy view instead:
+        # valid only while every merge on this layer is a copy-merge (the
+        # merged-store path), which never writes the base storage. H3's DiT
+        # backup alone is 38 GB of anonymous memory under clone().
+        if snapshot_base:
+            self.cpu_weight = base_layer.weight.detach().to("cpu").clone()
+            self._base_is_view = False
+        else:
+            self.cpu_weight = base_layer.weight.detach()
+            self._base_is_view = True
         # indicates adapter weights don't contain this layer
         # (which shouldn't normally happen, but we want to separate it from the case of erroneous merging)
         # Default to True to prevent using uninitialized weights; set to False when weights are loaded
@@ -205,6 +215,13 @@ def set_lora_weights(
         elif self.merged:
             self.unmerge_lora_weights()
 
+    def _ensure_base_snapshot_owned(self) -> None:
+        """An in-place merge is about to write the base storage; if the
+        snapshot is a zero-copy view into it, materialize the clone now."""
+        if self._base_is_view:
+            self.cpu_weight = self.cpu_weight.clone()
+            self._base_is_view = False
+
     @torch.no_grad()
     def _merge_lora_into_data(
         self,
@@ -274,6 +291,41 @@ def _should_merge_in_fp32(
                 return False
         return True
 
+    @torch.no_grad()
+    def compute_merged_weight(self) -> torch.Tensor:
+        """The merged weight as a new CPU tensor; the base is never written.
+
+        Same math as the in-place merge — computed on the device, in fp32
+        when the policy says so, rounded back once — so the bytes are
+        identical to what merge_lora_weights would have left in place.
+        """
+        base = self.weight.data
+        target_dtype = base.dtype
+        work = base.detach().to(get_local_torch_device())
+        if (
+            self._should_merge_in_fp32(self.lora_weights_list)
+            and work.is_floating_point()
+            and work.dtype != torch.float32
+        ):
+            work = work.to(torch.float32)
+        self._merge_lora_into_data(work, self.lora_weights_list)
+        return work.to("cpu", dtype=target_dtype)
+
+    def install_merged_weight(
+        self, merged: torch.Tensor, base_view: torch.Tensor
+    ) -> None:
+        """Adopt an externally held merged weight (e.g. a cache mapping).
+
+        The single place the cached-merge state transition happens: the
+        parameter points at `merged`, the layer counts as merged, and the
+        unmerge snapshot is the untouched base view — zero-copy, because
+        nothing wrote the base storage.
+        """
+        self.weight.data = merged
+        self.merged = True
+        self.cpu_weight = base_view.detach()
+        self._base_is_view = True
+
     @torch.no_grad()
     def merge_lora_weights(self, strength: float | None = None) -> None:
         if strength is not None:
@@ -294,6 +346,7 @@ def merge_lora_weights(self, strength: float | None = None) -> None:
         if self.disable_lora:
             return
 
+        self._ensure_base_snapshot_owned()
         if self.merged:
             self.unmerge_lora_weights()
 
@@ -476,8 +529,9 @@ def __init__(
         base_layer: ColumnParallelLinear,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
+        snapshot_base: bool = True,
     ) -> None:
-        super().__init__(base_layer, lora_rank, lora_alpha)
+        super().__init__(base_layer, lora_rank, lora_alpha, snapshot_base)
 
     def forward(self, input_: torch.Tensor) -> torch.Tensor:
         if self.merged or self.disable_lora:
@@ -538,8 +592,9 @@ def __init__(
         base_layer: MergedColumnParallelLinear,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
+        snapshot_base: bool = True,
     ) -> None:
-        super().__init__(base_layer, lora_rank, lora_alpha)
+        super().__init__(base_layer, lora_rank, lora_alpha, snapshot_base)
 
     def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
         return A
@@ -574,8 +629,9 @@ def __init__(
         base_layer: QKVParallelLinear,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
+        snapshot_base: bool = True,
     ) -> None:
-        super().__init__(base_layer, lora_rank, lora_alpha)
+        super().__init__(base_layer, lora_rank, lora_alpha, snapshot_base)
 
     def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
         return A
@@ -606,8 +662,9 @@ def __init__(
         base_layer: RowParallelLinear,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
+        snapshot_base: bool = True,
     ) -> None:
-        super().__init__(base_layer, lora_rank, lora_alpha)
+        super().__init__(base_layer, lora_rank, lora_alpha, snapshot_base)
 
     def forward(self, input_: torch.Tensor):
         if self.merged or self.disable_lora:
@@ -692,8 +749,9 @@ def __init__(
         base_layer: nn.Linear,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
+        snapshot_base: bool = True,
     ) -> None:
-        super().__init__(base_layer, lora_rank, lora_alpha)
+        super().__init__(base_layer, lora_rank, lora_alpha, snapshot_base)
 
     @torch.compile()
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -728,10 +786,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             return out
 
 
+def _use_owned_base_snapshot(snapshot_base: bool, device_type: str) -> bool:
+    return snapshot_base or device_type not in ("cpu", "meta")
+
+
 def wrap_with_lora_layer(
     layer: nn.Module,
     lora_rank: int | None = None,
     lora_alpha: int | None = None,
+    snapshot_base: bool = True,
 ) -> BaseLayerWithLoRA | None:
     """
     transform the given layer to its corresponding LoRA layer
@@ -750,10 +813,14 @@ def wrap_with_lora_layer(
     }
     for src_layer_type, lora_layer_type in supported_layer_types.items():
         if isinstance(layer, src_layer_type):  # type: ignore[arg-type]
+            effective_snapshot_base = _use_owned_base_snapshot(
+                snapshot_base, layer.weight.device.type
+            )
             ret = lora_layer_type(
                 layer,
                 lora_rank=lora_rank,
                 lora_alpha=lora_alpha,
+                snapshot_base=effective_snapshot_base,
             )
             return ret
     return None

python/sglang/multimodal_gen/runtime/pipelines_core/lora/lora_merge_cache.py

@@ -0,0 +1,200 @@
+"""File-backed store for LoRA-merged weights.
+
+Merging an adapter writes the base weight in place. Under layerwise offload
+the base weight is a view into the checkpoint mapping, so the write is a
+copy-on-write: every merged byte turns into anonymous host memory the kernel
+cannot reclaim. MiniMax-H3's DiT alone is 61.7 GB — a real 32 GB host dies on
+it, and on any host the pin budget collapses to zero before the offload
+managers ever see the weights.
+
+Written once to a per-layer cache file and mapped back, the same merged bytes
+become page cache: droppable, refaultable, and invisible to the anonymous
+accounting. The offload managers then classify them as mapped weights on
+their own — no coordination needed. Rehoming happens layer by layer inside
+the merge loop, so the anonymous high-water mark stays one layer wide, and a
+later start with the same (base, adapters, strengths) adopts the store
+without paying the merge at all.
+"""
+
+import hashlib
+import json
+import os
+import shutil
+
+import torch
+from safetensors.torch import load_file as safetensors_load_file
+from safetensors.torch import save_file as safetensors_save_file
+
+from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_MANIFEST = "manifest.json"
+_DISK_HEADROOM = 1.15
+
+
+def lora_merge_cache_key(
+    base_paths: list[str],
+    adapters: list[tuple[str, float, float | None]],
+) -> str:
+    """Key of one merged-weights combination.
+
+    `base_paths` are the component checkpoint paths (HF snapshot paths carry
+    the revision hash); `adapters` are ordered (lora_path, strength, alpha)
+    triples — order matters, merges compose in order.
+    """
+    parts = [os.path.realpath(p) for p in sorted(base_paths)]
+    for path, strength, alpha in adapters:
+        real = os.path.realpath(path)
+        try:
+            size = os.path.getsize(real)
+        except OSError:
+            size = -1
+        parts.append(f"{real}|{size}|{strength}|{alpha}")
+    return hashlib.sha1("||".join(parts).encode()).hexdigest()[:16]
+
+
+class LoraMergeCache:
+    """Streams merged weights into a cache directory, one file per layer."""
+
+    def __init__(self, key: str, expected_bytes: int) -> None:
+        self.root = os.path.join(
+            envs.SGLANG_DIFFUSION_CACHE_ROOT, "lora_merge_cache", key
+        )
+        self.manifest_path = os.path.join(self.root, _MANIFEST)
+        self.expected_bytes = expected_bytes
+        self._entries: dict[str, dict] = {}
+        self._writable: bool | None = None
+
+    # -- adoption (fast path) -------------------------------------------------
+
+    def is_complete(self) -> bool:
+        """A complete store from an earlier run of the same combination."""
+        try:
+            with open(self.manifest_path) as handle:
+                manifest = json.load(handle)
+        except (OSError, ValueError):
+            return False
+        entries = manifest.get("layers")
+        if not isinstance(entries, dict) or not entries:
+            return False
+        for meta in entries.values():
+            if not os.path.exists(os.path.join(self.root, meta.get("file", ""))):
+                return False
+        self._entries = entries
+        return True
+
+    def get(
+        self, name: str, shape: torch.Size, dtype: torch.dtype
+    ) -> torch.Tensor | None:
+        """The cached merged tensor for `name`, mapped from its file.
+
+        Purely a lookup: the caller decides what to do with the tensor. A
+        missing or mismatched entry returns None — mismatch also drops the
+        remaining entries, because one wrong file means the whole combination
+        key no longer describes this module.
+        """
+        meta = self._entries.get(name)
+        if meta is None:
+            return None
+        mapped = safetensors_load_file(os.path.join(self.root, meta["file"]))
+        tensor = mapped.get("weight")
+        if (
+            tensor is None
+            or tuple(tensor.shape) != tuple(shape)
+            or tensor.dtype != dtype
+        ):
+            logger.warning(
+                "LoRA merge cache entry for %s does not match the module; "
+                "ignoring the cache",
+                name,
+            )
+            self._entries = {}
+            return None
+        return tensor
+
+    # -- capture (first run) --------------------------------------------------
+
+    def _ensure_writable(self) -> bool:
+        if self._writable is not None:
+            return self._writable
+        try:
+            os.makedirs(self.root, exist_ok=True)
+            usage = shutil.disk_usage(self.root)
+            if usage.free < self.expected_bytes * _DISK_HEADROOM:
+                logger.warning(
+                    "LoRA merge cache needs %.1f GiB free under %s but only "
+                    "%.1f GiB is available; merged weights stay in anonymous "
+                    "host memory",
+                    self.expected_bytes * _DISK_HEADROOM / 1024**3,
+                    self.root,
+                    usage.free / 1024**3,
+                )
+                self._writable = False
+            else:
+                self._writable = True
+        except OSError as exc:
+            logger.warning("LoRA merge cache unavailable (%s)", exc)
+            self._writable = False
+        return self._writable
+
+    def put(self, name: str, merged: torch.Tensor) -> torch.Tensor | None:
+        """Write one merged tensor to its cache file and return the mapping.
+
+        The returned tensor is a view into the file — page cache the kernel
+        can drop — and the only thing the cache hands back; what to install it
+        into is the caller's business. None means the bytes could not be
+        cached (disk shortage, write failure) and the caller should keep its
+        own copy.
+        """
+        if not self._ensure_writable():
+            return None
+        fname = hashlib.sha1(name.encode()).hexdigest()[:16] + ".safetensors"
+        path = os.path.join(self.root, fname)
+        try:
+            tmp = f"{path}.tmp.{os.getpid()}"
+            safetensors_save_file({"weight": merged.contiguous()}, tmp)
+            os.replace(tmp, path)
+            mapped = safetensors_load_file(path)["weight"]
+        except Exception as exc:
+            logger.warning(
+                "Could not cache merged weight %s (%s); it stays in "
+                "anonymous host memory",
+                name,
+                exc,
+            )
+            try:
+                if os.path.exists(path):
+                    os.remove(path)
+            except OSError:
+                pass
+            return None
+        self._entries[name] = {
+            "file": fname,
+            "shape": list(merged.shape),
+            "dtype": str(merged.dtype),
+        }
+        return mapped
+
+    def finalize(self, extra: dict | None = None) -> None:
+        """Write the manifest; only a complete store is ever adopted."""
+        if not self._entries or not self._ensure_writable():
+            return
+        manifest = {"layers": self._entries}
+        if extra:
+            manifest.update(extra)
+        tmp = f"{self.manifest_path}.tmp.{os.getpid()}"
+        try:
+            with open(tmp, "w") as handle:
+                json.dump(manifest, handle)
+            os.replace(tmp, self.manifest_path)
+        except OSError as exc:
+            logger.warning("LoRA merge cache manifest not written (%s)", exc)
+            return
+        logger.info(
+            "Merged weights cached to %s (%d layers); anonymous host memory "
+            "no longer holds them",
+            self.root,
+            len(self._entries),
+        )