[Perf] Reuse verified H3 prepared weights across model loads

docs/design/minimax_h3/COLD_LOADING.md

@@ -0,0 +1,84 @@
+# H3 cold loading and reusable prepared weights
+
+## Scope and baseline
+
+Integration base: `fe67339ddf862df0441a4e844fc664d9cafdffd0` (`onecat/main`).
+This scope covers original H3, INT8 ConvRot and LightX2V/FastH3 adapters;
+checkpoint precision, attention, sampling and output semantics stay unchanged.
+
+The deployed four-card V100 host has 62 GiB RAM and uses disk-backed masters.
+A real original FastH3 request spent 973.19 seconds preparing the service and
+104.95 seconds generating. At 604 seconds into loading, worker disk counters
+already totalled 99.57 GiB read and 109.08 GiB written. The old path loads row
+storage, transforms it to column storage and snapshots it again. Temporary
+masters are reconstructed at each service launch. Shared text encoder loading
+also repeats across H3 variants. These observations are not a new benchmark.
+
+## Intended contract
+
+- Prepare final projection strides before loading, avoiding a second full copy.
+- Reuse exact per-component, per-TP-rank prepared CPU storage, including the
+  shared text encoder, without copying a second complete model for publication.
+- Key entries by source/checkpoint identity, precision, topology and relevant
+  adapter/configuration. Retain original model files and all existing variants.
+- Publish atomically only after loading/validation succeeds; validate metadata
+  and data on restore. Use private mappings to prevent writes changing cache.
+- Bound disk use and evict only inactive cache entries under leases. Corrupt or
+  incomplete entries must rebuild, never silently serve unverified weights.
+- Report preparation subphases and distinguish first construction from reuse.
+
+## Validation status
+
+Implementation and targeted tests are in progress. No new startup speed or
+quality claim is made. GPU validation waits for the user's current generation
+to finish; no active user task is interrupted. Real first-load/reload timing
+and output comparison are required before promotion.
+
+AI assistance: OpenAI Codex.
+
+## Candidate implementation
+
+`--prepared-weight-cache` enables exact rank-local prepared transformer and
+text-encoder storage under `VLLM_CACHE_ROOT/h3-prepared`. The default budget is
+128 GiB (`--prepared-weight-cache-gib`). Native callers remain opt-in. Studio
+uses the capability probe to enable it independently of generation fast mode.
+
+The initial fill allocates the final projection strides before loading. A
+completed entry includes checksums of all storage files, tensor metadata and
+its manifest. Subsequent loads validate it before binding private mappings.
+Keys include checkpoint identity, source, Torch/CUDA, processing configuration
+and TP rank/topology. Ordinary LoRA sidecars are applied after the base cache;
+FastH3 fusion is keyed by its exact adapter. Active entries are protected by
+lifetime leases. Only inactive prepared entries can be evicted; original model
+files are never removed. Missing, incomplete or corrupt entries are rebuilt.
+
+Loading reports actual tensor or byte counts separately from component counts.
+No overall time percentage or performance prediction is synthesized.
+
+CPU validation: 80 passed, 3 skipped across prepared weights, host residency,
+service, progress, FastH3 and Studio fast-path suites. These cover exact signed
+INT8 bytes, FP16 layout, shared storage offsets, cache corruption and active
+lease protection. Real GPU startup, output parity and speed remain pending;
+this candidate must not be promoted based only on the CPU result.
+
+## Video VAE loading
+
+The native video component is a 2.60-billion-element FP32 checkpoint. Its
+reference factory initializes parameters and then replaces them with a strict
+state-dict load. Prepared loading skips random fills only for Parameters that
+are subsequently covered by a successful checkpoint load. Other tensors and
+buffers keep their normal initialization; incomplete/custom loaders retry
+normally. Temporary initialization hooks are always restored.
+
+A complete CPU state dict can also be assigned without copying when dtype,
+shape, strides and storage offsets match and neither side has tied storage.
+Custom parameter metadata, non-strided tensors, explicit assignment options,
+precision conversions and incomplete dictionaries retain normal semantics.
+The checkpoint's private mappings do not allow writes to change the file.
+Audio VAE loading is unchanged.
+
+An isolated V100-host CPU check of the actual video component produced the
+same full parameter-and-buffer SHA256 with and without this optimization:
+`3145d6a4c1b6576045991a36cf7b3606ec327c0c36d4152b91cf4498e3a9dfed`.
+The final loader selected assignment and took 1.67 seconds in that component
+check. This excludes downstream staging and is not a full-model startup time.

tests/video/test_h3_host_memory.py

@@ -48,7 +48,7 @@ def test_both_vae_stagers_honor_the_host_policy(monkeypatch, pin_memory, kind):
     monkeypatch.setattr(
         vae, "_load_component_config", lambda path: {"sample_rate": 44100}
     )
-    monkeypatch.setattr(vae, "_load_remote_component", lambda *args: remote)
+    monkeypatch.setattr(vae, "_load_remote_component", lambda *args, **kwargs: remote)
     monkeypatch.setattr(
         vae,
         "PinnedModuleStager",

tests/video/test_h3_initialization.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.models.minimax_h3.initialization import (
+    load_without_random_parameter_init,
+)
+
+
+def test_only_checkpoint_replaced_parameters_skip_random_fills():
+    original = nn.init.kaiming_uniform_
+    state = {"weight": torch.ones(4, 3), "bias": torch.full((4,), 2.0)}
+
+    def factory():
+        model = nn.Linear(3, 4)
+        constant = torch.empty(5)
+        nn.init.uniform_(constant, 3, 3)
+        model.register_buffer("constant", constant, persistent=False)
+        model.load_state_dict(state, strict=True)
+        return model
+
+    model = load_without_random_parameter_init(factory)
+    assert nn.init.kaiming_uniform_ is original
+    assert torch.equal(model.weight, state["weight"])
+    assert torch.equal(model.constant, torch.full((5,), 3.0))
+    assert torch.equal(model(torch.ones(1, 3)), torch.full((1, 4), 5.0))
+
+
+def test_partial_checkpoint_retries_with_normal_initialization():
+    calls = []
+
+    def factory():
+        calls.append(True)
+        model = nn.Linear(3, 4)
+        model.load_state_dict({"weight": torch.ones(4, 3)}, strict=False)
+        return model
+
+    model = load_without_random_parameter_init(factory)
+    assert len(calls) == 2
+    assert torch.isfinite(model.bias).all()
+    assert torch.all(model.bias.abs() <= 1 / 3**0.5)
+
+
+def test_initializer_and_load_hooks_are_restored_on_error():
+    initialize, load = nn.init.uniform_, nn.Module.load_state_dict
+
+    def factory():
+        nn.Linear(3, 4).load_state_dict({}, strict=True)
+
+    with pytest.raises(RuntimeError):
+        load_without_random_parameter_init(factory)
+    assert nn.init.uniform_ is initialize
+    assert nn.Module.load_state_dict is load
+
+
+def test_assigning_checkpoint_storage_is_supported():
+    def factory():
+        model = nn.Linear(3, 4, bias=False)
+        model.load_state_dict({"weight": torch.ones(4, 3)}, assign=True)
+        return model
+
+    model = load_without_random_parameter_init(factory)
+    assert torch.equal(model.weight, torch.ones(4, 3))
+
+
+def test_matching_checkpoint_is_assigned_without_a_second_copy(tmp_path):
+    from safetensors.torch import load_file, save_file
+
+    path = tmp_path / "weights.safetensors"
+    save_file({"weight": torch.arange(12, dtype=torch.float32).reshape(4, 3)}, path)
+    original = path.read_bytes()
+    state = load_file(path)
+
+    def factory():
+        model = nn.Linear(3, 4, bias=False)
+        model.load_state_dict(state)
+        return model
+
+    model = load_without_random_parameter_init(factory)
+    assert model.weight.data_ptr() == state["weight"].data_ptr()
+    with torch.no_grad():
+        model.weight.add_(1)
+    assert path.read_bytes() == original  # safetensors CPU mappings are private.
+
+
+def test_dtype_conversion_keeps_original_copy_semantics():
+    state = {"weight": torch.arange(12, dtype=torch.float16).reshape(4, 3)}
+
+    def factory():
+        model = nn.Linear(3, 4, bias=False, dtype=torch.float32)
+        model.load_state_dict(state)
+        return model
+
+    model = load_without_random_parameter_init(factory)
+    assert model.weight.dtype == torch.float32
+    assert model.weight.data_ptr() != state["weight"].data_ptr()
+    assert torch.equal(model.weight, state["weight"].float())
+
+
+def test_tied_parameters_are_not_detached_by_assignment():
+    def factory():
+        model = nn.Module()
+        model.first = nn.Linear(3, 4, bias=False)
+        model.second = nn.Linear(3, 4, bias=False)
+        model.second.weight = model.first.weight
+        model.load_state_dict(
+            {"first.weight": torch.ones(4, 3), "second.weight": torch.full((4, 3), 2.0)}
+        )
+        return model
+
+    model = load_without_random_parameter_init(factory)
+    assert model.first.weight is model.second.weight
+    assert torch.equal(model.first.weight, torch.full((4, 3), 2.0))
+
+
+def test_assignment_does_not_introduce_new_parameter_aliases():
+    def factory():
+        model = nn.Sequential(nn.Linear(3, 4, bias=False), nn.Linear(3, 4, bias=False))
+        weight = torch.ones(4, 3)
+        model.load_state_dict({"0.weight": weight, "1.weight": weight})
+        return model
+
+    model = load_without_random_parameter_init(factory)
+    assert model[0].weight.data_ptr() != model[1].weight.data_ptr()
+
+
+def test_storage_marker_expires_after_parameter_conversion():
+    from vllm.model_executor.models.minimax_h3.initialization import (
+        uses_assigned_checkpoint_storage,
+    )
+
+    def factory():
+        model = nn.Linear(3, 4, bias=False)
+        model.load_state_dict({"weight": torch.ones(4, 3)})
+        return model
+
+    model = load_without_random_parameter_init(factory)
+    assert uses_assigned_checkpoint_storage(model)
+    model.double()
+    assert not uses_assigned_checkpoint_storage(model)
+
+
+def test_replica_identity_requires_all_workers_and_rejects_replacement(monkeypatch):
+    from vllm.model_executor.models.minimax_h3 import vae
+
+    monkeypatch.setattr(vae.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(vae.dist, "get_world_size", lambda: 2)
+    values: list[tuple[str, int] | None] = [("same", 1), ("same", 1)]
+    monkeypatch.setattr(
+        vae.dist,
+        "all_gather_object",
+        lambda result, value: result.__setitem__(slice(None), values),
+    )
+    assert vae._same_checkpoint_replicas(values[0])
+    values[1] = None
+    assert not vae._same_checkpoint_replicas(values[0])
+    values[1] = ("replacement", 2)
+    with pytest.raises(ValueError, match="changed between"):
+        vae._same_checkpoint_replicas(values[0])