[feat] Reuse Wan inference weights through normal loaders

docs/contributing/weight_cache.md

@@ -0,0 +1,168 @@
+# Retaining native weights across process restarts
+
+A weight-cache donor owns loaded GPU tensors. Consumers create fresh Python
+modules on `meta`, validate compatibility, and map those tensors through CUDA
+IPC. Exiting a consumer does not unload the donor. This foundation is independent
+of an edit runner and does not pickle or transfer a model object.
+
+The initial contract is Linux, one GPU, native strided tensors, inference only,
+and one active consumer pipeline at a time. The normal inference loader supports
+Wan 2.1 T2V 1.3B DiT, UMT5 text encoder and VAE together. CPU registered tensors,
+DTensor/FSDP, offload, compilation, quantization, LoRA and training are unsupported.
+These combinations raise rather than falling back to disk.
+
+## Complete Wan inference through the normal loader
+
+The donor example loads DiT, T5 and VAE once through `VideoGenerator`. One process
+owns all three components and serves one socket per component. Tokenizers and
+schedulers remain local; each consumer builds new model objects and request state.
+
+Start the donor on the same GPU, user and IPC namespace as the consumers:
+
+```bash
+python examples/benchmark/wan_weight_cache.py donor \
+  --model-path /path/to/Wan2.1-T2V-1.3B-Diffusers \
+  --checkpoint-id IMMUTABLE_REVISION --cache-dir /tmp/fv-wan-weights \
+  --output /tmp/wan-donor.json
+```
+
+Wait for `DiT, T5 and VAE donor sockets ready`, then run independent consumers:
+
+```bash
+python examples/benchmark/wan_weight_cache.py consumer \
+  --model-path /path/to/Wan2.1-T2V-1.3B-Diffusers \
+  --checkpoint-id IMMUTABLE_REVISION --cache-dir /tmp/fv-wan-weights \
+  --output /tmp/wan-consumer-1.json --tensor-output /tmp/wan-consumer-1.pt
+```
+
+Repeat with new output filenames. Use role `disk` for fresh-loading parity checks.
+Each worker runs two identical seeded 9-frame, 128x128, four-step requests. The
+in-process executor records full component fingerprints and IPC-handle witnesses.
+`--executor mp` also exercises the default executor selection; FastVideo resolves
+single-GPU `mp` to `UniprocExecutor`, which the report records explicitly.
+Stop the donor with Ctrl-C after consumers finish; it waits for active leases.
+
+For application code, use the ordinary public inference entry point:
+
+```python
+generator = VideoGenerator.from_pretrained(
+    local_checkpoint, num_gpus=1, tp_size=1, sp_size=1,
+    weight_cache_dir="/tmp/fv-wan-weights",
+    weight_cache_checkpoint="IMMUTABLE_REVISION",
+    dit_cpu_offload=False, dit_layerwise_offload=False,
+    text_encoder_cpu_offload=False, vae_cpu_offload=False,
+    lazy_module_load=False, inference_torch_compile=False,
+)
+try:
+    result = generator.generate_video(prompt, save_video=False)
+finally:
+    generator.shutdown()
+```
+
+The same fields have CLI options `--weight-cache-dir` and
+`--weight-cache-checkpoint`. A checkpoint identity without a directory keeps disk
+loading and records the donor keys. Supplying the directory makes attachment
+mandatory for all three components. There is no silent hydration fallback.
+
+`PipelineComponentLoader` retains normal config resolution and class dispatch.
+Its native loader adapters attach after configuration but before checkpoint
+hydration. Identity includes the resolved architecture, config constants, dtype,
+attention backend request and loader processing version. The caller must provide
+an immutable checkpoint revision or content identity; the cache does not hash
+all checkpoint files during attachment. Recreate the donor when the checkpoint,
+layout, or weight-processing implementation changes.
+
+The pipeline owns an `ExitStack` of attachment leases. Pipeline close and worker
+shutdown synchronize and detach all components; partial load failure releases
+already attached components. Direct `PipelineComponentLoader` users must supply
+`weight_cache_leases=stack` and keep that stack open throughout inference. Runtime
+LoRA merge/change APIs reject shared weights. Arbitrary user tensor mutations
+remain outside the read-only contract and cannot be prevented by CUDA IPC.
+
+The donor reserves the combined weight memory for its lifetime. Attached workers
+still need their own CUDA contexts, communication state and activation memory.
+Keep the donor and sequential workers inside the same GPU allocation; this does
+not retain GPU memory across allocation teardown or machine restart.
+
+## Native Wan VAE example
+
+Use the same GPU, user and IPC namespace for all processes. The model directory
+is a local Wan 2.1 1.3B Diffusers VAE component. Supply its immutable revision.
+Create a private socket directory:
+
+```bash
+install -d -m 700 /tmp/fv-weights
+python examples/benchmark/vae_weight_cache.py donor \
+  --model-path /path/to/checkpoint/vae --checkpoint-id IMMUTABLE_REVISION \
+  --socket /tmp/fv-weights/vae.sock --output /tmp/donor.json
+```
+
+After the socket exists, run consumers from another terminal. Each invocation is
+an independent process:
+
+```bash
+python examples/benchmark/vae_weight_cache.py consumer \
+  --model-path /path/to/checkpoint/vae --checkpoint-id IMMUTABLE_REVISION \
+  --socket /tmp/fv-weights/vae.sock --output /tmp/consumer-1.json \
+  --tensor-output /tmp/consumer-1.pt
+```
+
+Repeat with new output filenames. The same donor retains the tensors. The example
+reads configuration in each consumer, constructs fresh VAE caches and module
+objects, and attaches registered weights without checkpoint hydration. Its FP32
+codec check decodes identical input twice, checks cache reset and weight
+fingerprints, and can export output for an independent `disk`-role comparison.
+Use Ctrl-C on the donor after consumers finish; graceful shutdown waits for the
+active connection to close and writes a final weight fingerprint.
+
+## Python boundary
+
+The donor receives an already loaded, idle inference model:
+
+```python
+serve_weight_cache(model.eval(), key, socket_path, stop=stop_event)
+```
+
+Construction and attachment happen in the consumer:
+
+```python
+client = WeightCacheClient(socket_path, key, device)
+with client.load_model(factory) as attached:
+    output = attached.model(inputs)
+```
+
+`factory` constructs the component without reading checkpoint weights. It runs
+inside a meta-device context. Initialize required non-registered configuration
+constants outside that context. Model/request state belongs to the consumer;
+only registered parameters and buffers are shared. Non-persistent buffers and
+tied-parameter identities are included. Consumer bindings are restored to meta
+on context exit, after GPU synchronization. Do not retain mapped tensors outside
+the context.
+
+The donor must outlive every mapped tensor. Shared tensors have a read-only
+inference contract; CUDA does not enforce read-only mappings. In-place weight
+changes require private copies or a new donor. Donor crashes invalidate consumers
+and require recovery. Imports, CUDA contexts, compilation, graphs and request
+caches are not transferred or retained by this API.
+
+## Compatibility and transport
+
+`WeightCacheKey` requires checkpoint, component, configuration, dtype and
+weight-processing identities. The caller must provide the correct immutable
+checkpoint identity and increment the processing version when weight preparation
+changes. A method-body edit can reuse the same weights. Both peers also check the
+Torch version and physical GPU UUID. The complete registered layout checks class,
+names, shapes, dtypes, strides, offsets, aliases and buffer persistence before
+any bindings change. A mismatch or missing donor raises; no disk fallback exists.
+
+The Unix socket and its directory must be private to the current user. Each peer
+checks Linux peer credentials before decoding trusted same-user descriptors.
+Descriptors contain CUDA IPC handles, not GPU weight bytes. New PyTorch reduction
+arguments are created for every consumer to preserve per-mapping refcounts. A
+per-socket lock prevents a second donor from replacing an active donor's socket.
+
+The design follows the separation used by vLLM's
+[weight-cache donor](https://github.com/vllm-project/vllm/blob/b965ee564d1897fdaeddaeceecc316aa15a45312/vllm/model_executor/model_loader/weight_cache/daemon.py)
+and [IPC loader](https://github.com/vllm-project/vllm/blob/b965ee564d1897fdaeddaeceecc316aa15a45312/vllm/model_executor/model_loader/weight_cache/ipc_loader.py).
+vLLM additionally handles post-processing, quantization and parallel shards.
+Those capabilities are not implied by this initial native-tensor foundation.

examples/benchmark/vae_weight_cache.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Persistent native Wan VAE weights with independently restartable IPC consumers.
+
+Run donor in one process; run consumer repeatedly in fresh processes on the same
+GPU. The donor must stay alive. The example uses a local, immutable Wan 2.1 1.3B
+Diffusers VAE component and FP32. It shares weights, not mutable VAE caches.
+"""
+import argparse
+import hashlib
+import json
+from pathlib import Path
+import signal
+import threading
+import time
+
+import torch
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+import fastvideo.envs as envs
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.models.loader.component_loader import VAELoader
+from fastvideo.models.loader.utils import set_default_torch_dtype
+from fastvideo.models.loader.weight_cache import WeightCacheClient, WeightCacheKey, serve_weight_cache
+from fastvideo.models.wan.vae import AutoencoderKLWan
+from fastvideo.models.wan.vae_config import WanVAEConfig
+from fastvideo.utils import get_open_port
+
+
+def fingerprint(model):
+    digest = hashlib.sha256()
+    for name, tensor in [*model.named_parameters(remove_duplicate=False), *model.named_buffers(remove_duplicate=False)]:
+        digest.update(name.encode())
+        digest.update(str((tensor.dtype, tuple(tensor.shape), tuple(tensor.stride()), tensor.storage_offset())).encode())
+        digest.update(tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu().numpy().tobytes())
+    return digest.hexdigest()
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("role", choices=["donor", "consumer", "disk"])
+    parser.add_argument("--model-path", type=Path, required=True)
+    parser.add_argument("--checkpoint-id", required=True, help="Immutable checkpoint revision or content digest")
+    parser.add_argument("--socket", type=Path)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--tensor-output", type=Path)
+    args = parser.parse_args()
+    if args.role != "disk" and args.socket is None:
+        parser.error("--socket is required for donor/consumer")
+    config_bytes = (args.model_path / "config.json").read_bytes()
+    architecture = json.loads(config_bytes)
+    if architecture.get("_class_name") != "AutoencoderKLWan":
+        raise ValueError("This adapter supports native Wan VAE only")
+    key = WeightCacheKey(args.checkpoint_id, "vae", hashlib.sha256(config_bytes).hexdigest(),
+                         "torch.float32", "native-wan-fp32-v1")
+    envs.setdefault_external("MASTER_ADDR", "127.0.0.1")
+    envs.setdefault_external("MASTER_PORT", str(get_open_port()))
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+    device = torch.device("cuda", torch.cuda.current_device())
+    report = {"role": args.role, "checkpoint": args.checkpoint_id, "torch": torch.__version__}
+    def load_disk():
+        runtime = FastVideoArgs(model_path=str(args.model_path), vae_cpu_offload=False,
+                                pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision="fp32"))
+        return VAELoader().load(str(args.model_path), runtime).eval()
+    native_config = WanVAEConfig()
+    native_config.update_model_arch(architecture)
+    def construct():
+        return VAELoader().build_model(AutoencoderKLWan, native_config, torch.device("meta"))
+    def verify(model):
+        report["weight_sha256_before"] = fingerprint(model)
+        x = torch.randn((1, 16, 2, 12, 12), generator=torch.Generator().manual_seed(1024)).to(device)
+        with torch.no_grad():
+            first = model.decode(x)
+            second = model.decode(x)
+        torch.cuda.synchronize(device)
+        assert first.shape == (1, 3, 5, 96, 96) and torch.isfinite(first).all()
+        assert torch.equal(first, second), "Request-local VAE caches were not reset"
+        report["weight_sha256_after"] = fingerprint(model)
+        assert report["weight_sha256_before"] == report["weight_sha256_after"]
+        if args.tensor_output:
+            args.tensor_output.parent.mkdir(parents=True, exist_ok=True)
+            torch.save(first.cpu(), args.tensor_output)
+    try:
+        if args.role == "donor":
+            model = load_disk()
+            report["weight_sha256_before"] = fingerprint(model)
+            stop = threading.Event()
+            for sig in (signal.SIGINT, signal.SIGTERM):
+                signal.signal(sig, lambda *_: stop.set())
+            print("Serving donor weights", flush=True)
+            serve_weight_cache(model, key, args.socket, stop=stop)
+            report["weight_sha256_after"] = fingerprint(model)
+            assert report["weight_sha256_before"] == report["weight_sha256_after"]
+        else:
+            start = time.perf_counter()
+            allocated = torch.cuda.memory_allocated(device)
+            if args.role == "consumer":
+                with set_default_torch_dtype(torch.float32):
+                    with WeightCacheClient(args.socket, key, device).load_model(construct) as attachment:
+                        report["load_seconds"] = time.perf_counter() - start
+                        report["allocated_weight_delta_bytes"] = torch.cuda.memory_allocated(device) - allocated
+                        report["donor_pid"] = attachment.donor_pid
+                        report["ipc_handles"] = attachment.handles
+                        report["weight_source"] = "cuda_ipc"
+                        verify(attachment.model)
+            else:
+                model = load_disk()
+                torch.cuda.synchronize(device)
+                report["load_seconds"] = time.perf_counter() - start
+                report["allocated_weight_delta_bytes"] = torch.cuda.memory_allocated(device) - allocated
+                report["weight_source"] = "disk"
+                verify(model)
+        report["status"] = "ok"
+    except Exception as error:
+        report.update(status="error", error=repr(error))
+        raise
+    finally:
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(json.dumps(report, indent=2) + "\n")
+        cleanup_dist_env_and_memory()
+
+
+if __name__ == "__main__":
+    main()

examples/benchmark/wan_weight_cache.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Retain Wan DiT/T5/VAE weights; verify fresh normal-loader inference workers.
+
+Use a local Wan 2.1 T2V 1.3B Diffusers checkpoint on one GPU. Start a donor,
+then run consumers in separate processes with the same private cache directory.
+The tiny request checks execution/parity, not video quality or throughput.
+"""
+import argparse
+from dataclasses import asdict
+import hashlib
+import json
+from pathlib import Path
+import signal
+import threading
+import time
+
+import torch
+
+from fastvideo import VideoGenerator
+from fastvideo.models.loader.inference_weight_cache import COMPONENTS
+from fastvideo.models.loader.weight_cache import serve_weight_cache
+
+
+def fingerprints(modules):
+    result = {}
+    for component in COMPONENTS:
+        digest = hashlib.sha256()
+        model = modules[component]
+        for name, tensor in [*model.named_parameters(remove_duplicate=False),
+                             *model.named_buffers(remove_duplicate=False)]:
+            digest.update(name.encode())
+            digest.update(str((tuple(tensor.shape), tensor.dtype, tuple(tensor.stride()))).encode())
+            digest.update(tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu().numpy().tobytes())
+        result[component] = digest.hexdigest()
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("role", choices=("donor", "consumer", "disk"))
+    parser.add_argument("--model-path", required=True)
+    parser.add_argument("--checkpoint-id", required=True)
+    parser.add_argument("--cache-dir", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--tensor-output", type=Path)
+    parser.add_argument("--executor", choices=("uni", "mp"), default="uni")
+    args = parser.parse_args()
+    report = {"role": args.role, "torch": torch.__version__, "executor": args.executor}
+    generator = None
+    stop = threading.Event()
+    threads = []
+    failures = []
+    started = time.perf_counter()
+    try:
+        # Import/runtime allocations are separate from retained weight storage.
+        torch.cuda.init()
+        allocated = torch.cuda.memory_allocated()
+        generator = VideoGenerator.from_pretrained(
+            args.model_path, num_gpus=1, distributed_executor_backend=args.executor,
+            tp_size=1, sp_size=1, hsdp_replicate_dim=1, hsdp_shard_dim=1,
+            dit_cpu_offload=False, dit_layerwise_offload=False, vae_cpu_offload=False,
+            text_encoder_cpu_offload=False, image_encoder_cpu_offload=False,
+            lazy_module_load=False, inference_torch_compile=False,
+            vae_tiling=False, vae_sp=False,
+            weight_cache_dir=str(args.cache_dir) if args.role == "consumer" else None,
+            weight_cache_checkpoint=args.checkpoint_id,
+        )
+        torch.cuda.synchronize()
+        report["ready_seconds"] = time.perf_counter() - started
+        report["pipeline_allocation_delta_bytes"] = torch.cuda.memory_allocated() - allocated
+        # FastVideo resolves both uni and single-GPU mp to UniprocExecutor.
+        # Record the effective executor instead of assuming the requested name
+        # implies a spawned worker.
+        report["effective_executor"] = type(generator.executor).__name__
+        modules = None
+        if report["effective_executor"] == "UniprocExecutor":
+            pipeline = generator.executor.driver_worker.worker.pipeline
+            modules = pipeline.modules
+            report["keys"] = {name: asdict(modules[name]._fastvideo_weight_cache_key) for name in COMPONENTS}
+            report["attachments"] = {name: modules[name]._fastvideo_weight_cache_witness for name in COMPONENTS}
+            report["fingerprints_before"] = fingerprints(modules)
+        if args.role == "donor":
+            if modules is None:
+                raise ValueError("The donor must use the uni executor")
+            args.cache_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
+            for sig in (signal.SIGINT, signal.SIGTERM):
+                signal.signal(sig, lambda *_: stop.set())
+            def serve(component):
+                try:
+                    torch.cuda.set_device(0)
+                    serve_weight_cache(modules[component], modules[component]._fastvideo_weight_cache_key,
+                                       args.cache_dir / f"{component}.sock", stop=stop)
+                except BaseException as error:
+                    failures.append(error)
+                    stop.set()
+            for component in COMPONENTS:
+                thread = threading.Thread(target=serve, args=(component,))
+                thread.start()
+                threads.append(thread)
+            deadline = time.monotonic() + 30
+            while not all((args.cache_dir / f"{name}.sock").exists() for name in COMPONENTS):
+                if stop.wait(.1) or time.monotonic() > deadline:
+                    raise RuntimeError(f"Donor failed to become ready: {failures}")
+            report["donor_ready_seconds"] = time.perf_counter() - started
+            print("DiT, T5 and VAE donor sockets ready", flush=True)
+            while not stop.wait(.5):
+                pass
+            if failures:
+                raise failures[0]
+        else:
+            outputs = []
+            for _ in range(2):
+                result = generator.generate_video(
+                    "A red cube on a wooden table.", num_frames=9, height=128, width=128,
+                    num_inference_steps=4, seed=1024, save_video=False, return_frames=True,
+                )
+                samples = result["samples"].detach().cpu()
+                assert samples.shape == (1, 3, 9, 128, 128) and torch.isfinite(samples).all()
+                outputs.append(samples)
+            assert torch.equal(outputs[0], outputs[1]), "Repeated seeded requests differ"
+            report["output_shape"] = list(outputs[0].shape)
+            if args.tensor_output:
+                torch.save(outputs[0], args.tensor_output)
+        # Close servers before the final fingerprint so no consumer can still
+        # execute on these tensors during the donor's final verification.
+        stop.set()
+        for thread in threads:
+            thread.join()
+        if modules is not None:
+            report["fingerprints_after"] = fingerprints(modules)
+            assert report["fingerprints_before"] == report["fingerprints_after"]
+        report["status"] = "ok"
+    except BaseException as error:
+        report["status"] = "error"
+        report["error"] = repr(error)
+        raise
+    finally:
+        stop.set()
+        for thread in threads:
+            thread.join()
+        if generator is not None:
+            generator.executor.shutdown()
+        report["total_seconds"] = time.perf_counter() - started
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(json.dumps(report, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/fastvideo_args.py

@@ -162,6 +162,11 @@ class FastVideoArgs:
     trust_remote_code: bool = False
     revision: str | None = None
 
+    # Persistent native Wan inference weights. Directory selects IPC attachment;
+    # identity alone records donor keys during ordinary disk loading.
+    weight_cache_dir: str | None = None
+    weight_cache_checkpoint: str | None = None
+
     # Parallelism
     num_gpus: int = 1
     tp_size: int = -1
@@ -569,6 +574,15 @@ class FastVideoArgs:
             help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
         )
 
+        parser.add_argument("--weight-cache-dir",
+                            type=str,
+                            default=None,
+                            help="Private directory containing native Wan inference weight-donor sockets")
+        parser.add_argument("--weight-cache-checkpoint",
+                            type=str,
+                            default=None,
+                            help="Immutable checkpoint revision/content identity for retained weights")
+
         # Running mode
         parser.add_argument(
             "--mode",

fastvideo/models/loader/component_loader.py

@@ -7,7 +7,7 @@ import os
 import time
 from abc import ABC, abstractmethod
 from collections.abc import Generator, Iterable
-from contextlib import nullcontext
+from contextlib import ExitStack, nullcontext
 from copy import deepcopy
 from typing import Any, cast
 
@@ -39,6 +39,8 @@ from fastvideo.models.loader.text_encoder_quantization import (
     _process_quantized_text_encoder_weights,
     _resolve_text_encoder_checkpoint_path,
 )
+from fastvideo.models.loader.inference_weight_cache import (attach_cached_model, component_weight_cache,
+                                                           validate_cached_class)
 from fastvideo.models.loader.utils import set_default_torch_dtype
 from fastvideo.models.loader.weight_utils import (
     filter_duplicate_safetensors_files,
@@ -438,6 +440,11 @@ class TextEncoderLoader(ComponentLoader):
                 quant_cls = get_quantization_config(fastvideo_args.override_text_encoder_quant)
                 model_config.quant_config = quant_cls()
 
+            cached = attach_cached_model(model_cls, model_config, PRECISION_TO_TYPE[dtype],
+                                         lambda: model_cls(model_config))
+            if cached is not None:
+                return cached
+
             if getattr(model_cls, "supports_hf_from_pretrained", False):
                 model = model_cls.from_pretrained_local(  # type: ignore[attr-defined]
                     model_path,
@@ -836,6 +843,7 @@ class VAELoader(ComponentLoader):
             class_name = getattr(fastvideo_args, "_model_index_class_names", {}).get("vae")
         assert class_name is not None, (
             "Model config does not contain a _class_name attribute. Only diffusers format is supported.")
+        validate_cached_class(class_name)
         fastvideo_args.model_paths["vae"] = model_path
 
         from fastvideo.platforms import current_platform
@@ -934,7 +942,7 @@ class VAELoader(ComponentLoader):
             # LTX-2 uses CausalVideoAutoencoder with nested "vae" config
             if class_name == "CausalVideoAutoencoder" and "vae" in config:
                 vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-                vae = vae_cls(config).to(target_device)
+                vae = self.build_model(vae_cls, config, target_device)
                 if hasattr(vae, "set_tiling_config"):
                     vae_config = fastvideo_args.pipeline_config.vae_config
                     vae.set_tiling_config(
@@ -948,8 +956,33 @@ class VAELoader(ComponentLoader):
                 vae_config = fastvideo_args.pipeline_config.vae_config
                 vae_config.update_model_arch(config)
                 vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-                vae = vae_cls(vae_config).to(target_device)
+                cached = attach_cached_model(
+                    vae_cls, vae_config, torch.get_default_dtype(),
+                    lambda: self.build_model(vae_cls, vae_config, torch.device("meta")))
+                if cached is not None:
+                    return cached
+                vae = self.build_model(vae_cls, vae_config, target_device)
 
+        return self.load_weights(vae, model_path, class_name=class_name, config=config, target_device=target_device)
+
+    def build_model(self, vae_cls: type[nn.Module], config: Any, target_device: torch.device) -> nn.Module:
+        """Construct a native VAE without checkpoint weights; the caller owns the dtype context.
+
+        Checkpoint-backed special constructors in load() retain their existing
+        behavior. This boundary covers the native safetensors path, including Wan.
+        """
+        vae = vae_cls(config)
+        vae = vae.to(target_device)
+        return vae
+
+    def load_weights(self, vae: nn.Module, model_path: str, *, class_name: str,
+                     config: dict[str, Any], target_device: torch.device) -> nn.Module:
+        """Hydrate a constructed native VAE from disk and finalize it for inference.
+
+        Keep checkpoint remapping, strictness, overlays and memory-format changes
+        here so construction can be measured/reused independently. This remains a
+        copying disk loader; it does not attach IPC storage.
+        """
         # Find all safetensors files
         safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
         int8_convrot_path = None
@@ -1209,6 +1242,17 @@ class TransformerLoader(ComponentLoader):
 
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
 
+        default_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+        def construct_cached_transformer():
+            with set_default_torch_dtype(default_dtype):
+                return model_cls(config=dit_config, hf_config=hf_config)
+        record_resolved_attention_backend(dit_config)
+        cached = attach_cached_model(model_cls, dit_config, default_dtype, construct_cached_transformer,
+                                     extra_config=hf_config)
+        if cached is not None:
+            fastvideo_args.pipeline_config.dit_config = dit_config
+            return cached
+
         # Find all safetensors files
         safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
         if not safetensors_list:
@@ -1588,6 +1632,7 @@ class PipelineComponentLoader:
         component_model_path: str,
         transformers_or_diffusers: str,
         fastvideo_args: FastVideoArgs,
+        weight_cache_leases: ExitStack | None = None,
     ):
         """
         Load a pipeline module.
@@ -1629,4 +1674,9 @@ class PipelineComponentLoader:
                                  if resolved is not None else nullcontext())
 
         with attention_context:
-            return loader.load(component_model_path, fastvideo_args)
+            with component_weight_cache(module_name, fastvideo_args, weight_cache_leases) as cache:
+                model = loader.load(component_model_path, fastvideo_args)
+                if cache is not None:
+                    model._fastvideo_weight_cache_key = cache.key
+                    model._fastvideo_weight_cache_witness = cache.witness
+                return model

fastvideo/models/loader/inference_weight_cache.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Opt-in native Wan loader adapters for persistent CUDA IPC weights."""
+from collections.abc import Callable
+from contextlib import ExitStack, contextmanager
+from contextvars import ContextVar
+from dataclasses import asdict, dataclass, is_dataclass
+from enum import Enum
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+
+from fastvideo.models.loader.weight_cache import WeightCacheClient, WeightCacheKey
+
+COMPONENTS = ("transformer", "text_encoder", "vae")
+_CLASSES = {"transformer": "WanTransformer3DModel", "text_encoder": "UMT5EncoderModel",
+            "vae": "AutoencoderKLWan"}
+
+
+def validate_weight_cache_args(args) -> None:
+    """Reject modes that can mutate, relocate or reinterpret shared storage."""
+    if not args.inference_mode or str(args.mode) not in ("inference", "ExecutionMode.INFERENCE"):
+        raise ValueError("Weight reuse supports inference only")
+    if type(args.pipeline_config).__name__ != "WanT2V480PConfig":
+        raise ValueError("Weight reuse currently supports the native Wan 2.1 T2V 1.3B pipeline only")
+    if any(getattr(args, name) != 1 for name in ("num_gpus", "tp_size", "sp_size",
+                                                "hsdp_replicate_dim", "hsdp_shard_dim")):
+        raise ValueError("Weight reuse requires one GPU and single-rank parallelism")
+    unsupported = ("use_fsdp_inference", "dit_cpu_offload", "dit_layerwise_offload", "vae_cpu_offload",
+                   "text_encoder_cpu_offload", "lazy_module_load", "enable_torch_compile",
+                   "enable_torch_compile_text_encoder", "enable_torch_compile_vae", "inference_torch_compile",
+                   "lora_path", "init_weights_from_safetensors", "init_weights_from_safetensors_2",
+                   "override_text_encoder_safetensors", "override_text_encoder_quant",
+                   "override_transformer_cls_name", "override_transformer_safetensors")
+    enabled = [name for name in unsupported if getattr(args, name, None)]
+    if enabled:
+        raise ValueError(f"Weight reuse does not support: {', '.join(enabled)}")
+    configs = [args.pipeline_config.dit_config, *args.pipeline_config.text_encoder_configs]
+    if any(getattr(config, "quant_config", None) is not None for config in configs):
+        raise ValueError("Weight reuse does not support quantized components")
+    if args.pipeline_config.vae_config.use_light_vae:
+        raise ValueError("Weight reuse does not support light-VAE memory-format conversion")
+    if not args.weight_cache_checkpoint:
+        raise ValueError("Weight reuse requires an immutable weight_cache_checkpoint identity")
+
+
+def _json_default(value):
+    if is_dataclass(value):
+        return asdict(value)
+    if isinstance(value, torch.Tensor):
+        if value.is_meta:
+            raise ValueError("Configuration constants must be initialized outside meta construction")
+        return {"dtype": str(value.dtype), "shape": list(value.shape), "values": value.cpu().tolist()}
+    if isinstance(value, (set, frozenset)):
+        return sorted(value)
+    if isinstance(value, Enum):
+        return value.value
+    if isinstance(value, (torch.dtype, torch.device, Path)):
+        return str(value)
+    if callable(value):
+        return value.__module__ + "." + value.__qualname__
+    raise TypeError(f"Unsupported weight configuration identity: {type(value)}")
+
+
+@dataclass
+class _LoadContext:
+    component: str
+    args: Any
+    leases: ExitStack | None
+    key: WeightCacheKey | None = None
+    witness: dict | None = None
+
+
+_CURRENT: ContextVar[_LoadContext | None] = ContextVar("inference_weight_cache_load", default=None)
+
+
+@contextmanager
+def component_weight_cache(component, args, leases):
+    """The normal loader owns dispatch; the pipeline owns attachment lifetimes.
+
+    A checkpoint identity without a cache directory records donor keys while
+    retaining normal disk hydration. With a directory, attachment is mandatory.
+    """
+    if not (getattr(args, "weight_cache_dir", None) or getattr(args, "weight_cache_checkpoint", None)):
+        yield None
+        return
+    validate_weight_cache_args(args)
+    if component not in (*COMPONENTS, "scheduler", "tokenizer"):
+        raise ValueError(f"Unsupported weight-cache pipeline component: {component}")
+    if component not in COMPONENTS:
+        yield None
+        return
+    if args.weight_cache_dir and leases is None:
+        raise ValueError("IPC loading requires a caller-owned weight_cache_leases ExitStack")
+    context = _LoadContext(component, args, leases)
+    token = _CURRENT.set(context)
+    try:
+        yield context
+        if context.key is None:
+            raise ValueError(f"No native weight-cache adapter for {component}")
+    finally:
+        _CURRENT.reset(token)
+
+
+def validate_cached_class(class_name):
+    context = _CURRENT.get()
+    if context is not None and class_name != _CLASSES[context.component]:
+        raise ValueError(f"Unsupported cached {context.component} class: {class_name}")
+
+
+def attach_cached_model(model_cls, config, dtype, factory: Callable[[], nn.Module], *, extra_config=None):
+    """Called after config resolution and before any checkpoint hydration."""
+    context = _CURRENT.get()
+    if context is None:
+        return None
+    validate_cached_class(model_cls.__name__)
+    if getattr(config, "quant_config", None) is not None:
+        raise ValueError("Weight reuse does not support checkpoint-selected quantization")
+    identity = {"class": model_cls.__module__ + "." + model_cls.__qualname__,
+                "config": config, "extra": extra_config}
+    encoded = json.dumps(identity, default=_json_default, sort_keys=True).encode()
+    context.key = WeightCacheKey(context.args.weight_cache_checkpoint, context.component,
+                                hashlib.sha256(encoded).hexdigest(), str(dtype), "native-wan-loader-v1")
+    if not context.args.weight_cache_dir:
+        return None
+    client = WeightCacheClient(Path(context.args.weight_cache_dir) / f"{context.component}.sock", context.key)
+    attached = context.leases.enter_context(client.load_model(factory))
+    context.witness = {"donor_pid": attached.donor_pid, "handles": attached.handles}
+    return attached.model

fastvideo/models/loader/weight_cache.py

@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Single-GPU, same-user CUDA IPC weights for restartable inference consumers.
+
+The donor owns registered tensors; each consumer owns new Python/module state.
+The donor must outlive every mapped tensor. This prototype serves one active
+consumer at a time and never falls back to checkpoint loading.
+"""
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass, fields
+import contextlib
+import hashlib
+import os
+from pathlib import Path
+import pickle
+import socket
+import stat
+import struct
+import threading
+from typing import Any
+
+import torch
+from torch import nn
+from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor
+
+_VERSION = 1
+_MAX_MESSAGE = 64 * 1024 * 1024
+
+
+@dataclass(frozen=True)
+class WeightCacheKey:
+    """Caller supplies immutable checkpoint/config and weight-processing identities.
+
+    Method-body edits may reuse weights; edits affecting names/layout/preparation
+    need a different processing_version. Tensor compatibility is also checked.
+    """
+    checkpoint: str
+    component: str
+    config_hash: str
+    dtype: str
+    processing_version: str
+
+    def __post_init__(self):
+        if any(not isinstance(getattr(self, item.name), str) or not getattr(self, item.name)
+               for item in fields(self)):
+            raise ValueError("Every weight-cache identity field must be a nonempty string")
+
+
+@dataclass
+class WeightAttachment:
+    model: nn.Module
+    donor_pid: int
+    handles: dict[str, str]
+
+
+def _send(conn: socket.socket, payload: Any) -> None:
+    data = pickle.dumps(payload, protocol=5)
+    if len(data) > _MAX_MESSAGE:
+        raise ValueError("Weight-cache descriptor exceeds message limit")
+    conn.sendall(struct.pack("!I", len(data)) + data)
+
+
+def _read(conn: socket.socket, length: int) -> bytes:
+    data = bytearray()
+    while len(data) < length:
+        chunk = conn.recv(length - len(data))
+        if not chunk:
+            raise ConnectionError("Weight-cache peer disconnected")
+        data.extend(chunk)
+    return bytes(data)
+
+
+def _recv(conn: socket.socket) -> Any:
+    length, = struct.unpack("!I", _read(conn, 4))
+    if length > _MAX_MESSAGE:
+        raise ValueError("Weight-cache descriptor exceeds message limit")
+    return pickle.loads(_read(conn, length))
+
+
+def _same_user(conn: socket.socket) -> None:
+    # Descriptors contain pickle objects: authenticate the local peer before
+    # reading any payload. CUDA IPC is currently a Linux-only contract here.
+    _, uid, _ = struct.unpack("3i", conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))
+    if uid != os.getuid():
+        raise PermissionError("Weight-cache peer is not the socket owner")
+
+
+def _inventory(model: nn.Module) -> tuple[dict[str, torch.Tensor], dict[str, dict], dict[str, str]]:
+    tensors, specs, aliases, canonical = {}, {}, {}, {}
+    params = list(model.named_parameters(remove_duplicate=False))
+    buffers = list(model.named_buffers(remove_duplicate=False))
+    modules = dict(model.named_modules(remove_duplicate=False))
+    for kind, values in (("parameter", params), ("buffer", buffers)):
+        for name, tensor in values:
+            if type(tensor) not in (torch.Tensor, nn.Parameter) or tensor.layout != torch.strided:
+                raise ValueError(f"Only native strided tensors are supported: {name}")
+            if id(tensor) in canonical:
+                aliases[name] = canonical[id(tensor)]
+            else:
+                canonical[id(tensor)] = name
+                tensors[name] = tensor
+            parent, _, leaf = name.rpartition(".")
+            specs[name] = {"kind": kind, "shape": tuple(tensor.shape), "dtype": str(tensor.dtype),
+                           "stride": tuple(tensor.stride()), "offset": tensor.storage_offset(),
+                           "persistent": kind == "parameter" or leaf not in modules[parent]._non_persistent_buffers_set}
+    return tensors, specs, aliases
+
+
+def _manifest(model: nn.Module) -> dict:
+    _, specs, aliases = _inventory(model)
+    return {"model_class": type(model).__module__ + "." + type(model).__qualname__,
+            "specs": specs, "aliases": aliases}
+
+
+def _validate_manifest(model: nn.Module, manifest: dict) -> None:
+    if _manifest(model) != manifest:
+        raise ValueError("Cached model names, shapes, dtypes, strides, aliases or class do not match")
+
+
+def _bind(model: nn.Module, tensors: dict[str, torch.Tensor], manifest: dict) -> None:
+    """Validate everything before replacing any meta registration."""
+    _validate_manifest(model, manifest)
+    expected = set(manifest["specs"]) - set(manifest["aliases"])
+    if set(tensors) != expected:
+        raise ValueError("Incomplete cached tensor set")
+    for name, tensor in tensors.items():
+        spec = manifest["specs"][name]
+        if (tuple(tensor.shape), str(tensor.dtype), tuple(tensor.stride()), tensor.storage_offset()) != (
+                spec["shape"], spec["dtype"], spec["stride"], spec["offset"]):
+            raise ValueError(f"Cached tensor layout mismatch: {name}")
+    modules = dict(model.named_modules(remove_duplicate=False))
+    registered = {}
+    for name in [*tensors, *manifest["aliases"]]:
+        source = manifest["aliases"].get(name)
+        spec = manifest["specs"][name]
+        obj = registered[source] if source is not None else tensors[name]
+        if source is None and spec["kind"] == "parameter":
+            obj = nn.Parameter(obj, requires_grad=False)
+        parent, _, leaf = name.rpartition(".")
+        if spec["kind"] == "parameter":
+            modules[parent].register_parameter(leaf, obj)
+        else:
+            modules[parent].register_buffer(leaf, obj, persistent=spec["persistent"])
+        registered[name] = obj
+
+
+def _gpu_uuid(device: torch.device) -> str:
+    return str(torch.cuda.get_device_properties(device).uuid)
+
+
+def _private_directory(path: Path) -> None:
+    path.mkdir(mode=0o700, parents=True, exist_ok=True)
+    info = path.stat()
+    if info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) & 0o077:
+        raise PermissionError("Weight-cache socket directory must be owned by this user with mode 0700")
+
+
+def serve_weight_cache(model: nn.Module, key: WeightCacheKey, socket_path: str | Path,
+                       *, stop: threading.Event | None = None) -> None:
+    """Serve an already-loaded, idle inference model until stop is set.
+
+    One active consumer is supported. Graceful stop waits for its connection to
+    close before releasing the model. SIGKILL/donor crashes invalidate consumers.
+    Callers own signal handling and must not mutate or offload donor tensors.
+    """
+    # Keep the ordinary loader importable on platforms without Linux IPC.
+    # Only the explicitly requested donor needs flock.
+    import fcntl
+
+    if model.training:
+        raise ValueError("The donor model must be in evaluation mode")
+    tensors, _, _ = _inventory(model)
+    devices = {tensor.device for tensor in tensors.values()}
+    if not tensors or len(devices) != 1 or next(iter(devices)).type != "cuda":
+        raise ValueError("The donor requires native tensors on one CUDA GPU")
+    device = next(iter(devices))
+    torch.cuda.synchronize(device)
+    manifest = _manifest(model)
+    path = Path(socket_path)
+    _private_directory(path.parent)
+    stop = stop if stop is not None else threading.Event()
+    lock_fd = os.open(str(path) + ".lock", os.O_CREAT | os.O_RDWR, 0o600)
+    try:
+        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+    except BaseException:
+        os.close(lock_fd)
+        raise
+    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    try:
+        path.unlink(missing_ok=True)  # Only the lock holder removes a stale socket.
+        server.bind(str(path))
+        os.chmod(path, 0o600)
+        server.listen()
+        server.settimeout(0.5)
+        while not stop.is_set():
+            try:
+                conn, _ = server.accept()
+            except socket.timeout:
+                continue
+            with conn:
+                try:
+                    _same_user(conn)
+                    conn.settimeout(30)
+                    request = _recv(conn)
+                    if request != {"version": _VERSION, "key": asdict(key), "torch": torch.__version__,
+                                   "gpu_uuid": _gpu_uuid(device)}:
+                        _send(conn, {"error": "Weight-cache identity, runtime or GPU mismatch"})
+                        continue
+                    _send(conn, {"manifest": manifest, "pid": os.getpid()})
+                    if _recv(conn) != "attach":
+                        continue
+                    # Generate reductions PER CONSUMER: each export acquires its
+                    # own PyTorch CUDA IPC refcounts. Never reuse reduction args.
+                    entries = {name: reduce_tensor(tensor.detach())[1] for name, tensor in tensors.items()}
+                    _send(conn, entries)
+                    conn.settimeout(None)
+                    conn.recv(1)  # EOF/release: keep donor storage alive throughout the lease.
+                except (ConnectionError, OSError, EOFError):
+                    continue
+                except (ValueError, TypeError, pickle.UnpicklingError) as error:
+                    with contextlib.suppress(OSError):
+                        _send(conn, {"error": str(error)})
+                finally:
+                    torch.cuda.ipc_collect()
+    finally:
+        server.close()
+        path.unlink(missing_ok=True)
+        os.close(lock_fd)
+
+
+class WeightCacheClient:
+    """Attach fresh module objects to donor-owned GPU storage; no disk fallback."""
+    def __init__(self, socket_path: str | Path, key: WeightCacheKey, device: torch.device | None = None):
+        self.socket_path = Path(socket_path)
+        self.key = key
+        self.device = device if device is not None else torch.device("cuda", torch.cuda.current_device())
+
+    @contextmanager
+    def load_model(self, factory: Callable[[], nn.Module]) -> Iterator[WeightAttachment]:
+        """Construct on meta, validate the full layout, and bind IPC tensors.
+
+        Hold this context until all GPU work finishes. Shared tensors are for
+        read-only inference. Mutable request state must be constructed locally;
+        training, in-place weight edits, CPU buffers and DTensor are unsupported.
+        """
+        with torch.device("meta"):
+            model = factory()
+        meta_tensors, _, _ = _inventory(model)
+        template = _manifest(model)
+        if not meta_tensors or any(tensor.device.type != "meta" for tensor in meta_tensors.values()):
+            raise ValueError("Consumer factory must construct all registered tensors on meta")
+        _private_directory(self.socket_path.parent)
+        info = self.socket_path.stat()
+        if info.st_uid != os.getuid() or not stat.S_ISSOCK(info.st_mode) or stat.S_IMODE(info.st_mode) & 0o077:
+            raise PermissionError("Weight-cache socket must be private and owned by this user")
+        conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        conn.settimeout(30)
+        mapped = {}
+        bound = False
+        try:
+            conn.connect(str(self.socket_path))
+            _same_user(conn)
+            _send(conn, {"version": _VERSION, "key": asdict(self.key), "torch": torch.__version__,
+                         "gpu_uuid": _gpu_uuid(self.device)})
+            header = _recv(conn)
+            if "error" in header:
+                raise ValueError(header["error"])
+            _validate_manifest(model, header["manifest"])
+            _send(conn, "attach")
+            entries = _recv(conn)
+            index = self.device.index if self.device.index is not None else torch.cuda.current_device()
+            mapped, handles = {}, {}
+            for name, exported in entries.items():
+                args = list(exported)
+                args[6] = index  # Exporter and consumer can have different visible-device ordinals.
+                mapped[name] = rebuild_cuda_tensor(*args)
+                handles[name] = hashlib.sha256(args[7] or b"").hexdigest()
+            _bind(model, mapped, header["manifest"])
+            bound = True
+            yield WeightAttachment(model.eval(), header["pid"], handles)
+        finally:
+            # Synchronize even when inference raises before releasing the lease.
+            try:
+                torch.cuda.synchronize(self.device)
+                if bound:
+                    _bind(model, meta_tensors, template)
+                mapped.clear()
+            finally:
+                conn.close()

fastvideo/pipelines/composed_pipeline_base.py

@@ -9,6 +9,7 @@ import argparse
 import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
+from contextlib import ExitStack
 from functools import partial
 from typing import Any, cast
 
@@ -143,6 +144,7 @@ class ComposedPipelineBase(ABC):
         self._stages: list[PipelineStage] = []
         self._stage_name_mapping: dict[str, PipelineStage] = {}
         self._trace_mgr = None
+        self._weight_cache_leases = ExitStack()
 
         if required_config_modules is not None:
             self._required_config_modules = required_config_modules
@@ -170,7 +172,11 @@ class ComposedPipelineBase(ABC):
         # Load modules directly in initialization
         logger.info("Loading pipeline modules...")
         with self.profiler_controller.region("profiler_region_model_loading"):
-            self.modules = self.load_modules(fastvideo_args, loaded_modules)
+            try:
+                self.modules = self.load_modules(fastvideo_args, loaded_modules)
+            except BaseException:
+                self._weight_cache_leases.close()
+                raise
 
     def set_trainable(self) -> None:
         # Only train DiT
@@ -627,6 +633,7 @@ class ComposedPipelineBase(ABC):
                     component_model_path=component_model_path,
                     transformers_or_diffusers=transformers_or_diffusers,
                     fastvideo_args=fastvideo_args,
+                    weight_cache_leases=self._weight_cache_leases,
                 )
 
             if self._lazy_module_load_enabled(fastvideo_args) and module_name in self._lazy_module_names:
@@ -787,6 +794,9 @@ class ComposedPipelineBase(ABC):
         raise NotImplementedError("if training_mode is True, the pipeline must implement this method")
 
     def close(self) -> None:
+        leases = getattr(self, "_weight_cache_leases", None)
+        if leases is not None:
+            leases.close()
         detach_activation_trace(getattr(self, "_trace_mgr", None))
         self._trace_mgr = None
 

fastvideo/pipelines/lora_pipeline.py

@@ -424,6 +424,8 @@ class LoRAPipeline(ComposedPipelineBase):
             accumulate: Add this adapter to an already merged pure low-rank adapter.
         """
 
+        if getattr(self.fastvideo_args, "weight_cache_dir", None):
+            raise ValueError("LoRA edits cannot mutate shared IPC weights")
         if not math.isfinite(strength):
             raise ValueError(f"LoRA strength must be finite, got {strength}")
 
@@ -592,6 +594,8 @@ class LoRAPipeline(ComposedPipelineBase):
         NVFP4 layers are only supported while their BF16 weight is retained: a purged
         BF16 weight cannot be unmerged and remerged.
         """
+        if getattr(self.fastvideo_args, "weight_cache_dir", None):
+            raise ValueError("LoRA edits cannot mutate shared IPC weights")
         if _has_nvfp4_weights_without_bf16(self.trainable_transformer_modules):
             # TODO(David): Merge from the packed FP4 weight instead of rejecting.
             raise RuntimeError(
@@ -616,6 +620,8 @@ class LoRAPipeline(ComposedPipelineBase):
 
         NVFP4 is rejected: its FP4 buffers cannot be rebuilt from a purged BF16 weight.
         """
+        if getattr(self.fastvideo_args, "weight_cache_dir", None):
+            raise ValueError("LoRA edits cannot mutate shared IPC weights")
         if _has_quantized_nvfp4_weights(self.trainable_transformer_modules):
             # TODO(David): Preserve BF16 weights and requantize NVFP4 weights after LoRA unmerge.
             raise RuntimeError(

fastvideo/tests/loader/test_inference_weight_cache.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Normal-loader IPC dispatch fails closed and scopes leases to the pipeline."""
+from contextlib import ExitStack, contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fastvideo.fastvideo_args import ExecutionMode, FastVideoArgs
+from fastvideo.models.loader.component_loader import ComponentLoader, PipelineComponentLoader
+from fastvideo.models.loader import inference_weight_cache as cache
+from fastvideo.models.wan.pipeline_config import WanT2V480PConfig
+
+
+def args(**overrides):
+    config = dict(model_path="unused", pipeline_config=WanT2V480PConfig(), num_gpus=1,
+                  tp_size=1, sp_size=1, hsdp_replicate_dim=1, hsdp_shard_dim=1,
+                  dit_cpu_offload=False, dit_layerwise_offload=False, vae_cpu_offload=False,
+                  text_encoder_cpu_offload=False, image_encoder_cpu_offload=False, lazy_module_load=False,
+                  inference_torch_compile=False, weight_cache_checkpoint="immutable-revision")
+    config.update(overrides)
+    return FastVideoArgs(**config)
+
+
+@pytest.mark.parametrize("option,value", [
+    ("inference_mode", False), ("use_fsdp_inference", True), ("num_gpus", 2),
+    ("dit_cpu_offload", True), ("dit_layerwise_offload", True), ("text_encoder_cpu_offload", True),
+    ("vae_cpu_offload", True), ("lora_path", "adapter"), ("enable_torch_compile", True),
+    ("init_weights_from_safetensors", "override"), ("weight_cache_checkpoint", ""),
+])
+def test_unsupported_modes_rejected_before_disk_loading(option, value):
+    runtime = args(weight_cache_dir="unused")
+    setattr(runtime, option, value)
+    with pytest.raises(ValueError):
+        with cache.component_weight_cache("transformer", runtime, ExitStack()):
+            pytest.fail("Unsupported mode reached a loader")
+
+
+def test_training_mode_rejected_even_if_inference_boolean_is_stale():
+    runtime = args()
+    runtime.mode = ExecutionMode.FINETUNING
+    with pytest.raises(ValueError, match="inference only"):
+        cache.validate_weight_cache_args(runtime)
+
+
+def test_ipc_requires_explicit_lifetime_owner():
+    with pytest.raises(ValueError, match="ExitStack"):
+        with cache.component_weight_cache("vae", args(weight_cache_dir="unused"), None):
+            pytest.fail("IPC cannot escape an unmanaged lifetime")
+
+
+def test_unknown_adapter_cannot_silently_load_from_disk():
+    with pytest.raises(ValueError, match="No native"):
+        with cache.component_weight_cache("vae", args(), None):
+            pass
+
+
+def test_context_restores_after_failed_loader():
+    with pytest.raises(RuntimeError):
+        with cache.component_weight_cache("vae", args(), None):
+            raise RuntimeError("Failed constructor")
+    assert cache.attach_cached_model(None, None, None, lambda: None) is None
+
+
+@pytest.mark.parametrize("component", cache.COMPONENTS)
+def test_normal_dispatch_attaches_all_components_without_disk_hydration(monkeypatch, component):
+    runtime = args(weight_cache_dir="/private/cache")
+    cls = type(cache._CLASSES[component], (torch.nn.Linear,), {})
+    calls = []
+
+    class Loader:
+        def load(self, path, runtime):
+            model = cache.attach_cached_model(cls, {"width": 2}, torch.float32, lambda: cls(2, 2))
+            assert model is not None, "Would have reached disk hydration"
+            return model
+
+    class Client:
+        def __init__(self, path, key):
+            assert str(path) == f"/private/cache/{component}.sock"
+            self.key = key
+
+        @contextmanager
+        def load_model(self, factory):
+            with torch.device("meta"):
+                model = factory()
+            calls.append("attach")
+            try:
+                yield SimpleNamespace(model=model, donor_pid=123, handles={"weight": "handle"})
+            finally:
+                calls.append("release")
+
+    monkeypatch.setattr(ComponentLoader, "for_module_type", lambda *a: Loader())
+    monkeypatch.setattr(cache, "WeightCacheClient", Client)
+    leases = ExitStack()
+    model = PipelineComponentLoader.load_module(component, "unused", "diffusers", runtime,
+                                               weight_cache_leases=leases)
+    assert calls == ["attach"]
+    assert model._fastvideo_weight_cache_key.component == component
+    assert model._fastvideo_weight_cache_witness["donor_pid"] == 123
+    leases.close()
+    assert calls == ["attach", "release"]
+
+
+def test_key_changes_for_dtype_configuration_and_checkpoint():
+    cls = type("AutoencoderKLWan", (torch.nn.Linear,), {})
+    def key(config, dtype=torch.float32, checkpoint="revision"):
+        with cache.component_weight_cache("vae", args(weight_cache_checkpoint=checkpoint), None) as context:
+            assert cache.attach_cached_model(cls, config, dtype, lambda: None) is None
+        return context.key
+    baseline = key({"width": 2})
+    assert baseline == key({"width": 2})
+    assert baseline != key({"width": 3})
+    assert baseline != key({"width": 2}, torch.bfloat16)
+    assert baseline != key({"width": 2}, checkpoint="other")
+
+
+@pytest.mark.parametrize("method", ["set_lora_adapter", "merge_lora_weights", "unmerge_lora_weights"])
+def test_runtime_lora_apis_cannot_modify_shared_weights(method):
+    from fastvideo.pipelines.lora_pipeline import LoRAPipeline
+    pipeline = SimpleNamespace(fastvideo_args=args(weight_cache_dir="unused"))
+    kwargs = {"lora_nickname": "test"} if method == "set_lora_adapter" else {}
+    with pytest.raises(ValueError, match="cannot mutate"):
+        getattr(LoRAPipeline, method)(pipeline, **kwargs)
+
+
+def test_pipeline_close_releases_leases_and_is_idempotent():
+    from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+    closed = []
+    leases = ExitStack()
+    leases.callback(lambda: closed.append(True))
+    pipeline = SimpleNamespace(_weight_cache_leases=leases, _trace_mgr=None)
+    ComposedPipelineBase.close(pipeline)
+    ComposedPipelineBase.close(pipeline)
+    assert closed == [True]
+
+
+def test_partial_pipeline_load_failure_releases_existing_attachments(monkeypatch):
+    from contextlib import nullcontext
+    from fastvideo.pipelines import composed_pipeline_base as base
+    released = []
+    pipeline = SimpleNamespace()
+    def failing_load(*args):
+        pipeline._weight_cache_leases.callback(lambda: released.append("first-component"))
+        raise ValueError("Later component incompatible")
+    pipeline.load_modules = failing_load
+    monkeypatch.setattr(base, "maybe_init_distributed_environment_and_model_parallel", lambda *a: None)
+    monkeypatch.setattr(base, "get_world_group", lambda: SimpleNamespace(local_rank=0))
+    monkeypatch.setattr(base, "get_or_create_profiler", lambda *a: SimpleNamespace(region=lambda *a: nullcontext()))
+    runtime = SimpleNamespace(inference_mode=False, tp_size=1, sp_size=1)
+    with pytest.raises(ValueError, match="Later component"):
+        base.ComposedPipelineBase.__init__(pipeline, "unused", runtime, required_config_modules=[])
+    assert released == ["first-component"]
+
+
+def test_unsupported_vae_special_constructor_rejected_before_checkpoint_access(monkeypatch):
+    from fastvideo.models.loader import component_loader
+    runtime = args(weight_cache_dir="unused")
+    monkeypatch.setattr(component_loader, "get_diffusers_config",
+                        lambda **kwargs: {"_class_name": "LingBotWorld2WanVAE"})
+    with ExitStack() as leases, pytest.raises(ValueError, match="Unsupported cached vae class"):
+        PipelineComponentLoader.load_module("vae", "unused", "diffusers", runtime,
+                                           weight_cache_leases=leases)

fastvideo/tests/loader/test_vae_loading_boundaries.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Real safetensors hydration remains equivalent to the combined native loader."""
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from fastvideo.models.loader import component_loader
+from fastvideo.models.loader.component_loader import PipelineComponentLoader, VAELoader
+from fastvideo.models.loader.utils import set_default_torch_dtype
+
+
+class TinyVAE(torch.nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.weight = torch.nn.Parameter(torch.zeros(2))
+        self.register_buffer("scale", torch.zeros(2))
+
+    def forward(self, x):
+        return x * self.weight + self.scale
+
+
+@pytest.mark.parametrize("sharded", [False, True])
+def test_native_load_matches_separate_construction_and_hydration(tmp_path, monkeypatch, sharded):
+    config = SimpleNamespace(update_model_arch=lambda _: None)
+    args = SimpleNamespace(
+        model_paths={}, vae_cpu_offload=False,
+        pipeline_config=SimpleNamespace(vae_config=config, vae_precision="fp32"),
+    )
+    monkeypatch.setattr(component_loader, "get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(component_loader, "get_diffusers_config", lambda **_: {"_class_name": "AutoencoderKLWan"})
+    monkeypatch.setattr(component_loader.ModelRegistry, "resolve_model_cls", lambda _: (TinyVAE, None))
+    tensors = {"weight": torch.tensor([2., 3.]), "scale": torch.tensor([5., 7.])}
+    if sharded:
+        for key, tensor in tensors.items():
+            save_file({key: tensor}, str(tmp_path / f"{key}.safetensors"))
+    else:
+        save_file(tensors, str(tmp_path / "model.safetensors"))
+    loader = VAELoader()
+    combined = PipelineComponentLoader.load_module("vae", str(tmp_path), "diffusers", args)
+    with set_default_torch_dtype(torch.float32):
+        separate = loader.build_model(TinyVAE, config, torch.device("cpu"))
+    assert torch.count_nonzero(separate.weight) == 0
+    separate = loader.load_weights(separate, str(tmp_path), class_name="AutoencoderKLWan", config={},
+                                   target_device=torch.device("cpu"))
+    x = torch.tensor([11., 13.])
+    torch.testing.assert_close(combined(x), torch.tensor([27., 46.]), rtol=0, atol=0)
+    torch.testing.assert_close(combined(x), separate(x), rtol=0, atol=0)
+    assert not combined.training and not separate.training
+    for name, tensor in separate.state_dict().items():
+        torch.testing.assert_close(tensor, tensors[name], rtol=0, atol=0)
+
+
+def test_construction_does_not_read_checkpoint(monkeypatch):
+    def unexpected(*args, **kwargs):
+        raise AssertionError("construction accessed checkpoint")
+    monkeypatch.setattr(component_loader, "safetensors_load_file", unexpected)
+    vae = VAELoader().build_model(TinyVAE, None, torch.device("cpu"))
+    assert torch.count_nonzero(vae.weight) == 0
+
+
+def test_hydration_preserves_strict_loading_and_missing_checkpoint_errors(tmp_path):
+    loader = VAELoader()
+    vae = TinyVAE(None)
+    with pytest.raises(ValueError, match="No safetensors"):
+        loader.load_weights(vae, str(tmp_path), class_name="AutoencoderKL", config={},
+                            target_device=torch.device("cpu"))
+    save_file({"weight": torch.ones(2)}, str(tmp_path / "model.safetensors"))
+    with pytest.raises(RuntimeError, match="scale"):
+        loader.load_weights(vae, str(tmp_path), class_name="AutoencoderKL", config={},
+                            target_device=torch.device("cpu"))

fastvideo/tests/loader/test_weight_cache.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+"""IPC bindings preserve registration/alias identity and reject stale layouts."""
+from dataclasses import replace
+import multiprocessing as mp
+from pathlib import Path
+import time
+
+import pytest
+import torch
+
+from fastvideo.models.loader.weight_cache import (
+    WeightCacheClient, WeightCacheKey, _bind, _inventory, _manifest, serve_weight_cache,
+)
+
+
+class Tiny(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.tensor([2., 3.]))
+        self.alias = self.weight
+        self.register_buffer("offset", torch.tensor([5., 7.]), persistent=False)
+        self.requests = 0
+
+    def forward(self, x):
+        self.requests += 1
+        return x * self.weight + self.offset
+
+
+def key():
+    return WeightCacheKey("immutable-test-revision", "tiny", "config-sha", "torch.float32", "native-v1")
+
+
+def test_binding_preserves_ties_and_nonpersistent_buffers():
+    source = Tiny().eval()
+    tensors, _, _ = _inventory(source)
+    with torch.device("meta"):
+        target = Tiny()
+    _bind(target, tensors, _manifest(source))
+    assert target.weight is target.alias
+    assert target.weight.data_ptr() == source.weight.data_ptr()
+    assert "offset" not in target.state_dict()
+    torch.testing.assert_close(target(torch.tensor([11., 13.])), torch.tensor([27., 46.]), rtol=0, atol=0)
+    assert source.requests == 0 and target.requests == 1
+
+
+def test_incompatible_layout_is_rejected_before_binding():
+    source = Tiny()
+    with torch.device("meta"):
+        target = Tiny()
+    manifest = _manifest(source)
+    manifest["specs"]["weight"]["shape"] = (3,)
+    with pytest.raises(ValueError, match="do not match"):
+        _bind(target, _inventory(source)[0], manifest)
+    assert all(tensor.is_meta for tensor in target.parameters())
+
+
+def test_missing_cache_never_falls_back_to_disk(tmp_path):
+    with pytest.raises(FileNotFoundError):
+        with WeightCacheClient(tmp_path / "missing.sock", key(), torch.device("cuda:0")).load_model(Tiny):
+            pytest.fail("A missing donor must not load from disk")
+
+
+def test_identity_requires_every_field():
+    with pytest.raises(ValueError, match="nonempty"):
+        replace(key(), checkpoint="")
+
+
+def _donor(path, stop):
+    model = Tiny().cuda().eval()
+    serve_weight_cache(model, key(), path, stop=stop)
+    torch.testing.assert_close(model.weight.cpu(), torch.tensor([2., 3.]), rtol=0, atol=0)
+
+
+def _consumer(path, queue, mismatched=False):
+    client_key = replace(key(), checkpoint="wrong-revision") if mismatched else key()
+    try:
+        baseline = torch.cuda.memory_allocated()
+        with WeightCacheClient(path, client_key).load_model(Tiny) as attached:
+            model = attached.model
+            assert model.weight is model.alias
+            result = model(torch.tensor([11., 13.], device="cuda")).cpu().tolist()
+            assert model.requests == 1
+            queue.put({"result": result, "pid": attached.donor_pid, "handles": attached.handles,
+                       "allocated": torch.cuda.memory_allocated() - baseline})
+        assert model.weight.is_meta  # The model cannot be used beyond its lease.
+    except ValueError as error:
+        if not mismatched:
+            raise
+        queue.put({"rejected": str(error)})
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA IPC requires one GPU")
+def test_donor_survives_two_fresh_consumers_and_rejected_identity(tmp_path):
+    ctx = mp.get_context("spawn")
+    stop = ctx.Event()
+    queue = ctx.Queue()
+    path = str(tmp_path / "weights.sock")
+    donor = ctx.Process(target=_donor, args=(path, stop))
+    donor.start()
+    records = []
+    try:
+        deadline = time.monotonic() + 90
+        while not Path(path).exists():
+            assert donor.is_alive(), "Donor died before readiness"
+            assert time.monotonic() < deadline, "Donor readiness timeout"
+            time.sleep(0.1)
+        for mismatched in (False, True, False):
+            worker = ctx.Process(target=_consumer, args=(path, queue, mismatched))
+            worker.start()
+            try:
+                worker.join(90)
+                assert worker.exitcode == 0
+                records.append(queue.get(timeout=10))
+            finally:
+                if worker.is_alive():
+                    worker.terminate()
+                    worker.join(10)
+        assert records[0]["result"] == records[2]["result"] == [27., 46.]
+        assert records[0]["pid"] == records[2]["pid"] == donor.pid
+        assert records[0]["handles"] == records[2]["handles"]
+        assert "rejected" in records[1]
+        assert donor.is_alive()
+    finally:
+        stop.set()
+        donor.join(10)
+        if donor.is_alive():
+            donor.terminate()
+            donor.join(10)
+        queue.close()
+    assert donor.exitcode == 0
+
+
+def test_incomplete_tensor_set_does_not_partially_bind():
+    source = Tiny()
+    with torch.device("meta"):
+        target = Tiny()
+    tensors = _inventory(source)[0]
+    tensors.pop("offset")
+    with pytest.raises(ValueError, match="Incomplete"):
+        _bind(target, tensors, _manifest(source))
+    assert target.weight.is_meta and target.offset.is_meta
+
+
+def test_donor_rejects_training_models_before_serving(tmp_path):
+    with pytest.raises(ValueError, match="evaluation"):
+        serve_weight_cache(Tiny(), key(), tmp_path / "weights.sock")
+    assert not (tmp_path / "weights.sock").exists()
+
+
+def test_donor_rejects_cpu_weights_before_serving(tmp_path):
+    with pytest.raises(ValueError, match="one CUDA GPU"):
+        serve_weight_cache(Tiny().eval(), key(), tmp_path / "weights.sock")
+    assert not (tmp_path / "weights.sock").exists()

fastvideo/worker/gpu_worker.py

@@ -196,8 +196,7 @@ class Worker:
         logger.info("Worker %d shutting down...", self.rank, local_main_process_only=False)
         # Clean up resources
         if hasattr(self, 'pipeline') and self.pipeline is not None:
-            # Clean up pipeline resources if needed
-            pass
+            self.pipeline.close()
 
         # Destroy the distributed environment
         cleanup_dist_env_and_memory(shutdown_ray=False)

mkdocs.yml

@@ -255,6 +255,7 @@ nav:
     - Testing: contributing/testing.md
     - Environment Variables: contributing/env_vars.md
     - Performance Benchmarks: contributing/performance_benchmarks.md
+    - Retained Weights: contributing/weight_cache.md
     - Profiling: contributing/profiling.md
     - Activation Trace: contributing/activation_trace.md
     - Coding Agents: contributing/coding_agents.md