[Core] Expose HWR payload validation in diffusion startup

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

benchmarks/host_weight_runtime/README.md

@@ -1,57 +0,0 @@
-# Host weight dependency memory diagnostic
-
-`safetensors_retention.py` isolates repeated CPU `get_tensor()` calls from
-access to one cached tensor view. It imports neither vLLM nor HWR and uses a
-synthetic float32 payload controlled by `--tensor-elements` (default: 64).
-Tensor size can affect retention: a negative result for 16 elements does not
-rule out growth for 64 elements. The selected size is recorded in the JSON
-arguments and used for both file creation and the correctness check.
-HWR itself calls `get_tensor()` when
-acquiring a lease, then exposes cached views; this probe does not measure HWR
-requests or repeated lease acquisition.
-
-Run one mode per fresh process in the same Python environment and CPU affinity:
-
-```bash
-git rev-parse HEAD > revision.txt
-git status --short > working-tree.txt
-git diff > working-tree.patch
-# Select an allowed CPU from your process affinity; 0 is only an example.
-CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
-  benchmarks/host_weight_runtime/safetensors_retention.py \
-  --mode get_tensor --iterations 10000 --sample-every 5000 > feasibility.json
-```
-
-For a size comparison, run 16 and 64 elements twice each in fresh processes,
-with cached reuse as a control at both sizes:
-
-```bash
-for repetition in 1 2; do
-  for elements in 16 64; do
-    for mode in get_tensor reuse; do
-      CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
-        benchmarks/host_weight_runtime/safetensors_retention.py \
-        --mode "$mode" --tensor-elements "$elements" \
-        > "$mode-$elements-$repetition.json"
-    done
-  done
-done
-```
-
-The diagnostic sets Torch to one CPU thread, warms up before the baseline,
-and reports preparation separately from timed loop work. Checkpoints run GC
-and report Linux private/anonymous memory, RSS, and open descriptors. It checks
-that transient tensor objects and the final cached view are released. Linux
-`/proc/self/smaps_rollup` must be readable; invalid arguments and missing proc
-support fail explicitly. Keep raw JSON and the repository snapshot with results.
-
-Compare the final loop sample with the warmed baseline, then inspect the closed
-sample. Private-memory growth after Python tensor objects disappear is evidence
-of retention in the dependency/allocator path, not proof of a leak's root cause.
-RSS can include file-backed pages; it is not a private-memory metric. Samples
-also include small diagnostic bookkeeping allocations. Do not assert a portable
-memory threshold or turn this synthetic loop into a per-request service estimate.
-
-Do not drop shared caches, trim allocators, or add serving-time GC to make the
-numbers look smaller. A dependency replacement or version change requires its
-own reproducible comparison and correctness/lifecycle validation.

benchmarks/host_weight_runtime/safetensors_retention.py

@@ -1,120 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Linux CPU diagnostic: repeated safetensors materialization versus view reuse."""
-
-from __future__ import annotations
-
-import argparse
-import gc
-import json
-import os
-import platform
-import tempfile
-import time
-import weakref
-from pathlib import Path
-
-
-def positive_int(value: str) -> int:
-    parsed = int(value)
-    if parsed <= 0:
-        raise argparse.ArgumentTypeError("must be positive")
-    return parsed
-
-
-def sample(phase: str, iterations: int, loop_seconds: float) -> dict[str, object]:
-    gc.collect()
-    counters = {}
-    for line in Path("/proc/self/smaps_rollup").read_text().splitlines():
-        fields = line.split()
-        if len(fields) == 3 and fields[2] == "kB":
-            counters[fields[0].removesuffix(":")] = int(fields[1]) * 1024
-    return {
-        "phase": phase,
-        "iterations": iterations,
-        "loop_seconds": loop_seconds,
-        "rss_bytes": counters["Rss"],
-        "private_bytes": counters["Private_Clean"] + counters["Private_Dirty"],
-        "anonymous_bytes": counters["Anonymous"],
-        "fd_count": len(tuple(Path("/proc/self/fd").iterdir())),
-    }
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--mode", choices=("get_tensor", "reuse"), required=True)
-    parser.add_argument(
-        "--tensor-elements", type=positive_int, default=64, help="Number of float32 elements in the synthetic tensor"
-    )
-    parser.add_argument("--iterations", type=positive_int, default=1_000_000)
-    parser.add_argument("--sample-every", type=positive_int, default=100_000)
-    parser.add_argument("--warmup", type=positive_int, default=1000)
-    args = parser.parse_args()
-    if not Path("/proc/self/smaps_rollup").is_file():
-        parser.error("requires Linux /proc/self/smaps_rollup")
-
-    preparation_start = time.monotonic()
-    import safetensors
-    import torch
-    from safetensors import safe_open
-    from safetensors.torch import save_file
-
-    torch.set_num_threads(1)
-    torch.set_num_interop_threads(1)
-    samples = [sample("before_file", 0, 0.0)]
-    with tempfile.TemporaryDirectory(prefix="safetensors-retention-") as directory:
-        path = Path(directory) / "synthetic.safetensors"
-        save_file({"weight": torch.arange(args.tensor_elements, dtype=torch.float32)}, str(path))
-        with safe_open(path, framework="pt", device="cpu") as reader:
-            cached = reader.get_tensor("weight")
-            assert torch.equal(cached, torch.arange(args.tensor_elements, dtype=torch.float32))
-            for _ in range(args.warmup):
-                tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached
-                del tensor
-            samples.append(sample("baseline", 0, 0.0))
-            preparation_seconds = time.monotonic() - preparation_start
-            completed = 0
-            loop_seconds = 0.0
-            while completed < args.iterations:
-                count = min(args.sample_every, args.iterations - completed)
-                start = time.monotonic()
-                for _ in range(count):
-                    tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached
-                last_tensor = weakref.ref(tensor)
-                del tensor
-                loop_seconds += time.monotonic() - start
-                completed += count
-                snapshot = sample("loop", completed, loop_seconds)
-                snapshot["last_tensor_alive"] = last_tensor() is not None
-                if args.mode == "get_tensor":
-                    assert last_tensor() is None, "diagnostic retained a transient tensor"
-                samples.append(snapshot)
-            del cached
-        del reader
-        samples.append(sample("closed", completed, loop_seconds))
-        assert last_tensor() is None, "last tensor survived reader/view teardown"
-
-    print(
-        json.dumps(
-            {
-                "schema_version": 1,
-                "arguments": vars(args),
-                "python": platform.python_version(),
-                "torch": torch.__version__,
-                "safetensors": safetensors.__version__,
-                "kernel": platform.release(),
-                "hostname": platform.node(),
-                "pid": os.getpid(),
-                "cpu_affinity": sorted(os.sched_getaffinity(0)),
-                "torch_threads": torch.get_num_threads(),
-                "pythonmalloc": os.environ.get("PYTHONMALLOC", "default"),
-                "preparation_seconds": preparation_seconds,
-                "samples": samples,
-            },
-            indent=2,
-        )
-    )
-
-
-if __name__ == "__main__":
-    main()

docs/design/feature/host_weight_runtime.md

@@ -42,11 +42,6 @@ V1 does not include:
 
 ## Motivation and use cases
 
-For investigating CPU memory retained by dependency tensor materialization,
-see the [standalone safetensors diagnostic](https://github.com/vllm-project/vllm-omni/blob/main/benchmarks/host_weight_runtime/README.md).
-It distinguishes repeated dependency calls from reuse of cached views and does
-not establish per-request HWR leakage.
-
 Model loading can create the same final host representation repeatedly. This
 is especially expensive when loading performs checkpoint decoding, tensor
 renaming, TP slicing, quantization, packing, or scale construction before GPU
@@ -72,6 +67,26 @@ views and mapped ranges. A separate transport still decides whether to use
 registered mmap, private pinned staging, synchronous copies, or asynchronous
 H2D transfer.
 
+## Payload validation at startup
+
+Eligible diffusion loaders accept `--host-weight-runtime-validation` (or the
+`host_weight_runtime_validation` offline/stage configuration field):
+
+- `manifest_and_metadata` is the existing default. It validates identity,
+  metadata, tensor structure, and sizes, but does not detect payload-only
+  corruption.
+- `full_checksum` reads and hashes every payload on warm acquisition before
+  restoring weights. This adds startup I/O/CPU work proportional to artifact
+  bytes; it is not a per-inference check.
+
+For example, add `--host-weight-runtime-validation full_checksum` to an enabled
+HWR deployment. Preferred mode falls back to canonical loading and can publish
+a replacement after detecting corruption; required mode fails startup. Readers
+can use different validation levels in the same domain without changing artifact
+identity. Neither level authenticates an untrusted writer or prevents mutation
+after validation. Checksumming does not imply automatic recovery of a running
+service.
+
 ## Resolution behavior
 
 The loader resolves the immutable canonical source and computes the exact

tests/config/test_omni_config.py

@@ -193,6 +193,37 @@ def test_resolve_execution_mode_rejects_unknown_execution_type():
         omni_config_module._resolve_execution_mode("unknown_execution_type")
 
 
+@pytest.mark.parametrize("validation", ["manifest_and_metadata", "full_checksum"])
+def test_hwr_validation_survives_config_projection(validation: str):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    override = StageDeployConfig(stage_id=0, host_weight_runtime_validation=validation)
+    config = _from_pipeline_key(
+        "dreamzero",
+        deploy_config_path="dreamzero_tp1_cfg2",
+        cli_overrides={"host_weight_runtime_validation": validation},
+    )
+    stage = config.stage_by_id(0)
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.diffusion_config.host_weight_runtime_validation == validation
+    projection = omni_config_module._DiffusionConfigProjection.from_kwargs(
+        host_weight_runtime_validation=override.host_weight_runtime_validation
+    )
+    assert projection.host_weight_runtime_validation == validation
+    assert OmniDiffusionConfig(host_weight_runtime_validation=validation).host_weight_runtime_validation == validation
+    assert omni_config_module._DiffusionConfigProjection().host_weight_runtime_validation == "manifest_and_metadata"
+    assert OmniDiffusionConfig().host_weight_runtime_validation == "manifest_and_metadata"
+
+
+@pytest.mark.parametrize("validation", ["fs_verity", "unknown", None])
+def test_hwr_validation_rejects_unsupported_values(validation):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    for config_class in (OmniDiffusionConfig, omni_config_module._DiffusionConfigProjection):
+        with pytest.raises(ValueError, match="host_weight_runtime_validation"):
+            config_class(host_weight_runtime_validation=validation)
+
+
 def test_from_pipeline_config_preserves_current_pipeline_config_object():
     omni_config = _from_pipeline_key("minicpmo_4_5")
     pipeline = _resolve_pipeline_or_skip("minicpmo_4_5")

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -260,6 +260,78 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
     assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1
 
 
+@pytest.mark.parametrize("validation", ["manifest_and_metadata", "full_checksum"])
+@pytest.mark.parametrize("mode", ["preferred", "required"])
+def test_hwr_loader_payload_integrity_policy(tmp_path: Path, monkeypatch, validation: str, mode: str):
+    from vllm_omni.diffusion.offloader.startup import take_offload_startup_state
+
+    canonical = tmp_path / "canonical"
+    canonical.mkdir()
+    expected = torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)
+    save_file({"weight": expected}, str(canonical / "model.safetensors"))
+    root = tmp_path / "store"
+
+    def make_loader(runtime_mode: str):
+        config = _hwr_config(canonical, root, mode=runtime_mode)
+        config.host_weight_runtime_validation = validation
+        loader = DiffusersPipelineLoader(LoadConfig(), config)
+        pipeline = _HWRPipeline(canonical)
+        monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
+        return loader, pipeline
+
+    def close_plan(pipeline):
+        state = take_offload_startup_state(pipeline)
+        if state is not None and state.host_weight_plan is not None:
+            carrier = state.host_weight_plan.lease_carrier
+            if carrier is not None:
+                carrier.close()
+
+    cold_loader, cold = make_loader("preferred")
+    cold_loader.load_model(load_device="cpu", device=torch.device("cpu"))
+    assert cold.load_count == 1
+    close_plan(cold)
+    payloads = list((root / "artifacts").glob("*/*.safetensors"))
+    assert len(payloads) == 1
+    payload = payloads[0]
+    data = bytearray(payload.read_bytes())
+    data[-1] ^= 1  # Keep the header and file length valid; change only a BF16 payload byte.
+    original_mode = payload.stat().st_mode
+    payload.chmod(0o600)
+    payload.write_bytes(data)
+    payload.chmod(original_mode)
+
+    loader, pipeline = make_loader(mode)
+    try:
+        if validation == "full_checksum" and mode == "required":
+            with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
+                loader.load_model(load_device="cpu", device=torch.device("cpu"))
+            assert pipeline.load_count == 0
+            assert loader.take_host_weight_plan() is None
+        else:
+            loader.load_model(load_device="cpu", device=torch.device("cpu"))
+            if validation == "manifest_and_metadata":
+                assert pipeline.load_count == 0
+                assert not torch.equal(pipeline.transformer.weight, expected)
+            else:
+                assert pipeline.load_count == 1
+                assert torch.equal(pipeline.transformer.weight, expected)
+    finally:
+        close_plan(pipeline)
+
+    if validation == "full_checksum":
+        if mode == "required":
+            assert list((root / "deny").iterdir())
+        else:
+            assert list((root / "quarantine").iterdir())
+            warm_loader, warm = make_loader(mode)
+            try:
+                warm_loader.load_model(load_device="cpu", device=torch.device("cpu"))
+                assert warm.load_count == 0
+                assert torch.equal(warm.transformer.weight, expected)
+            finally:
+                close_plan(warm)
+
+
 def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present():
     cfg = SimpleNamespace(
         lora_backend="distill",

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -503,6 +503,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
             "--omni",
             "--enable-distributed-layerwise-offload",
             "--dlo-no-use-allgather",
+            "--host-weight-runtime-validation",
+            "full_checksum",
             "--host-weight-runtime-mode",
             "preferred",
             "--host-weight-runtime-root",
@@ -519,6 +521,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
     assert explicit_kwargs["host_weight_runtime_mode"] == "preferred"
     assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
     assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80
+    assert explicit_kwargs["host_weight_runtime_validation"] == "full_checksum"
+    assert engine_args["host_weight_runtime_validation"] == "full_checksum"
     assert engine_args["host_weight_runtime_mode"] == "preferred"
     assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
     assert engine_args["dlo_host_registration_limit_gib"] == 80

vllm_omni/config/omni_config.py

@@ -739,6 +739,7 @@ class _DiffusionConfigProjection:
     dlo_resident_layers: int = Field(default=0, ge=0)
     host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
     host_weight_runtime_root: str | None = None
+    host_weight_runtime_validation: str = "manifest_and_metadata"
     dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
@@ -892,6 +893,7 @@ class _DiffusionConfigProjection:
         validate_host_weight_runtime_options(
             mode=self.host_weight_runtime_mode,
             root=self.host_weight_runtime_root,
+            validation=self.host_weight_runtime_validation,
         )
         self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
             limit_gib=self.dlo_host_registration_limit_gib,

vllm_omni/config/stage_config.py

@@ -483,6 +483,7 @@ class StageDeployConfig:
     dlo_resident_layers: int | None = None
     host_weight_runtime_mode: str | None = None
     host_weight_runtime_root: str | None = None
+    host_weight_runtime_validation: str | None = None
     dlo_host_registration_limit_gib: float | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None

vllm_omni/diffusion/data.py

@@ -100,7 +100,9 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]
     return normalized
 
 
-def validate_host_weight_runtime_options(*, mode: object, root: object) -> None:
+def validate_host_weight_runtime_options(
+    *, mode: object, root: object, validation: object = "manifest_and_metadata"
+) -> None:
     """Validate HWR policy without touching the configured storage domain.
 
     Filesystem locality and store construction belong to the eligible loader
@@ -109,6 +111,8 @@ def validate_host_weight_runtime_options(*, mode: object, root: object) -> None:
     """
     if mode not in {"disabled", "preferred", "required"}:
         raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required")
+    if validation not in ("manifest_and_metadata", "full_checksum"):
+        raise ValueError("host_weight_runtime_validation must be manifest_and_metadata or full_checksum")
     if mode != "disabled" and (not isinstance(root, str) or not root.strip()):
         raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
 
@@ -831,6 +835,7 @@ class OmniDiffusionConfig:
     # existing loader/storage path.
     host_weight_runtime_mode: str = "disabled"
     host_weight_runtime_root: str | None = None
+    host_weight_runtime_validation: str = "manifest_and_metadata"
     # Optional per-worker ceiling for registering final-layout HWR mappings.
     # Zero adds no ceiling; pin_cpu_memory controls whether registration is tried.
     dlo_host_registration_limit_gib: float = 0.0
@@ -1241,6 +1246,7 @@ class OmniDiffusionConfig:
         validate_host_weight_runtime_options(
             mode=self.host_weight_runtime_mode,
             root=self.host_weight_runtime_root,
+            validation=self.host_weight_runtime_validation,
         )
         self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
             limit_gib=self.dlo_host_registration_limit_gib,

vllm_omni/diffusion/model_loader/host_weight_loader.py

@@ -26,6 +26,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan, Te
 logger = init_logger(__name__)
 
 if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
     from vllm_omni.diffusion.model_loader.host_weights.identity_adapter import FinalLayoutIdentityContext
     from vllm_omni.diffusion.model_loader.host_weights.source_identity import (
         NodeSourceDigestCache,
@@ -40,6 +41,8 @@ class _HWRCommitError(RuntimeError):
 class HWRLoaderMixin:
     """Optional final-layout HWR behavior shared by diffusion loaders."""
 
+    od_config: OmniDiffusionConfig
+
     @staticmethod
     def _identity_value(value: object) -> object:
         """Convert config objects into deterministic identity metadata."""
@@ -319,10 +322,12 @@ class HWRLoaderMixin:
             HostWeightLeaseCarrier,
             HostWeightRuntime,
             HostWeightRuntimeConfig,
+            IntegrityPolicy,
             ProductionPolicy,
             ResolutionOutcome,
             RuntimeMode,
             StorageDomainPolicy,
+            ValidationLevel,
         )
         from vllm_omni.host_weight_runtime.filesystem import detect_storage_class
 
@@ -363,6 +368,11 @@ class HWRLoaderMixin:
             HostWeightRuntimeConfig(
                 mode=mode,
                 domain=StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)),
+                integrity=IntegrityPolicy(
+                    local_lookup=ValidationLevel(
+                        getattr(self.od_config, "host_weight_runtime_validation", "manifest_and_metadata")
+                    ),
+                ),
                 production=ProductionPolicy(
                     allow_local_build=False,
                     allow_post_load_publish=True,

vllm_omni/engine/arg_utils.py

@@ -582,6 +582,7 @@ class OrchestratorArgs:
     dlo_resident_layers: int = 0
     host_weight_runtime_mode: str = "disabled"
     host_weight_runtime_root: str | None = None
+    host_weight_runtime_validation: str = "manifest_and_metadata"
     dlo_host_registration_limit_gib: float = 0.0
     boundary_ratio: float | None = None
     flow_shift: float | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1145,6 +1145,7 @@ class AsyncOmniEngine:
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
             "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
             "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
+            "host_weight_runtime_validation": kwargs.get("host_weight_runtime_validation", "manifest_and_metadata"),
             "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (

vllm_omni/entrypoints/cli/serve.py

@@ -770,6 +770,15 @@ class OmniServeCommand(CLISubcommand):
                 "preferred first."
             ),
         )
+        omni_config_group.add_argument(
+            "--host-weight-runtime-validation",
+            choices=("manifest_and_metadata", "full_checksum"),
+            default="manifest_and_metadata",
+            help=(
+                "HWR warm-acquisition validation. The default checks metadata, not payload corruption; "
+                "full_checksum reads and hashes every payload before restoring weights."
+            ),
+        )
         omni_config_group.add_argument(
             "--host-weight-runtime-root",
             type=str,