[Core] Expose HWR payload validation in diffusion startup

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

docs/design/feature/host_weight_runtime.md

@@ -67,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,