From f88d757a20d507fb30560ea2fdea84185d875c68 Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Sun, 6 Sep 2026 15:56:17 +0800
Subject: [PATCH 01/10] [Bugfix] Escalate diffusion worker shutdown and retain
 survivors

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit c79c841e00d9b33cfd428aa60f5d74d10fe4720b)
---
 .../test_multiproc_engine_concurrency.py      | 133 +++++++++++++++++-
 .../diffusion/executor/multiproc_executor.py  |  71 +++++++---
 2 files changed, 185 insertions(+), 19 deletions(-)

diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py
index efe6bf9a320..ade8d8fa725 100644
--- a/tests/diffusion/test_multiproc_engine_concurrency.py
+++ b/tests/diffusion/test_multiproc_engine_concurrency.py
@@ -4,8 +4,10 @@
 import asyncio
 import multiprocessing as mp
 import queue
+import signal
 import threading
 import time
+import weakref
 from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock
 
@@ -1330,6 +1332,7 @@ def test_worker_joins_share_one_global_deadline(self, monkeypatch):
         class FakeProcess:
             def __init__(self, name):
                 self.name = name
+                self.pid = 123
                 self.alive = True
                 self.terminated = False
                 self.join_timeouts = []
@@ -1358,6 +1361,132 @@ def terminate(self):
         assert first.terminated and second.terminated
         assert not first.is_alive() and not second.is_alive()
 
+    @pytest.mark.parametrize("cooperative", [False, True])
+    def test_real_workers_exit_and_are_reaped(self, monkeypatch, cooperative):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        ctx = mp.get_context("fork")
+        ready, stop = ctx.Event(), ctx.Event()
+
+        def run_worker():
+            signal.signal(signal.SIGTERM, signal.SIG_IGN)
+            ready.set()
+            stop.wait()
+
+        dead = ctx.Process(target=lambda: None)
+        worker = ctx.Process(target=run_worker)
+        monkeypatch.setattr(executor_module, "_WORKER_SHUTDOWN_GRACE_S", 0.2)
+        monkeypatch.setattr(executor_module, "_WORKER_TERMINATE_GRACE_S", 0.1)
+        monkeypatch.setattr(executor_module, "_WORKER_KILL_GRACE_S", 5.0, raising=False)
+        dead.start()
+        worker.start()
+        try:
+            dead.join(5)
+            assert not dead.is_alive()
+            assert ready.wait(5), "worker did not install its signal handler"
+            mq = Mock()
+            if cooperative:
+                mq.enqueue.side_effect = lambda *args, **kwargs: stop.set()
+            else:
+                mq.enqueue.side_effect = OSError("shutdown queue unavailable")
+            cleaner = executor_module._ExecutorShutdownCleaner(mq, 2, [dead, worker])
+
+            cleaner()
+
+            assert not worker.is_alive()
+            assert worker.exitcode == (0 if cooperative else -signal.SIGKILL)
+            assert cleaner.processes == []
+            assert cleaner.broadcast_mq is None
+            cleaner()
+        finally:
+            for proc in (dead, worker):
+                if proc.is_alive():
+                    proc.kill()
+                proc.join(5)
+                proc.close()
+
+    @pytest.mark.parametrize("failed_action", ["terminate", "kill"])
+    def test_kill_phase_shares_deadline_and_continues_after_os_errors(self, monkeypatch, failed_action):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        first, second = Mock(pid=101), Mock(pid=102)
+        first.name, second.name = "first", "second"
+        for proc in (first, second):
+            proc.is_alive.side_effect = [True, True, True, False]
+        getattr(first, failed_action).side_effect = OSError("signal failed")
+        first.join.side_effect = [OSError("join failed"), None, None]
+        monotonic = Mock(side_effect=[100, 100, 110, 120, 120, 124, 130, 130, 134])
+        monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic))
+        cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second])
+
+        cleaner()
+
+        first.kill.assert_called_once()
+        second.kill.assert_called_once()
+        assert [call.args[0] for call in first.join.call_args_list] == [15, 5, 5]
+        assert [call.args[0] for call in second.join.call_args_list] == [5, 1, 1]
+        assert cleaner.processes == []
+
+    def test_executor_retains_survivor_for_explicit_shutdown_retry(self, monkeypatch):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        survivor = Mock(pid=123)
+        survivor.name = "surviving-worker"
+        survivor.is_alive.return_value = True
+        cleaner = executor_module._ExecutorShutdownCleaner(processes=[survivor])
+        executor, _, _ = _make_executor()
+        executor._shutdown_cleaner = cleaner
+        executor._processes = [survivor]
+        executor._finalizer = weakref.finalize(executor, cleaner)
+        executor._pump_stop = threading.Event()
+        executor._futures_lock = threading.RLock()
+        executor._rpc_futures = {}
+        executor._output_futures = {}
+        executor._batch_split_map = {}
+        log_error = Mock()
+        monkeypatch.setattr(executor_module.logger, "error", log_error)
+
+        executor.shutdown()
+
+        assert not executor._finalizer.alive
+        assert executor._shutdown_cleaner is cleaner
+        assert executor._processes == [survivor]
+        log_error.assert_called_once()
+        assert log_error.call_args.args[1] == [("surviving-worker", 123)]
+        survivor.is_alive.return_value = False
+
+        executor.shutdown()
+        executor.shutdown()
+
+        assert executor._shutdown_cleaner is None
+        assert executor._processes == []
+
+    def test_concurrent_cleanup_does_not_repeat_process_operations(self):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        joining, release = threading.Event(), threading.Event()
+        proc = Mock(pid=123)
+        proc.is_alive.return_value = True
+
+        def join(timeout):
+            joining.set()
+            assert release.wait(5)
+            proc.is_alive.return_value = False
+
+        proc.join.side_effect = join
+        cleaner = executor_module._ExecutorShutdownCleaner(processes=[proc])
+        thread = threading.Thread(target=cleaner)
+        thread.start()
+        try:
+            assert joining.wait(5)
+            cleaner()
+            proc.join.assert_called_once()
+        finally:
+            release.set()
+            thread.join(5)
+        assert not thread.is_alive()
+        assert cleaner.processes == []
+
 
 # ───────── monitor thread & death sentinel integration tests ─────────
 
@@ -1406,7 +1535,7 @@ def test_worker_monitor_sets_is_failed_and_calls_callbacks_on_death(self):
         executor._result_mq = None
         executor._shutdown_cleaner = None
         # Use a no-op so shutdown() doesn't crash on None resources.
-        executor._finalizer = lambda: None
+        executor._finalizer = weakref.finalize(executor, lambda: None)
         # ------------------------------------------------------------------
         # Attributes added by remove_bubble_v2 (async D2H); shutdown() iterates
         # over them, so they need to exist even when constructed via __new__.
@@ -1441,7 +1570,7 @@ def test_worker_monitor_noop_when_already_closed(self):
         executor._broadcast_mq = None
         executor._result_mq = None
         executor._shutdown_cleaner = None
-        executor._finalizer = lambda: None
+        executor._finalizer = weakref.finalize(executor, lambda: None)
 
         proc = _make_short_lived_process()
         executor._processes = [proc]
diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py
index 39e85abcc5c..57b23656cf8 100644
--- a/vllm_omni/diffusion/executor/multiproc_executor.py
+++ b/vllm_omni/diffusion/executor/multiproc_executor.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import concurrent.futures
@@ -10,7 +13,7 @@
 import time
 import weakref
 from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from multiprocessing.synchronize import Event
 from typing import TYPE_CHECKING, Any, cast
 
@@ -42,6 +45,7 @@
 _DLO_DP_WAVE_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", 600.0))
 _WORKER_SHUTDOWN_GRACE_S = 15.0
 _WORKER_TERMINATE_GRACE_S = 5.0
+_WORKER_KILL_GRACE_S = 5.0
 _RESULT_PUMP_JOIN_TIMEOUT_S = 2.0
 
 
@@ -80,33 +84,61 @@ class _ExecutorShutdownCleaner:
     broadcast_mq: MessageQueue | None = None
     num_workers: int = 0
     processes: list[mp.Process] | None = None
+    _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
 
     def __call__(self) -> None:
         """Clean up background resources."""
+        # The worker monitor and an explicit shutdown may race. A retry must
+        # not signal/join the same Process objects while cleanup is in flight.
+        if not self._lock.acquire(blocking=False):
+            return
+        try:
+            self._cleanup()
+        finally:
+            self._lock.release()
+
+    def _cleanup(self) -> None:
         if self.broadcast_mq is not None:
             try:
                 for _ in range(self.num_workers):
                     self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE, timeout=1.0)
 
-                self.broadcast_mq = None
             except Exception as exc:
                 logger.warning("Failed to send shutdown signal: %s", exc)
+            finally:
+                self.broadcast_mq = None
 
         if self.processes:
-            join_deadline = time.monotonic() + _WORKER_SHUTDOWN_GRACE_S
-            for proc in self.processes:
-                if not proc.is_alive():
-                    continue
-                proc.join(max(0.0, join_deadline - time.monotonic()))
-
             alive = [proc for proc in self.processes if proc.is_alive()]
-            for proc in alive:
-                logger.warning("Terminating diffusion worker %s after timeout", proc.name)
-                proc.terminate()
+            for action, grace in (
+                (None, _WORKER_SHUTDOWN_GRACE_S),
+                ("terminate", _WORKER_TERMINATE_GRACE_S),
+                ("kill", _WORKER_KILL_GRACE_S),
+            ):
+                if not alive:
+                    break
+                if action is not None:
+                    for proc in alive:
+                        try:
+                            logger.warning("Calling %s on diffusion worker %s (pid=%s)", action, proc.name, proc.pid)
+                            getattr(proc, action)()
+                        except OSError:
+                            logger.exception("Failed to %s diffusion worker %s (pid=%s)", action, proc.name, proc.pid)
 
-            terminate_deadline = time.monotonic() + _WORKER_TERMINATE_GRACE_S
-            for proc in alive:
-                proc.join(max(0.0, terminate_deadline - time.monotonic()))
+                deadline = time.monotonic() + grace
+                for proc in alive:
+                    try:
+                        proc.join(max(0.0, deadline - time.monotonic()))
+                    except OSError:
+                        logger.exception("Failed to join diffusion worker %s (pid=%s)", proc.name, proc.pid)
+                alive = [proc for proc in alive if proc.is_alive()]
+
+            self.processes = alive
+            if alive:
+                logger.error(
+                    "Diffusion worker cleanup incomplete after kill: %s; retaining processes for shutdown retry",
+                    [(proc.name, proc.pid) for proc in alive],
+                )
 
 
 class MultiprocDiffusionExecutor(DiffusionExecutor):
@@ -1006,8 +1038,12 @@ def check_health(self) -> None:
     def shutdown(self) -> None:
         self._closed = True
         self._pump_stop.set()
+        cleaner = self._shutdown_cleaner
         try:
-            self._finalizer()
+            if self._finalizer.alive:
+                self._finalizer()
+            elif cleaner is not None:
+                cleaner()
         finally:
             pump_threads = getattr(self, "_result_pump_threads", [])
             for thread in pump_threads:
@@ -1031,5 +1067,6 @@ def shutdown(self) -> None:
                 self._rpc_futures.clear()
                 self._output_futures.clear()
                 self._batch_split_map.clear()
-            self._shutdown_cleaner = None
-            self._processes = []
+            self._processes = (cleaner.processes or []) if cleaner is not None else []
+            if not self._processes:
+                self._shutdown_cleaner = None

From 61d5af91ab7bfa1f597568667cc208d475f0b2c7 Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Sun, 6 Sep 2026 16:06:21 +0800
Subject: [PATCH 02/10] [Bugfix] Bound HWR domain initialization lock waits

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit d455411e9e44d459c1d9b6644a3909f30d23dcb6)
---
 docs/design/feature/host_weight_runtime.md    |  11 +-
 docs/design/module/host_weight_runtime.md     |  12 +-
 .../test_runtime_resolution.py                | 108 ++++++++++++++++++
 .../host_weight_runtime/filesystem/store.py   |  10 +-
 vllm_omni/host_weight_runtime/runtime.py      |   2 +-
 5 files changed, 134 insertions(+), 9 deletions(-)

diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md
index a9bb1b0b838..ac7bc7b3f9d 100644
--- a/docs/design/feature/host_weight_runtime.md
+++ b/docs/design/feature/host_weight_runtime.md
@@ -313,8 +313,15 @@ One process per exact identity owns a build; other workers wait and then acquire
 leases for the published artifact. Publication is invisible until all payloads
 and metadata are validated, hashed, fsynced, and atomically renamed.
 
-`coordination_timeout_seconds` bounds filesystem lock acquisition. It does not
-cancel synchronous validation, a producer that has already started, or atomic
+`coordination_timeout_seconds` bounds domain-initialization and lookup/build
+lock acquisition. Store construction and each later resolution or publication
+operation have separate budgets from the same wait policy, rather than one
+end-to-end startup deadline. A domain-init timeout follows the retryable domain
+failure policy below. After contention ends, a fresh runtime construction can
+retry initialization; the timed-out runtime retains its original failure.
+
+The coordination budget does not cancel filesystem I/O, synchronous validation,
+a producer that has already started, or atomic
 publication. A hung in-process producer therefore blocks its owning process and
 must be handled by external process supervision. Enforceable producer
 cancellation requires a future process-isolated producer contract.
diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md
index ac92711fc4d..c975616671b 100644
--- a/docs/design/module/host_weight_runtime.md
+++ b/docs/design/module/host_weight_runtime.md
@@ -136,8 +136,16 @@ so only when it is marked retryable; unsupported capabilities, identity
 collisions, producer failures, and publication failures remain visible as
 startup failures.
 
-`coordination_timeout_seconds` bounds acquisition of lookup and build locks; it
-is not a hard wall-clock deadline for in-process validation, production, or
+`coordination_timeout_seconds` bounds acquisition of the domain-initialization
+lock and lookup/build locks. Store construction and each later resolution or
+publication operation receive separate coordination budgets from the same
+`WaitPolicy`; they do not share an end-to-end startup deadline. A domain-init
+lock timeout is a retryable domain failure: preferred mode may fall back and
+required mode fails. Retrying initialization requires constructing a fresh
+runtime after contention ends; a waiter never cancels or unlocks the owner.
+
+The coordination budget is not a hard wall-clock deadline for filesystem I/O,
+in-process validation, production, or
 atomic publication. Once this V1 implementation becomes the producer, the
 synchronous producer and publication run to completion. A hung producer blocks
 its owning process and requires external process supervision; enforceable
diff --git a/tests/host_weight_runtime/test_runtime_resolution.py b/tests/host_weight_runtime/test_runtime_resolution.py
index 21bab37dee4..cd2de28165d 100644
--- a/tests/host_weight_runtime/test_runtime_resolution.py
+++ b/tests/host_weight_runtime/test_runtime_resolution.py
@@ -4,6 +4,10 @@
 
 from __future__ import annotations
 
+import multiprocessing as mp
+import threading
+import time
+from multiprocessing.connection import Connection
 from pathlib import Path
 from typing import cast
 
@@ -15,6 +19,7 @@
     CanonicalJson,
     ComponentIdentity,
     FailureCode,
+    HostWeightError,
     HostWeightFailure,
     HostWeightLease,
     HostWeightLeaseCarrier,
@@ -40,12 +45,14 @@
     StoreStatus,
     TensorWriteSpec,
     ValidationLevel,
+    WaitPolicy,
     WeightArtifactIdentity,
     WeightProductionSpec,
     WeightRepresentation,
     WeightSourceIdentity,
 )
 from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class
+from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, lock_is_active
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -518,6 +525,107 @@ def test_retryable_store_initialization_failure_obeys_runtime_mode(
     assert resolution.report.attempts[0].failure.retryable
 
 
+def _initialize_contended_runtime(config: HostWeightRuntimeConfig, output: Connection) -> None:
+    try:
+        output.send("starting")
+        started = time.monotonic()
+        runtime = HostWeightRuntime.from_config(config)
+        resolution = runtime.resolve(_identity())
+        try:
+            output.send((time.monotonic() - started, resolution.report))
+        finally:
+            if resolution.lease is not None:
+                resolution.lease.close()
+    finally:
+        output.close()
+
+
+@pytest.mark.parametrize(
+    ("mode", "expected"),
+    [
+        (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK),
+        (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED),
+    ],
+)
+def test_domain_lock_timeout_and_fresh_runtime_recovery(
+    tmp_path: Path, mode: RuntimeMode, expected: ResolutionOutcome
+) -> None:
+    config = HostWeightRuntimeConfig(mode=mode, domain=_domain(tmp_path), wait=WaitPolicy(0.1))
+    initial = HostWeightRuntime.from_config(config)
+    built = initial.resolve(_identity(), producer=CountingProducer(_identity()))
+    assert built.lease is not None
+    built.lease.close()
+    metadata = {name: (tmp_path / name).read_bytes() for name in ("domain.json", "domain-policy.json")}
+    lock_path = tmp_path / "locks" / "domain-init.lock"
+    ctx = mp.get_context("spawn")
+    reader, writer = ctx.Pipe(duplex=False)
+    process = ctx.Process(target=_initialize_contended_runtime, args=(config, writer))
+
+    with FileLock(lock_path, exclusive=True, deadline=None):
+        process.start()
+        writer.close()
+        try:
+            # Import/bootstrap is outside the short coordination budget.
+            assert reader.poll(60), "waiter did not start"
+            assert reader.recv() == "starting"
+            assert reader.poll(5), "domain initialization did not respect the coordination timeout"
+            elapsed, report = reader.recv()
+            assert 0.1 <= elapsed < 5
+            assert report.outcome is expected
+            failure = report.attempts[0].failure
+            assert failure is not None
+            assert failure.stage is ResolutionStage.DOMAIN
+            assert failure.code is FailureCode.DOMAIN_UNAVAILABLE
+            assert failure.retryable
+            assert "timed out waiting for" in failure.message
+            assert "domain-init.lock" in failure.message
+            assert lock_is_active(lock_path), "waiter must not unlock the owner"
+            process.join(5)
+            assert process.exitcode == 0
+        finally:
+            if process.is_alive():
+                process.kill()
+            process.join(5)
+            process.close()
+            reader.close()
+
+    assert {name: (tmp_path / name).read_bytes() for name in metadata} == metadata
+    recovered = HostWeightRuntime.from_config(config).resolve(_identity())
+    assert recovered.report.outcome is ResolutionOutcome.LOCAL_HIT
+    assert recovered.lease is not None
+    recovered.lease.close()
+
+
+def test_direct_store_waits_for_domain_lock_released_before_deadline(tmp_path: Path) -> None:
+    config = HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path))
+    assert config.domain is not None
+    initial = HostWeightRuntime.from_config(config)
+    assert initial.store is not None
+    lock_path = tmp_path / "locks" / "domain-init.lock"
+    with FileLock(lock_path, exclusive=True, deadline=None) as owner:
+        release = threading.Timer(0.1, owner.close)
+        release.start()
+        try:
+            store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=WaitPolicy(5))
+        finally:
+            release.join(5)
+    assert store.root == tmp_path
+
+
+def test_direct_store_domain_timeout_is_typed(tmp_path: Path) -> None:
+    config = HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path))
+    assert config.domain is not None
+    HostWeightRuntime.from_config(config)
+    lock_path = tmp_path / "locks" / "domain-init.lock"
+    with FileLock(lock_path, exclusive=True, deadline=None):
+        with pytest.raises(HostWeightError) as caught:
+            FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=WaitPolicy(0.01))
+        assert caught.value.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+        assert caught.value.failure.retryable
+        assert lock_is_active(lock_path)
+    assert not lock_is_active(lock_path)
+
+
 @pytest.mark.parametrize(
     ("stage", "code"),
     [
diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py
index 0f01e36dac5..ac04881f900 100644
--- a/vllm_omni/host_weight_runtime/filesystem/store.py
+++ b/vllm_omni/host_weight_runtime/filesystem/store.py
@@ -23,7 +23,7 @@
 import torch
 from safetensors import SafetensorError, safe_open
 
-from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel
+from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel, WaitPolicy
 from ..errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionStage
 from ..identity import CanonicalJson, WeightArtifactIdentity, canonical_json
 from ..inspection import (
@@ -429,6 +429,8 @@ def __init__(
         domain: StorageDomainPolicy,
         capacity: CapacityPolicy,
         integrity: IntegrityPolicy,
+        *,
+        wait: WaitPolicy = WaitPolicy(),
     ) -> None:
         self.domain_policy = domain
         self.capacity_policy = capacity
@@ -464,7 +466,7 @@ def __init__(
                 )
             )
         try:
-            self._initialize_domain()
+            self._initialize_domain(deadline=time.monotonic() + wait.coordination_timeout_seconds)
         except HostWeightError:
             raise
         except OSError as exc:
@@ -485,7 +487,7 @@ def __init__(
                 )
             ) from exc
 
-    def _initialize_domain(self) -> None:
+    def _initialize_domain(self, *, deadline: float) -> None:
         try:
             self.filesystem_type = detect_filesystem_type(self.root)
             detected_storage_class = _storage_class_for_filesystem_type(self.filesystem_type)
@@ -545,7 +547,7 @@ def _initialize_domain(self) -> None:
                         )
                     )
 
-        with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=None):
+        with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=deadline):
             domain_path = self.root / _DOMAIN_FILE
             if domain_path.exists():
                 domain_metadata = _read_json_file(domain_path)
diff --git a/vllm_omni/host_weight_runtime/runtime.py b/vllm_omni/host_weight_runtime/runtime.py
index fbc726669cf..17ababce9eb 100644
--- a/vllm_omni/host_weight_runtime/runtime.py
+++ b/vllm_omni/host_weight_runtime/runtime.py
@@ -70,7 +70,7 @@ def from_config(
         from .filesystem import FilesystemHostWeightStore  # noqa: PLC0415
 
         try:
-            store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity)
+            store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=config.wait)
         except HostWeightError as exc:
             return cls(
                 config,

From 438e8231633ae39784136542c0ffe8d1e9313c2c Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Sun, 6 Sep 2026 16:12:19 +0800
Subject: [PATCH 03/10] [Bugfix] Skip HWR store-size scans when no limit is
 configured

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit 00d69846608b428f59bb7ebd3d16f7b99a893259)
---
 vllm_omni/host_weight_runtime/filesystem/store.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py
index ac04881f900..ec6b49f3c79 100644
--- a/vllm_omni/host_weight_runtime/filesystem/store.py
+++ b/vllm_omni/host_weight_runtime/filesystem/store.py
@@ -1249,8 +1249,10 @@ def _check_capacity(self, projected_artifact_bytes: int, additional_bytes: int)
                     retryable=True,
                 )
             )
-        store_bytes = self._tree_bytes(self.root)
-        if policy.max_store_bytes is not None and store_bytes + additional_bytes > policy.max_store_bytes:
+        if (
+            policy.max_store_bytes is not None
+            and self._tree_bytes(self.root) + additional_bytes > policy.max_store_bytes
+        ):
             raise HostWeightError(
                 _failure(
                     ResolutionStage.CAPACITY,

From 02ab15ff30ac3a006ea7ec7baaf1a357b098350b Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Sun, 6 Sep 2026 16:27:51 +0800
Subject: [PATCH 04/10] [Core] Add explicit cleanup for quarantined HWR
 artifacts

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit c6a0c7f6b00e6413f90bcd139405f24f1d39b631)
---
 docs/design/module/host_weight_runtime.md     |  37 ++++
 .../test_filesystem_store.py                  | 209 ++++++++++++++++++
 .../host_weight_runtime/filesystem/store.py   | 102 ++++++---
 3 files changed, 320 insertions(+), 28 deletions(-)

diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md
index c975616671b..3b1b90b89e1 100644
--- a/docs/design/module/host_weight_runtime.md
+++ b/docs/design/module/host_weight_runtime.md
@@ -307,6 +307,43 @@ Explicit cleanup uses the locked move and retries exact-key
 `.cleanup.*` tombstones left by an interrupted removal; a later build performs
 the same tombstone reconciliation before producing replacement content.
 
+Failure-quarantined entries remain until explicitly selected for removal;
+ordinary `cleanup(identity)` does not remove that diagnostic history.
+`FilesystemHostWeightStore.cleanup_quarantined(storage_name)` accepts one
+quarantined inventory basename and preserves the current artifact, deny marker,
+and other failure-quarantined entries. It takes the same nonblocking build and
+exclusive artifact locks, so an active builder or lease returns a typed refusal.
+Malformed names are rejected before mutation, and symlinks are not followed.
+
+Selected entries move through the existing `.cleanup.*` tombstone protocol.
+Retrying the original name reconciles pending cleanup tombstones for that key,
+even if the original entry disappeared during an earlier attempt. A missing
+selection is otherwise an idempotent success. Existing cleanup intents for the
+key may be completed along with the selection; other quarantine history is
+retained. Parent synchronization is retried even when an earlier deletion
+removed the last tombstone before its sync failed.
+
+For a configured store, operators can inspect and deliberately select an entry:
+
+```python
+from vllm_omni.host_weight_runtime import ArtifactInventoryState, HostWeightError
+
+for entry in store.inspect_domain().inventory:
+    if entry.state is ArtifactInventoryState.QUARANTINED:
+        print(entry.storage_name, entry.size_bytes)
+
+selected_storage_name = input("Quarantined entry to remove: ")
+failure = store.cleanup_quarantined(selected_storage_name)
+if failure is not None:
+    raise HostWeightError(failure)
+```
+
+Use the same authoritative domain and capacity policy as the service. Review
+inventory sizes and inspect again after cleanup before explicitly retrying a
+build or startup. These are logical-byte, point-in-time observations; concurrent
+activity can affect the difference and physical disk/RAM reclamation is not
+implied. This operation adds no automatic retention policy or service retry.
+
 An operational failure while inspecting a noncooperative competing publication
 is a retryable storage failure. The competing entry remains authoritative and
 is not mislabeled as corrupt or quarantined merely because its manifest could
diff --git a/tests/host_weight_runtime/test_filesystem_store.py b/tests/host_weight_runtime/test_filesystem_store.py
index 7fe716c5ddd..d4051fd0e81 100644
--- a/tests/host_weight_runtime/test_filesystem_store.py
+++ b/tests/host_weight_runtime/test_filesystem_store.py
@@ -440,6 +440,22 @@ def _corrupt_test_payload(artifact: Path) -> None:
     os.chmod(artifact, 0o555)
 
 
+def _quarantine_and_rebuild(store: FilesystemHostWeightStore, identity: WeightArtifactIdentity) -> Path:
+    before = set(store.quarantine_dir.iterdir())
+    _corrupt_test_payload(store.artifacts_dir / identity.key)
+    result = store.get_or_build(
+        BuildRequest(identity),
+        FakeProducer(identity),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 10,
+    )
+    assert result.status is StoreStatus.BUILT
+    assert result.lease is not None
+    result.lease.close()
+    (quarantined,) = set(store.quarantine_dir.iterdir()) - before
+    return quarantined
+
+
 def _build_process(root: str, counter: str, initial_lookup_barrier: object, result_queue: object) -> None:
     try:
         identity = _identity()
@@ -2496,6 +2512,199 @@ def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path:
     assert not list(free_limited.artifacts_dir.iterdir())
 
 
+def test_cleanup_quarantined_preserves_current_artifact_other_entries_and_deny(tmp_path: Path) -> None:
+    store = _make_store(tmp_path / "store")
+    identity, artifact = _publish_test_artifact(store)
+    first = _quarantine_and_rebuild(store, identity)
+    second = _quarantine_and_rebuild(store, identity)
+    ready_inode = artifact.stat().st_ino
+    lock_inodes = {path.name: path.stat().st_ino for path in store.locks_dir.iterdir()}
+
+    assert store.cleanup_quarantined(first.name) is None
+    assert set(store.quarantine_dir.iterdir()) == {second}
+    assert artifact.stat().st_ino == ready_inode
+    hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+    assert hit.lease is not None
+    assert torch.equal(hit.lease.tensors["layer.weight"], _weight())
+    hit.lease.close()
+
+    _corrupt_test_payload(artifact)
+    assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.INVALID
+    deny = store.deny_dir / f"{identity.key}.json"
+    deny_bytes = deny.read_bytes()
+    assert store.cleanup_quarantined(second.name) is None
+    assert deny.read_bytes() == deny_bytes
+    assert artifact.stat().st_ino == ready_inode
+    assert {path.name: path.stat().st_ino for path in store.locks_dir.iterdir()} == lock_inodes
+    assert store.cleanup_quarantined(first.name) is None
+
+
+@pytest.mark.parametrize("active", ["build", "lease"])
+def test_cleanup_quarantined_refuses_active_key(tmp_path: Path, active: str) -> None:
+    store = _make_store(tmp_path / "store")
+    identity, _ = _publish_test_artifact(store)
+    quarantined = _quarantine_and_rebuild(store, identity)
+    owner: FileLock | HostWeightLease
+    if active == "build":
+        owner = FileLock(store._build_lock_path(identity.key), exclusive=True, deadline=None)
+        expected = FailureCode.ACTIVE_BUILD_TIMEOUT
+    else:
+        lease = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).lease
+        assert lease is not None
+        owner = lease
+        expected = FailureCode.ACTIVE_LEASE
+    try:
+        failure = store.cleanup_quarantined(quarantined.name)
+        assert failure is not None and failure.code is expected and failure.retryable
+        assert quarantined.is_dir()
+    finally:
+        owner.close()
+    assert store.cleanup_quarantined(quarantined.name) is None
+
+
+@pytest.mark.parametrize("recovery", ["cleanup", "build"])
+def test_cleanup_quarantined_retries_post_move_removal_failure(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, recovery: str
+) -> None:
+    store = _make_store(tmp_path / "store")
+    identity, artifact = _publish_test_artifact(store)
+    quarantined = _quarantine_and_rebuild(store, identity)
+    original_rmtree = shutil.rmtree
+
+    def fail_removal(path: str | os.PathLike[str]) -> None:
+        if Path(path).parent == store.quarantine_dir:
+            raise OSError(errno.EIO, "injected removal failure")
+        original_rmtree(path)
+
+    with monkeypatch.context() as fault:
+        fault.setattr(shutil, "rmtree", fail_removal)
+        failure = store.cleanup_quarantined(quarantined.name)
+    assert failure is not None and failure.code is FailureCode.QUARANTINE_FAILED and failure.retryable
+    assert not quarantined.exists()
+    assert len(list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))) == 1
+
+    if recovery == "cleanup":
+        assert store.cleanup_quarantined(quarantined.name) is None
+    else:
+        # A real build (not a warm hit) reconciles the same tombstone protocol.
+        _corrupt_test_payload(artifact)
+        result = store.get_or_build(
+            BuildRequest(identity),
+            FakeProducer(identity),
+            validation=ValidationLevel.FULL_CHECKSUM,
+            deadline=time.monotonic() + 10,
+        )
+        assert result.status is StoreStatus.BUILT
+        assert result.lease is not None
+        assert torch.equal(result.lease.tensors["layer.weight"], _weight())
+        result.lease.close()
+    assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+@pytest.mark.parametrize("after_removal", [False, True])
+def test_cleanup_quarantined_retries_parent_sync_failure(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, after_removal: bool
+) -> None:
+    store = _make_store(tmp_path / "store")
+    identity, _ = _publish_test_artifact(store)
+    quarantined = _quarantine_and_rebuild(store, identity)
+    original_sync = filesystem_store_module._fsync_directory
+
+    def fail_sync(path: Path) -> None:
+        if path == store.quarantine_dir and (
+            not after_removal or not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+        ):
+            raise OSError(errno.EIO, "injected quarantine sync failure")
+        original_sync(path)
+
+    with monkeypatch.context() as fault:
+        fault.setattr(filesystem_store_module, "_fsync_directory", fail_sync)
+        failure = store.cleanup_quarantined(quarantined.name)
+    assert failure is not None and failure.code is FailureCode.QUARANTINE_FAILED and failure.retryable
+    assert not quarantined.exists()
+    assert bool(list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))) is not after_removal
+
+    # An already absent selection must still retry the failed durability step.
+    with monkeypatch.context() as fault:
+        fault.setattr(filesystem_store_module, "_fsync_directory", fail_sync)
+        assert store.cleanup_quarantined(quarantined.name) is not None
+    assert store.cleanup_quarantined(quarantined.name) is None
+    assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+@pytest.mark.parametrize("entry_kind", ["symlink", "file"])
+def test_cleanup_quarantined_rejects_non_directory_entry(tmp_path: Path, entry_kind: str) -> None:
+    store = _make_store(tmp_path / "store")
+    identity, artifact = _publish_test_artifact(store)
+    entry = store.quarantine_dir / f"{identity.key}.validation_failure.test"
+    if entry_kind == "symlink":
+        entry.symlink_to(artifact, target_is_directory=True)
+    else:
+        entry.write_text("not an artifact directory")
+    failure = store.cleanup_quarantined(entry.name)
+    assert failure is not None and failure.code is FailureCode.QUARANTINE_FAILED
+    assert not failure.retryable
+    assert entry.exists()
+    hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+    assert hit.lease is not None
+    assert torch.equal(hit.lease.tensors["layer.weight"], _weight())
+    hit.lease.close()
+
+
+@pytest.mark.parametrize(
+    "name",
+    ["../escape", "/tmp/escape", "f" * 64, "g" * 64 + ".x", "f" * 64 + ".", "f" * 64 + ".x/../y", "f" * 64 + ".x\0"],
+)
+def test_cleanup_quarantined_rejects_invalid_selection_before_mutation(tmp_path: Path, name: str) -> None:
+    store = _make_store(tmp_path / "store")
+    locks = set(store.locks_dir.iterdir())
+    with pytest.raises(ValueError, match="inventory basename"):
+        store.cleanup_quarantined(name)
+    assert set(store.locks_dir.iterdir()) == locks
+
+
+def test_cleanup_quarantined_restores_bounded_store_recovery_headroom(tmp_path: Path) -> None:
+    store = _make_store(tmp_path / "store", capacity=CapacityPolicy(max_store_bytes=16384))
+    identity = _identity()
+    for _ in range(32):
+        result = store.get_or_build(
+            BuildRequest(identity),
+            FakeProducer(identity),
+            validation=ValidationLevel.FULL_CHECKSUM,
+            deadline=time.monotonic() + 10,
+        )
+        if result.status is StoreStatus.FAILED:
+            break
+        assert result.lease is not None
+        result.lease.close()
+        _corrupt_test_payload(store.artifacts_dir / identity.key)
+    assert result.failure is not None and result.failure.code is FailureCode.STORE_LIMIT_EXCEEDED
+    quarantined = [
+        entry for entry in store.inspect_domain().inventory if entry.state is ArtifactInventoryState.QUARANTINED
+    ]
+    assert len(quarantined) > 1
+    assert store.cleanup(identity) is None
+    still_full = store.get_or_build(
+        BuildRequest(identity),
+        FakeProducer(identity),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 10,
+    )
+    assert still_full.failure is not None and still_full.failure.code is FailureCode.STORE_LIMIT_EXCEEDED
+
+    assert store.cleanup_quarantined(quarantined[0].storage_name) is None
+    recovered = store.get_or_build(
+        BuildRequest(identity),
+        FakeProducer(identity),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 10,
+    )
+    assert recovered.status is StoreStatus.BUILT
+    assert recovered.lease is not None
+    assert torch.equal(recovered.lease.tensors["layer.weight"], _weight())
+    recovered.lease.close()
+
+
 def test_domain_capacity_policy_is_authoritative(tmp_path: Path) -> None:
     root = tmp_path / "store"
     _make_store(root, capacity=CapacityPolicy(max_store_bytes=4096))
diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py
index ec6b49f3c79..be44178abec 100644
--- a/vllm_omni/host_weight_runtime/filesystem/store.py
+++ b/vllm_omni/host_weight_runtime/filesystem/store.py
@@ -1470,43 +1470,45 @@ def _tree_bytes(root: Path) -> int:
                         total += file_stat.st_size
         return total
 
-    def cleanup(self, identity: WeightArtifactIdentity) -> HostWeightFailure | None:
-        """Explicitly remove an inactive artifact without waiting for leases."""
+    @contextlib.contextmanager
+    def _cleanup_locks(self, key: str) -> Iterator[None]:
+        """Exclude cooperative builders and leases without waiting for them."""
         try:
-            try:
-                build_lock = FileLock(
-                    self._build_lock_path(identity.key),
-                    exclusive=True,
-                    deadline=None,
-                    nonblocking=True,
-                )
-            except FileLockTimeoutError:
-                return _failure(
+            build_lock = FileLock(self._build_lock_path(key), exclusive=True, deadline=None, nonblocking=True)
+        except FileLockTimeoutError as exc:
+            raise HostWeightError(
+                _failure(
                     ResolutionStage.LIFECYCLE,
                     FailureCode.ACTIVE_BUILD_TIMEOUT,
-                    f"artifact {identity.key} has an active builder",
+                    f"artifact {key} has an active builder",
                     retryable=True,
                 )
-            with build_lock:
-                try:
-                    artifact_lock = FileLock(
-                        self._artifact_lock_path(identity.key),
-                        exclusive=True,
-                        deadline=None,
-                        nonblocking=True,
-                    )
-                except FileLockTimeoutError:
-                    return _failure(
+            ) from exc
+        with build_lock:
+            try:
+                artifact_lock = FileLock(self._artifact_lock_path(key), exclusive=True, deadline=None, nonblocking=True)
+            except FileLockTimeoutError as exc:
+                raise HostWeightError(
+                    _failure(
                         ResolutionStage.LIFECYCLE,
                         FailureCode.ACTIVE_LEASE,
-                        f"artifact {identity.key} has an active lease",
+                        f"artifact {key} has an active lease",
                         retryable=True,
                     )
-                with artifact_lock:
-                    self._repair_quarantine_transitions_locked(identity.key)
-                    self._quarantine_locked(identity.key, reason="cleanup")
-                    self._remove_cleanup_tombstones_locked(identity.key)
-                    self._remove_deny_locked(identity.key)
+                ) from exc
+            with artifact_lock:
+                yield
+
+    def cleanup(self, identity: WeightArtifactIdentity) -> HostWeightFailure | None:
+        """Explicitly remove an inactive artifact without waiting for leases."""
+        try:
+            with self._cleanup_locks(identity.key):
+                self._repair_quarantine_transitions_locked(identity.key)
+                self._quarantine_locked(identity.key, reason="cleanup")
+                self._remove_cleanup_tombstones_locked(identity.key)
+                self._remove_deny_locked(identity.key)
+        except HostWeightError as exc:
+            return exc.failure
         except OSError as exc:
             return _failure(
                 ResolutionStage.LIFECYCLE,
@@ -1516,6 +1518,50 @@ def cleanup(self, identity: WeightArtifactIdentity) -> HostWeightFailure | None:
             )
         return None
 
+    def cleanup_quarantined(self, storage_name: str) -> HostWeightFailure | None:
+        """Remove one quarantine inventory entry, preserving the current artifact.
+
+        A retry also finishes earlier cleanup tombstones for the same key.
+        Active builders or leases are reported rather than interrupted.
+        """
+        key, separator, suffix = storage_name.partition(".")
+        if (
+            not separator
+            or not suffix
+            or "\0" in storage_name
+            or _ARTIFACT_KEY_RE.fullmatch(key) is None
+            or Path(storage_name).name != storage_name
+        ):
+            raise ValueError("quarantined storage_name must be a single inventory basename with an artifact key")
+        try:
+            with self._cleanup_locks(key):
+                entry = self.quarantine_dir / storage_name
+                exists = _path_exists_without_following(entry)
+                if exists and not stat.S_ISDIR(entry.lstat().st_mode):
+                    return _failure(
+                        ResolutionStage.LIFECYCLE,
+                        FailureCode.QUARANTINE_FAILED,
+                        f"quarantine entry is not a directory: {storage_name}",
+                    )
+                self._repair_quarantine_transitions_locked(key)
+                if exists and not storage_name.startswith(f"{key}.cleanup."):
+                    tombstone = self.quarantine_dir / f"{key}.cleanup.{uuid.uuid4().hex}"
+                    _move_artifact_locked(entry, tombstone)
+                self._remove_cleanup_tombstones_locked(key)
+                # A previous attempt may have removed the tombstone but failed
+                # its parent sync. Retry that sync even when no entry remains.
+                _fsync_directory(self.quarantine_dir)
+        except HostWeightError as exc:
+            return exc.failure
+        except OSError as exc:
+            return _failure(
+                ResolutionStage.LIFECYCLE,
+                FailureCode.QUARANTINE_FAILED,
+                f"failed to remove quarantined entry {storage_name}: {exc}",
+                retryable=True,
+            )
+        return None
+
     def inspect_domain(self) -> DomainInspection:
         usage = shutil.disk_usage(self.root)
         inventory: list[ArtifactInventoryEntry] = []

From fd563dc3e4212c8c6d371da2e373070d03bf788a Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Sun, 6 Sep 2026 16:44:52 +0800
Subject: [PATCH 05/10] [Core] Add opt-in serialized HWR build admission

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit 90815cc86632562192156a79446beaf250e66021)

Signed-off-by: roG0d <baonudesifeizhai@gmail.com>
---
 docs/design/feature/host_weight_runtime.md    |  10 +
 docs/design/module/host_weight_runtime.md     |  38 ++-
 .../test_build_admission.py                   | 305 ++++++++++++++++++
 vllm_omni/host_weight_runtime/config.py       |   6 +
 .../host_weight_runtime/filesystem/store.py   |  39 ++-
 5 files changed, 394 insertions(+), 4 deletions(-)
 create mode 100644 tests/host_weight_runtime/test_build_admission.py

diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md
index ac7bc7b3f9d..a34d278052e 100644
--- a/docs/design/feature/host_weight_runtime.md
+++ b/docs/design/feature/host_weight_runtime.md
@@ -313,6 +313,16 @@ One process per exact identity owns a build; other workers wait and then acquire
 leases for the published artifact. Publication is invisible until all payloads
 and metadata are validated, hashed, fsynced, and atomically renamed.
 
+By default, different identities may build concurrently. A domain configured
+with `CapacityPolicy(build_admission="serialized", max_store_bytes=...)` admits
+only one producer/publication operation at a time. Warm hits remain concurrent,
+and admission uses the existing coordination deadline. All workers must agree
+on the serialized domain's schema-2 capacity policy; use a new root when opting
+in, since in-place migration of an active domain is unsupported. This prevents
+competing producers from racing allocation checks but does not provide a
+filesystem-wide byte quota or automatic stale-data reclamation. See the
+[module capacity contract](../module/host_weight_runtime.md#validation-and-capacity).
+
 `coordination_timeout_seconds` bounds domain-initialization and lookup/build
 lock acquisition. Store construction and each later resolution or publication
 operation have separate budgets from the same wait policy, rather than one
diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md
index 3b1b90b89e1..b060ed0ac9a 100644
--- a/docs/design/module/host_weight_runtime.md
+++ b/docs/design/module/host_weight_runtime.md
@@ -356,9 +356,41 @@ point-in-time evidence only; kernel locks are not a persistent owner registry.
 Capacity policy covers all store-owned bytes, including ready, temporary, and
 quarantined data. The local writer preflights `max_artifact_bytes`,
 `max_store_bytes`, and `min_free_bytes`, and preallocates payloads where the
-filesystem supports it. `ENOSPC` remains a normal typed store failure because
-concurrent preflight is inherently racy. Automatic eviction and strict
-concurrent reservations are not implemented.
+filesystem supports it. `ENOSPC` remains a normal typed store failure.
+
+`CapacityPolicy.build_admission` defaults to `concurrent`, preserving parallel
+production for different identities and best-effort capacity preflight. An
+explicit `serialized` policy, which requires `max_store_bytes`, admits one
+cooperative build/publication operation per domain through
+`locks/domain-build.lock`. Production lock order becomes domain admission,
+per-key build, then artifact. Validated warm hits bypass admission; queued
+same-key callers still recheck and join a completed artifact. Admission waiting
+uses the caller's existing coordination deadline and returns a retryable build
+timeout without interrupting its owner.
+
+The kernel releases admission on process exit. An explicit same-key retry uses
+the existing stale-temp recovery; files from unrelated dead builds continue to
+count toward capacity. A synchronous producer that hangs still requires
+external process supervision. Admission release errors are logged without
+revising an already determined publication result.
+
+Concurrent domains retain their exact schema-1 capacity document. Serialized
+domains use a schema-2 capacity document with `build_admission=serialized`, so
+mixed policies and old readers fail initialization rather than bypass admission.
+Opt into serialization using a new root with an agreed policy; in-place policy
+migration while old workers may exist is unsupported. For example:
+
+```python
+from vllm_omni.host_weight_runtime import CapacityPolicy
+
+capacity = CapacityPolicy(max_store_bytes=192 * 1024**3, build_admission="serialized")
+```
+
+Serialization prevents cooperative producers from racing each other's
+allocation checks. It is not a filesystem quota: control metadata, finalization,
+and outside writers can still affect total usage. Automatic eviction and
+parallel byte reservations are not implemented. This policy trades parallel
+cold-build throughput for predictable producer admission.
 
 For tmpfs, artifact bytes are host-memory consumption and may consume swap;
 they must not be reported operationally as ordinary disk capacity. The store
diff --git a/tests/host_weight_runtime/test_build_admission.py b/tests/host_weight_runtime/test_build_admission.py
new file mode 100644
index 00000000000..8bd06bbb6cf
--- /dev/null
+++ b/tests/host_weight_runtime/test_build_admission.py
@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU process and policy coverage for domain-wide producer admission."""
+
+from __future__ import annotations
+
+import errno
+import json
+import multiprocessing as mp
+import os
+import time
+from multiprocessing.connection import Connection
+from pathlib import Path
+from typing import Literal
+
+import pytest
+import torch
+
+from tests.host_weight_runtime.test_filesystem_store import FakeProducer, _identity, _make_store, _publish_test_artifact
+from vllm_omni.host_weight_runtime import (
+    BuildRequest,
+    CapacityPolicy,
+    FailureCode,
+    HostWeightError,
+    HostWeightRuntime,
+    HostWeightRuntimeConfig,
+    ProductionMetadata,
+    ResolutionOutcome,
+    RuntimeMode,
+    StoreResult,
+    StoreStatus,
+    TensorWriteSpec,
+    ValidationLevel,
+    WaitPolicy,
+    WeightArtifactIdentity,
+    WeightProductionSpec,
+)
+from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, lock_is_active
+from vllm_omni.host_weight_runtime.protocols import ArtifactWriter
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class AdmissionProducer:
+    def __init__(self, identity: WeightArtifactIdentity, control: Connection | None = None) -> None:
+        self._spec = FakeProducer(identity).spec
+        self.control = control
+
+    @property
+    def spec(self) -> WeightProductionSpec:
+        return self._spec
+
+    def produce(self, writer: ArtifactWriter) -> ProductionMetadata:
+        action = "produce"
+        if self.control is not None:
+            self.control.send("entered")
+            if not self.control.poll(60):
+                raise TimeoutError("test did not release producer")
+            action = self.control.recv()
+        spec = TensorWriteSpec("payload", (65536,), torch.bfloat16)
+        with writer.open_tensor_file("weights.safetensors", (spec,)) as output:
+            output.write_tensor("payload", torch.arange(65536, dtype=torch.float32).to(torch.bfloat16))
+            if action == "crash":
+                os._exit(27)
+        return ProductionMetadata("test-producer-v1", "test-restorer-v1")
+
+
+def _admission_process(root: Path, policy: CapacityPolicy, rank: int, control: Connection) -> None:
+    try:
+        store = _make_store(root, capacity=policy)
+        identity = _identity(tp_rank=rank)
+        original_lookup = store.lookup
+        first_lookup = True
+
+        def notify_lookup(
+            identity: WeightArtifactIdentity, *, validation: ValidationLevel, deadline: float | None = None
+        ) -> StoreResult:
+            nonlocal first_lookup
+            result = original_lookup(identity, validation=validation, deadline=deadline)
+            if first_lookup:
+                # Prove the caller observed a miss before the parent lets the
+                # active producer publish, independent of process scheduling.
+                control.send("looked_up")
+                first_lookup = False
+            return result
+
+        store.lookup = notify_lookup  # type: ignore[method-assign]
+        control.send("ready")
+        assert control.poll(60) and control.recv() == "start"
+        result = store.get_or_build(
+            BuildRequest(identity),
+            AdmissionProducer(identity, control),
+            validation=ValidationLevel.FULL_CHECKSUM,
+            deadline=time.monotonic() + 30,
+        )
+        if result.lease is not None:
+            result.lease.close()
+        control.send((result.status.value, result.failure.code.value if result.failure is not None else None))
+    finally:
+        control.close()
+
+
+@pytest.mark.parametrize("admission", ["concurrent", "serialized"])
+@pytest.mark.parametrize("same_key", [False, True])
+def test_cross_identity_admission_and_capacity(
+    tmp_path: Path, admission: Literal["concurrent", "serialized"], same_key: bool
+) -> None:
+    policy = CapacityPolicy(max_store_bytes=192 * 1024, build_admission=admission)
+    store = _make_store(tmp_path / "store", capacity=policy)
+    ctx = mp.get_context("spawn")
+    pairs = [ctx.Pipe() for _ in range(2)]
+    processes = [
+        ctx.Process(target=_admission_process, args=(store.root, policy, 0 if same_key else rank, pair[1]))
+        for rank, pair in enumerate(pairs)
+    ]
+    try:
+        for process, pair in zip(processes, pairs, strict=True):
+            process.start()
+            pair[1].close()
+        first, second = [pair[0] for pair in pairs]
+        for connection in (first, second):
+            assert connection.poll(60) and connection.recv() == "ready"
+        first.send("start")
+        assert first.poll(10) and first.recv() == "looked_up"
+        assert first.poll(10) and first.recv() == "entered"
+        second.send("start")
+        assert second.poll(10) and second.recv() == "looked_up"
+        if admission == "concurrent" and not same_key:
+            assert second.poll(10) and second.recv() == "entered"
+        else:
+            assert not second.poll(0.2), "a second producer entered while coordination should exclude it"
+        first.send("produce")
+        assert first.poll(30) and first.recv() == (StoreStatus.BUILT.value, None)
+        if same_key:
+            assert second.poll(30) and second.recv() == (StoreStatus.JOINED.value, None)
+        else:
+            if admission == "serialized":
+                assert second.poll(10) and second.recv() == "entered"
+            second.send("produce")
+            assert second.poll(30)
+            assert second.recv() == (StoreStatus.FAILED.value, FailureCode.STORE_LIMIT_EXCEEDED.value)
+        assert policy.max_store_bytes is not None
+        assert store.inspect_domain().store_bytes < policy.max_store_bytes
+        assert not list(store.tmp_dir.iterdir())
+        assert {path.name for path in store.artifacts_dir.iterdir()} == {_identity(tp_rank=0).key}
+        for process in processes:
+            process.join(10)
+            assert process.exitcode == 0
+    finally:
+        for process in processes:
+            if process.pid is not None:
+                if process.is_alive():
+                    process.kill()
+                process.join(5)
+                process.close()
+        for pair in pairs:
+            for connection in pair:
+                connection.close()
+
+
+def test_admission_timeout_preserves_owner_and_warm_hit_bypasses_lock(tmp_path: Path) -> None:
+    policy = CapacityPolicy(max_store_bytes=512 * 1024, build_admission="serialized")
+    store = _make_store(tmp_path / "store", capacity=policy)
+    identity, _ = _publish_test_artifact(store)
+    other = _identity(tp_rank=1)
+    lock_path = store.locks_dir / "domain-build.lock"
+    with FileLock(lock_path, exclusive=True, deadline=None):
+        hit = store.get_or_build(
+            BuildRequest(identity),
+            FakeProducer(identity),
+            validation=ValidationLevel.FULL_CHECKSUM,
+            deadline=time.monotonic() + 0.02,
+        )
+        assert hit.status is StoreStatus.HIT and hit.lease is not None
+        hit.lease.close()
+        timeout = store.get_or_build(
+            BuildRequest(other),
+            FakeProducer(other),
+            validation=ValidationLevel.FULL_CHECKSUM,
+            deadline=time.monotonic() + 0.02,
+        )
+        assert timeout.status is StoreStatus.TIMEOUT
+        assert timeout.failure is not None and timeout.failure.code is FailureCode.ACTIVE_BUILD_TIMEOUT
+        assert timeout.failure.retryable
+        assert lock_is_active(lock_path)
+    _publish_test_artifact(store, other)
+
+
+def test_producer_failure_releases_admission(tmp_path: Path) -> None:
+    store = _make_store(
+        tmp_path / "store", capacity=CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
+    )
+    identity = _identity()
+    failed = store.get_or_build(
+        BuildRequest(identity),
+        FakeProducer(identity, write_mode="incomplete"),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 5,
+    )
+    assert failed.status is StoreStatus.FAILED
+    assert not lock_is_active(store.locks_dir / "domain-build.lock")
+    _publish_test_artifact(store, identity)
+
+
+@pytest.mark.parametrize(
+    ("mode", "expected"),
+    [(RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK), (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED)],
+)
+def test_admission_timeout_obeys_runtime_mode(tmp_path: Path, mode: RuntimeMode, expected: ResolutionOutcome) -> None:
+    policy = CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
+    store = _make_store(tmp_path / "store", capacity=policy)
+    runtime = HostWeightRuntime.from_config(
+        HostWeightRuntimeConfig(mode=mode, domain=store.domain_policy, capacity=policy, wait=WaitPolicy(0.01))
+    )
+    identity = _identity()
+    with FileLock(store.locks_dir / "domain-build.lock", exclusive=True, deadline=None):
+        resolution = runtime.resolve(identity, producer=FakeProducer(identity))
+    assert resolution.report.outcome is expected
+    failure = resolution.report.attempts[-1].failure
+    assert failure is not None and failure.code is FailureCode.ACTIVE_BUILD_TIMEOUT
+
+
+def test_crashed_producer_releases_admission_and_same_key_retry_recovers(tmp_path: Path) -> None:
+    policy = CapacityPolicy(max_store_bytes=192 * 1024, build_admission="serialized")
+    store = _make_store(tmp_path / "store", capacity=policy)
+    ctx = mp.get_context("spawn")
+    parent, child = ctx.Pipe()
+    process = ctx.Process(target=_admission_process, args=(store.root, policy, 0, child))
+    process.start()
+    child.close()
+    try:
+        assert parent.poll(60) and parent.recv() == "ready"
+        parent.send("start")
+        assert parent.poll(10) and parent.recv() == "looked_up"
+        assert parent.poll(10) and parent.recv() == "entered"
+        parent.send("crash")
+        process.join(10)
+        assert process.exitcode == 27
+    finally:
+        if process.is_alive():
+            process.kill()
+        process.join(5)
+        process.close()
+        parent.close()
+    assert not lock_is_active(store.locks_dir / "domain-build.lock")
+    assert list(store.tmp_dir.iterdir())
+    identity = _identity(tp_rank=0)
+    recovered = store.get_or_build(
+        BuildRequest(identity),
+        AdmissionProducer(identity),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 10,
+    )
+    assert recovered.status is StoreStatus.BUILT and recovered.lease is not None
+    assert torch.equal(recovered.lease.tensors["payload"], torch.arange(65536, dtype=torch.float32).to(torch.bfloat16))
+    recovered.lease.close()
+    assert not list(store.tmp_dir.iterdir())
+
+
+def test_admission_release_error_preserves_publication(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    store = _make_store(
+        tmp_path / "store", capacity=CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
+    )
+    original_close = FileLock.close
+
+    def fail_admission_close(lock: FileLock) -> None:
+        original_close(lock)
+        if lock.path.name == "domain-build.lock":
+            raise OSError(errno.EIO, "injected admission release failure")
+
+    monkeypatch.setattr(FileLock, "close", fail_admission_close)
+    _publish_test_artifact(store)
+    assert not lock_is_active(store.locks_dir / "domain-build.lock")
+
+
+def test_admission_policy_is_authoritative_and_legacy_document_is_unchanged(tmp_path: Path) -> None:
+    concurrent = CapacityPolicy(max_store_bytes=16384)
+    serialized = CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
+    legacy = _make_store(tmp_path / "legacy", capacity=concurrent)
+    path = legacy.root / "domain-policy.json"
+    original = path.read_bytes()
+    assert json.loads(original) == {
+        "schema_version": 1,
+        "policy_version": 1,
+        "max_artifact_bytes": None,
+        "max_store_bytes": 16384,
+        "min_free_bytes": 0,
+        "eviction": "none",
+    }
+    with pytest.raises(HostWeightError, match="incompatible schema"):
+        _make_store(legacy.root, capacity=serialized)
+    assert path.read_bytes() == original
+    strict = _make_store(tmp_path / "serialized", capacity=serialized)
+    document = json.loads((strict.root / "domain-policy.json").read_bytes())
+    assert document["schema_version"] == 2 and document["build_admission"] == "serialized"
+    assert _make_store(strict.root, capacity=serialized).domain_uuid == strict.domain_uuid
+    with pytest.raises(HostWeightError, match="incompatible schema"):
+        _make_store(strict.root, capacity=concurrent)
+
+
+def test_admission_policy_validation() -> None:
+    with pytest.raises(ValueError, match="requires max_store_bytes"):
+        CapacityPolicy(build_admission="serialized")
+    with pytest.raises(ValueError, match="must be concurrent or serialized"):
+        CapacityPolicy(build_admission="unknown")  # type: ignore[arg-type]
diff --git a/vllm_omni/host_weight_runtime/config.py b/vllm_omni/host_weight_runtime/config.py
index 0bacd69d200..6ab6c7de771 100644
--- a/vllm_omni/host_weight_runtime/config.py
+++ b/vllm_omni/host_weight_runtime/config.py
@@ -8,6 +8,7 @@
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
+from typing import Literal
 
 
 class RuntimeMode(str, Enum):
@@ -66,6 +67,7 @@ class CapacityPolicy:
     max_store_bytes: int | None = None
     min_free_bytes: int = 0
     eviction: str = "none"
+    build_admission: Literal["concurrent", "serialized"] = "concurrent"
 
     def __post_init__(self) -> None:
         for name, value in (
@@ -78,6 +80,10 @@ def __post_init__(self) -> None:
             raise ValueError("min_free_bytes must not be negative")
         if not isinstance(self.eviction, str) or self.eviction != "none":
             raise ValueError("automatic host weight eviction is not implemented")
+        if self.build_admission not in ("concurrent", "serialized"):
+            raise ValueError("build_admission must be concurrent or serialized")
+        if self.build_admission == "serialized" and self.max_store_bytes is None:
+            raise ValueError("serialized build admission requires max_store_bytes")
 
 
 @dataclass(frozen=True)
diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py
index be44178abec..34403626676 100644
--- a/vllm_omni/host_weight_runtime/filesystem/store.py
+++ b/vllm_omni/host_weight_runtime/filesystem/store.py
@@ -629,7 +629,7 @@ def _initialize_domain(self, *, deadline: float) -> None:
                 self.policy_version = 1
 
     def _capacity_document(self, *, policy_version: int) -> dict[str, object]:
-        return {
+        document: dict[str, object] = {
             "schema_version": 1,
             "policy_version": policy_version,
             "max_artifact_bytes": self.capacity_policy.max_artifact_bytes,
@@ -637,6 +637,9 @@ def _capacity_document(self, *, policy_version: int) -> dict[str, object]:
             "min_free_bytes": self.capacity_policy.min_free_bytes,
             "eviction": self.capacity_policy.eviction,
         }
+        if self.capacity_policy.build_admission == "serialized":
+            document.update(schema_version=2, build_admission="serialized")
+        return document
 
     def _build_lock_path(self, key: str) -> Path:
         return self.locks_dir / f"{key}.build.lock"
@@ -1015,6 +1018,40 @@ def get_or_build(
         if initial.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}:
             return initial
 
+        if self.capacity_policy.build_admission == "concurrent":
+            return self._build_after_miss(identity, producer, validation=validation, deadline=deadline)
+        try:
+            admission = FileLock(self.locks_dir / "domain-build.lock", exclusive=True, deadline=deadline)
+        except FileLockTimeoutError as exc:
+            return StoreResult(
+                StoreStatus.TIMEOUT,
+                failure=_failure(
+                    ResolutionStage.PRODUCTION, FailureCode.ACTIVE_BUILD_TIMEOUT, str(exc), retryable=True
+                ),
+            )
+        except OSError as exc:
+            return StoreResult(
+                StoreStatus.FAILED,
+                failure=_failure(ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE, str(exc), retryable=True),
+            )
+        try:
+            return self._build_after_miss(identity, producer, validation=validation, deadline=deadline)
+        finally:
+            try:
+                admission.close()
+            except OSError:
+                # Releasing coordination must not revise the completed build's
+                # publication outcome or replace its primary failure.
+                logger.warning("Failed to release domain build admission under %s", self.root, exc_info=True)
+
+    def _build_after_miss(
+        self,
+        identity: WeightArtifactIdentity,
+        producer: WeightProducer,
+        *,
+        validation: ValidationLevel,
+        deadline: float,
+    ) -> StoreResult:
         try:
             build_lock = FileLock(self._build_lock_path(identity.key), exclusive=True, deadline=deadline)
         except FileLockTimeoutError as exc:

From e974eeeec02255d1cc864ab375a45899b8f9db5f Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Sun, 6 Sep 2026 16:54:01 +0800
Subject: [PATCH 06/10] [Core] Expose HWR payload validation in diffusion
 startup

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit d989aa6fe2af1acf39e6996270f566e23e0e333e)
---
 docs/design/feature/host_weight_runtime.md    | 20 ++++++
 tests/config/test_omni_config.py              | 31 ++++++++
 .../model_loader/test_diffusers_loader.py     | 72 +++++++++++++++++++
 .../test_async_omni_diffusion_config.py       |  4 ++
 vllm_omni/config/omni_config.py               |  2 +
 vllm_omni/config/stage_config.py              |  1 +
 vllm_omni/diffusion/data.py                   |  8 ++-
 .../model_loader/host_weight_loader.py        | 10 +++
 vllm_omni/engine/arg_utils.py                 |  1 +
 vllm_omni/engine/async_omni_engine.py         |  1 +
 vllm_omni/entrypoints/cli/serve.py            |  9 +++
 11 files changed, 158 insertions(+), 1 deletion(-)

diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md
index a34d278052e..8427fa462a9 100644
--- a/docs/design/feature/host_weight_runtime.md
+++ b/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
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index f91e6af42b9..fddb952a125 100644
--- a/tests/config/test_omni_config.py
+++ b/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")
diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py
index 555cb19bb4d..cc306776242 100644
--- a/tests/diffusion/model_loader/test_diffusers_loader.py
+++ b/tests/diffusion/model_loader/test_diffusers_loader.py
@@ -260,6 +260,78 @@ def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
     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",
diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py
index 616d5486eb4..53a1f778bd5 100644
--- a/tests/entrypoints/test_async_omni_diffusion_config.py
+++ b/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
diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index 1a1c0c3ae8e..0595507cc39 100644
--- a/vllm_omni/config/omni_config.py
+++ b/vllm_omni/config/omni_config.py
@@ -739,6 +739,7 @@ class keeps the diffusion-only knobs from ``OmniDiffusionConfig`` without
     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 @@ def __post_init__(self) -> None:
         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,
diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py
index 06840e55ac1..ae4ee6fc49a 100644
--- a/vllm_omni/config/stage_config.py
+++ b/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
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 00764ea46d8..298fa8d97d5 100644
--- a/vllm_omni/diffusion/data.py
+++ b/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 @@ def __post_init__(self):
         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,
diff --git a/vllm_omni/diffusion/model_loader/host_weight_loader.py b/vllm_omni/diffusion/model_loader/host_weight_loader.py
index 4b98535558d..01427909d91 100644
--- a/vllm_omni/diffusion/model_loader/host_weight_loader.py
+++ b/vllm_omni/diffusion/model_loader/host_weight_loader.py
@@ -26,6 +26,7 @@
 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 @@ def _resolve_hwr(
             HostWeightLeaseCarrier,
             HostWeightRuntime,
             HostWeightRuntimeConfig,
+            IntegrityPolicy,
             ProductionPolicy,
             ResolutionOutcome,
             RuntimeMode,
             StorageDomainPolicy,
+            ValidationLevel,
         )
         from vllm_omni.host_weight_runtime.filesystem import detect_storage_class
 
@@ -363,6 +368,11 @@ def _resolve_hwr(
             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,
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 812f4c3e7d2..bb298b967ae 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/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
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index 608d605e9b1..989d470176e 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -1145,6 +1145,7 @@ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list:
             "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": (
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index ea56acb1af3..48354d8943c 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -770,6 +770,15 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu
                 "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,

From fa2c0ef057252838b461e13ce2a2c581d75bdb21 Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Sun, 6 Sep 2026 16:57:46 +0800
Subject: [PATCH 07/10] [Misc] Add standalone safetensors retention diagnostic

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit 86166fd988887e5592e08d7ecf9fa23b70865b82)
---
 benchmarks/host_weight_runtime/README.md      |  49 ++++++++
 .../safetensors_retention.py                  | 117 ++++++++++++++++++
 docs/design/feature/host_weight_runtime.md    |   5 +
 3 files changed, 171 insertions(+)
 create mode 100644 benchmarks/host_weight_runtime/README.md
 create mode 100644 benchmarks/host_weight_runtime/safetensors_retention.py

diff --git a/benchmarks/host_weight_runtime/README.md b/benchmarks/host_weight_runtime/README.md
new file mode 100644
index 00000000000..9a834ff1ff8
--- /dev/null
+++ b/benchmarks/host_weight_runtime/README.md
@@ -0,0 +1,49 @@
+# 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 16-element float32 payload. 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 comparison, run each mode twice in fresh processes with the same controls:
+
+```bash
+for repetition in 1 2; 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" > "$mode-$repetition.json"
+  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.
diff --git a/benchmarks/host_weight_runtime/safetensors_retention.py b/benchmarks/host_weight_runtime/safetensors_retention.py
new file mode 100644
index 00000000000..bcdb6427a7a
--- /dev/null
+++ b/benchmarks/host_weight_runtime/safetensors_retention.py
@@ -0,0 +1,117 @@
+# 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("--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(16, 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(16, 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()
diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md
index 8427fa462a9..28a01974802 100644
--- a/docs/design/feature/host_weight_runtime.md
+++ b/docs/design/feature/host_weight_runtime.md
@@ -42,6 +42,11 @@ V1 does not include:
 
 ## Motivation and use cases
 
+For investigating CPU memory retained by dependency tensor materialization,
+see the [standalone safetensors diagnostic](../../../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

From add3a030b04602fbd7c0d9532add4c2e2b584fae Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Sun, 6 Sep 2026 21:53:50 +0800
Subject: [PATCH 08/10] [Misc] Make retention probe tensor size configurable

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit 743953b821e5352f3958c781811d53c28702584b)
---
 benchmarks/host_weight_runtime/README.md      | 20 +++++++++++++------
 .../safetensors_retention.py                  |  7 +++++--
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/benchmarks/host_weight_runtime/README.md b/benchmarks/host_weight_runtime/README.md
index 9a834ff1ff8..8998111518f 100644
--- a/benchmarks/host_weight_runtime/README.md
+++ b/benchmarks/host_weight_runtime/README.md
@@ -2,7 +2,11 @@
 
 `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 16-element float32 payload. HWR itself calls `get_tensor()` when
+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.
 
@@ -18,14 +22,18 @@ CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
   --mode get_tensor --iterations 10000 --sample-every 5000 > feasibility.json
 ```
 
-For a comparison, run each mode twice in fresh processes with the same controls:
+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 mode in get_tensor reuse; do
-    CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
-      benchmarks/host_weight_runtime/safetensors_retention.py \
-      --mode "$mode" > "$mode-$repetition.json"
+  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
 ```
diff --git a/benchmarks/host_weight_runtime/safetensors_retention.py b/benchmarks/host_weight_runtime/safetensors_retention.py
index bcdb6427a7a..55795e7ef80 100644
--- a/benchmarks/host_weight_runtime/safetensors_retention.py
+++ b/benchmarks/host_weight_runtime/safetensors_retention.py
@@ -43,6 +43,9 @@ def sample(phase: str, iterations: int, loop_seconds: float) -> dict[str, object
 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)
@@ -61,10 +64,10 @@ def main() -> None:
     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(16, dtype=torch.float32)}, str(path))
+        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(16, dtype=torch.float32))
+            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

From 6d8bd67263d6ee695db7fcaf81a64c2406186a00 Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Sun, 6 Sep 2026 17:06:41 +0800
Subject: [PATCH 09/10] [Core] Add explicit HWR registration bypass and
 progress logs

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit 729ccd7f11df83800b86c402ee1881c28a595c8c)

Signed-off-by: roG0d <baonudesifeizhai@gmail.com>
---
 .../distributed_layerwise_offload.md          | 15 +++++++++
 tests/config/test_omni_config.py              | 33 +++++++++++++++++++
 .../offloader/test_cuda_host_registration.py  | 27 +++++++++++++++
 .../test_distributed_layerwise_backend.py     | 19 ++++++++---
 .../test_async_omni_diffusion_config.py       |  4 +++
 vllm_omni/config/omni_config.py               |  2 ++
 vllm_omni/config/stage_config.py              |  1 +
 vllm_omni/diffusion/data.py                   |  9 +++++
 vllm_omni/diffusion/offloader/base.py         |  6 ++++
 .../offloader/cuda_host_registration.py       | 11 ++++++-
 .../distributed_layerwise_backend.py          |  9 +++++
 vllm_omni/engine/arg_utils.py                 |  1 +
 vllm_omni/engine/async_omni_engine.py         |  1 +
 vllm_omni/entrypoints/cli/serve.py            | 11 +++++++
 14 files changed, 144 insertions(+), 5 deletions(-)

diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md
index 730d5547252..315a1defce8 100644
--- a/docs/design/feature/offloader/distributed_layerwise_offload.md
+++ b/docs/design/feature/offloader/distributed_layerwise_offload.md
@@ -289,6 +289,13 @@ store, tensor ownership, or H2D payload. On success, each tensor view copies
 directly into the existing rotating HBM block buffers and the two private host
 staging slots are not allocated.
 
+`--dlo-host-registration-mode` defaults to `auto`, preserving this behavior.
+Select `disabled` to bypass HWR mapping registration and use bounded host staging
+regardless of mapping size or budget. This leaves the existing pin-memory policy
+for staging buffers unchanged. The equivalent offline/stage configuration field
+is `dlo_host_registration_mode`; an explicit disable requires enabled HWR with
+no-AllGather DLO.
+
 `--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling
 over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned
 memory policy, unsupported platform/capability, over-budget mapping, or fully
@@ -296,6 +303,14 @@ rolled-back registration error selects the existing two-slot staging path. A
 partial registration that cannot be rolled back aborts startup because closing
 the lease would unmap memory still owned by the platform.
 
+Automatic registration logs its entry with device and budget, then each CUDA
+region's ordinal/total and byte size before calling `cudaHostRegister`. A
+completion message follows each successful call. These identify the last
+entered boundary if progress stops; they do not impose a timeout. A synchronous
+CUDA call cannot be safely cancelled by a Python timer, and a stuck process
+still requires external supervision. The explicit disable option avoids this
+registration path; it does not repair the underlying driver hang.
+
 Direct checkpoint mmap remains unchanged and continues to use staging. It may
 require loader-owned per-block transforms, while the HWR artifact already
 contains final runtime bytes. DLO AllGather never receives an HWR final-layout
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index fddb952a125..8feed4904b1 100644
--- a/tests/config/test_omni_config.py
+++ b/tests/config/test_omni_config.py
@@ -224,6 +224,39 @@ def test_hwr_validation_rejects_unsupported_values(validation):
             config_class(host_weight_runtime_validation=validation)
 
 
+@pytest.mark.parametrize("registration_mode", ["auto", "disabled"])
+def test_hwr_registration_policy_reaches_offloader(tmp_path: Path, registration_mode: str):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.offloader.base import OffloadConfig
+
+    options = {
+        "enable_distributed_layerwise_offload": True,
+        "dlo_use_allgather": False,
+        "host_weight_runtime_mode": "preferred",
+        "host_weight_runtime_root": str(tmp_path / "unused-store"),
+        "dlo_host_registration_mode": registration_mode,
+    }
+    config = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2", cli_overrides=options)
+    stage = config.stage_by_id(0)
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.diffusion_config.dlo_host_registration_mode == registration_mode
+    offline = OmniDiffusionConfig(**options)
+    transport = OffloadConfig.from_od_config(offline)
+    assert transport.dlo_host_registration_mode == registration_mode
+    assert transport.pin_cpu_memory
+    assert transport.dlo_host_registration_limit_gib == 0.0
+    assert not (tmp_path / "unused-store").exists()
+
+
+@pytest.mark.parametrize("registration_mode", ["unknown", None, "disabled"])
+def test_hwr_registration_policy_rejects_invalid_or_ineligible_config(registration_mode):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    for config_class in (OmniDiffusionConfig, omni_config_module._DiffusionConfigProjection):
+        with pytest.raises(ValueError, match="dlo_host_registration_mode"):
+            config_class(dlo_host_registration_mode=registration_mode)
+
+
 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")
diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py
index a3efbb40a12..bbed5a3e3a1 100644
--- a/tests/diffusion/offloader/test_cuda_host_registration.py
+++ b/tests/diffusion/offloader/test_cuda_host_registration.py
@@ -124,6 +124,33 @@ def test_registration_rejects_writable_or_over_budget_before_cuda(monkeypatch: p
     assert runtime.registered == []
 
 
+@pytest.mark.parametrize("second_result", [0, 7])
+def test_registration_logs_entry_before_each_runtime_call(monkeypatch: pytest.MonkeyPatch, second_result: int) -> None:
+    runtime = _FakeRuntime([0, second_result])
+    messages: list[str] = []
+    monkeypatch.setattr(registration_module.logger, "info", lambda message, *args: messages.append(message % args))
+    original_register = runtime.cudaHostRegister
+
+    def register(pointer: int, size: int, flags: int) -> int:
+        ordinal = len(runtime.registered) + 1
+        assert messages[-1] == f"Registering HWR host range {ordinal}/2: {size} bytes"
+        return original_register(pointer, size, flags)
+
+    monkeypatch.setattr(runtime, "cudaHostRegister", register)
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+    regions = (_region("first", 0x1000, 4096), _region("second", 0x3000, 4096))
+    if second_result:
+        with pytest.raises(HostRegistrationError, match="error-7"):
+            CudaHostRegistration.create(regions, max_bytes=None)
+        assert runtime.unregistered == [0x1000]
+    else:
+        registration = CudaHostRegistration.create(regions, max_bytes=None)
+        assert registration.close() == ()
+    assert sum(message.startswith("Registering ") for message in messages) == 2
+    assert sum(message.startswith("Registered ") for message in messages) == (1 if second_result else 2)
+
+
 def test_registration_requires_read_only_capability(monkeypatch: pytest.MonkeyPatch) -> None:
     runtime = _FakeRuntime([0], read_only_supported=False)
     monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py
index 05d7d93b583..0a838a36952 100644
--- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py
+++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py
@@ -96,7 +96,7 @@ def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime
         hook.offload_layer()
         assert not hook.is_materialized
 
-    def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch):
+    def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch, patched_offload_runtime):
         current_block = nn.Linear(2, 2)
         next_block = nn.Linear(2, 2)
         expected = {name: tensor.detach().clone() for name, tensor in next_block.state_dict().items()}
@@ -1373,10 +1373,14 @@ def test_rank_local_disable_drains_pending_transfers(self, patched_offload_runti
 
         synchronize.assert_called()
 
-    def test_hwr_registration_failure_falls_back_to_bounded_staging(
+    @pytest.mark.parametrize("registration_mode", ["auto", "disabled"])
+    @pytest.mark.parametrize("limit_gib", [0.0, 1.5])
+    def test_hwr_registration_selection_uses_bounded_staging(
         self,
         monkeypatch,
         patched_offload_runtime,
+        registration_mode,
+        limit_gib,
     ):
         plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
         backend = DistributedLayerwiseOffloadBackend(
@@ -1385,6 +1389,8 @@ def test_hwr_registration_failure_falls_back_to_bounded_staging(
                 pin_cpu_memory=True,
                 dp_size=1,
                 dlo_use_allgather=False,
+                dlo_host_registration_mode=registration_mode,
+                dlo_host_registration_limit_gib=limit_gib,
             ),
             torch.device("cpu"),
             host_weight_plan=plan,
@@ -1394,12 +1400,14 @@ def test_hwr_registration_failure_falls_back_to_bounded_staging(
 
         def fail_registration(*args, **kwargs):
             del args, kwargs
+            assert registration_mode == "auto", "disabled mode entered platform registration"
             raise HostRegistrationError("registration unavailable in CPU test")
 
         def allocate_unpinned_staging(hooks, resident_group=None):
             # Registration selection still observes pin_cpu_memory=True. The
             # CPU-only test avoids asking the host for CUDA-pinned allocations.
             for hook in hooks:
+                assert hook.pin_memory, "registration policy disabled staging pinning"
                 hook.pin_memory = False
             return original_staging_allocator(hooks, resident_group)
 
@@ -1412,6 +1420,7 @@ def allocate_unpinned_staging(hooks, resident_group=None):
         assert backend._using_rank_local_mmap
         assert not backend._using_registered_mmap
         assert backend._host_registration is None
+        assert backend.config.pin_cpu_memory
         assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
         assert not lease.closed
 
@@ -2630,7 +2639,9 @@ def fail_second_hook(*args, **kwargs):
             assert registry is None or registry.get_hook("distributed_layerwise_offload") is None
             torch.testing.assert_close(block.weight, expected)
 
-    def test_multirank_enable_failure_cleanup_skips_restore_collective(self, monkeypatch, mocker):
+    def test_multirank_enable_failure_cleanup_skips_restore_collective(
+        self, monkeypatch, mocker, patched_offload_runtime
+    ):
         backend = DistributedLayerwiseOffloadBackend(
             OffloadConfig(
                 strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
@@ -2811,7 +2822,7 @@ def test_encoder_allgather_rejects_undeclared_replication(self, patched_offload_
         with pytest.raises(ValueError, match="not declared replicated"):
             backend.enable(pipeline)
 
-    def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self):
+    def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self, patched_offload_runtime):
         """Every rank must reject an unsafe encoder group, including stub ranks."""
         backend = DistributedLayerwiseOffloadBackend(
             OffloadConfig(
diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py
index 53a1f778bd5..ed9fcff2908 100644
--- a/tests/entrypoints/test_async_omni_diffusion_config.py
+++ b/tests/entrypoints/test_async_omni_diffusion_config.py
@@ -509,6 +509,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
             "preferred",
             "--host-weight-runtime-root",
             "/var/cache/vllm-omni/hwr",
+            "--dlo-host-registration-mode",
+            "disabled",
             "--dlo-host-registration-limit-gib",
             "80",
         ]
@@ -526,6 +528,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
     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
+    assert explicit_kwargs["dlo_host_registration_mode"] == "disabled"
+    assert engine_args["dlo_host_registration_mode"] == "disabled"
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():
diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index 0595507cc39..9d0108bf178 100644
--- a/vllm_omni/config/omni_config.py
+++ b/vllm_omni/config/omni_config.py
@@ -740,6 +740,7 @@ class keeps the diffusion-only knobs from ``OmniDiffusionConfig`` without
     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_mode: str = "auto"
     dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
@@ -897,6 +898,7 @@ def __post_init__(self) -> None:
         )
         self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
             limit_gib=self.dlo_host_registration_limit_gib,
+            mode=self.dlo_host_registration_mode,
             enable_dlo=self.enable_distributed_layerwise_offload,
             use_allgather=self.dlo_use_allgather,
             hwr_mode=self.host_weight_runtime_mode,
diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py
index ae4ee6fc49a..cd90e4c490a 100644
--- a/vllm_omni/config/stage_config.py
+++ b/vllm_omni/config/stage_config.py
@@ -484,6 +484,7 @@ class StageDeployConfig:
     host_weight_runtime_mode: str | None = None
     host_weight_runtime_root: str | None = None
     host_weight_runtime_validation: str | None = None
+    dlo_host_registration_mode: str | None = None
     dlo_host_registration_limit_gib: float | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 298fa8d97d5..ce11fbe3b1f 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -123,8 +123,15 @@ def validate_dlo_host_registration_options(
     enable_dlo: bool,
     use_allgather: bool,
     hwr_mode: object,
+    mode: object = "auto",
 ) -> float:
     """Validate the optional transport budget without probing CUDA or HWR."""
+    if mode not in ("auto", "disabled"):
+        raise ValueError("dlo_host_registration_mode must be auto or disabled")
+    if mode == "disabled" and (not enable_dlo or use_allgather or hwr_mode == "disabled"):
+        raise ValueError(
+            "disabled dlo_host_registration_mode requires enabled no-AllGather DLO and Host Weight Runtime"
+        )
     if not isinstance(limit_gib, (int, float, str)):
         raise TypeError(f"dlo_host_registration_limit_gib must be a number; got {type(limit_gib).__name__}")
     value = float(limit_gib)
@@ -838,6 +845,7 @@ class OmniDiffusionConfig:
     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_mode: str = "auto"
     dlo_host_registration_limit_gib: float = 0.0
 
     pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
@@ -1250,6 +1258,7 @@ def __post_init__(self):
         )
         self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
             limit_gib=self.dlo_host_registration_limit_gib,
+            mode=self.dlo_host_registration_mode,
             enable_dlo=self.enable_distributed_layerwise_offload,
             use_allgather=self.dlo_use_allgather,
             hwr_mode=self.host_weight_runtime_mode,
diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py
index bb047a7ebc0..c1afb2c17e1 100644
--- a/vllm_omni/diffusion/offloader/base.py
+++ b/vllm_omni/diffusion/offloader/base.py
@@ -73,12 +73,15 @@ class OffloadConfig:
     # Optional per-worker ceiling for registering an HWR mmap. Zero means no
     # additional ceiling; pin_cpu_memory controls whether registration is tried.
     dlo_host_registration_limit_gib: float = 0.0
+    dlo_host_registration_mode: str = "auto"
     # ``None`` preserves the model's legacy plan-driven component topology;
     # a frozenset is an explicit compact-API selection.
     components: frozenset[str] | None = None
     dlo_transfers: dict[str, DLOTransfer] | None = None
 
     def __post_init__(self) -> None:
+        if self.dlo_host_registration_mode not in ("auto", "disabled"):
+            raise ValueError("dlo_host_registration_mode must be auto or disabled")
         if self.components is not None:
             self.components = parse_offload_components(self.components)
         if self.dlo_transfers is None:
@@ -174,7 +177,9 @@ def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
         dlo_transfers = dict(resolved.transfers)
         dlo_resident_layers = resolved.resident_layers
         dit_uses_allgather = resolved.uses_allgather(DIT_COMPONENT)
+        registration_mode = getattr(od_config, "dlo_host_registration_mode", "auto")
         dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
+            mode=registration_mode,
             limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0),
             enable_dlo=enable_distributed_layerwise_offload,
             use_allgather=dit_uses_allgather,
@@ -210,6 +215,7 @@ def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
             dlo_use_allgather=dit_uses_allgather,
             dlo_resident_layers=dlo_resident_layers,
             dlo_host_registration_limit_gib=dlo_host_registration_limit_gib,
+            dlo_host_registration_mode=registration_mode,
             components=components,
             dlo_transfers=dlo_transfers,
         )
diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py
index 418b5d9fb0d..b5347b6c4f4 100644
--- a/vllm_omni/diffusion/offloader/cuda_host_registration.py
+++ b/vllm_omni/diffusion/offloader/cuda_host_registration.py
@@ -6,12 +6,14 @@
 
 import ctypes
 import mmap
+import time
 from collections import defaultdict
 from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Protocol
 
 import torch
+from vllm.logger import init_logger
 
 from vllm_omni.host_weight_runtime import MappedHostRegion
 
@@ -21,6 +23,8 @@
     HostRegistrationError,
 )
 
+logger = init_logger(__name__)
+
 _CUDA_HOST_REGISTER_READ_ONLY = 0x08
 _CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113
 
@@ -187,7 +191,9 @@ def create(
 
         registered: list[_AddressRange] = []
         try:
-            for region in mapped:
+            for index, region in enumerate(mapped, start=1):
+                logger.info("Registering HWR host range %d/%d: %d bytes", index, len(mapped), region.size)
+                started = time.perf_counter()
                 error = runtime.cudaHostRegister(region.start, region.size, _CUDA_HOST_REGISTER_READ_ONLY)
                 if int(error) != 0:
                     raise HostRegistrationError(
@@ -195,6 +201,9 @@ def create(
                         f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}"
                     )
                 registered.append(region)
+                logger.info(
+                    "Registered HWR host range %d/%d in %.3f s", index, len(mapped), time.perf_counter() - started
+                )
         except Exception as exc:
             rollback_errors: list[str] = []
             rollback_failed: list[_AddressRange] = []
diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py
index e626d87a543..f04060e0b6e 100644
--- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py
+++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py
@@ -1080,6 +1080,9 @@ def _try_register_hwr_mmap(self, source_tensors: tuple[torch.Tensor, ...]) -> bo
         lease = self._host_weight_lease
         if lease is None:
             return False
+        if self.config.dlo_host_registration_mode == "disabled":
+            logger.info("HWR mmap registration disabled by transport policy; using bounded host staging")
+            return False
         if not self.config.pin_cpu_memory:
             logger.info("HWR mmap registration disabled by pin_cpu_memory=False; using bounded host staging")
             return False
@@ -1089,6 +1092,12 @@ def _try_register_hwr_mmap(self, source_tensors: tuple[torch.Tensor, ...]) -> bo
 
         limit_gib = self.config.dlo_host_registration_limit_gib
         max_bytes = int(limit_gib * 1024**3) if limit_gib > 0 else None
+        logger.info(
+            "Starting HWR mmap registration on %s: %d source range(s), budget_bytes=%s",
+            self.device,
+            len(lease.mapped_regions),
+            max_bytes,
+        )
         started = time.perf_counter()
         try:
             registration = register_host_mappings(
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index bb298b967ae..7f918ece33e 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -583,6 +583,7 @@ class OrchestratorArgs:
     host_weight_runtime_mode: str = "disabled"
     host_weight_runtime_root: str | None = None
     host_weight_runtime_validation: str = "manifest_and_metadata"
+    dlo_host_registration_mode: str = "auto"
     dlo_host_registration_limit_gib: float = 0.0
     boundary_ratio: float | None = None
     flow_shift: float | None = None
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index 989d470176e..4c32297ab8b 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -1146,6 +1146,7 @@ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list:
             "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_mode": kwargs.get("dlo_host_registration_mode", "auto"),
             "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": (
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index 48354d8943c..01961acc69f 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -789,6 +789,17 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu
                 "the same persistent path for population and serving."
             ),
         )
+        omni_config_group.add_argument(
+            "--dlo-host-registration-mode",
+            choices=("auto", "disabled"),
+            default="auto",
+            help=(
+                "HWR transport registration policy for eligible no-AllGather DLO. "
+                "auto attempts registered direct H2D under the existing pin-memory and budget policy; "
+                "disabled selects bounded staging without registering the HWR mappings. "
+                "This preserves the pin-memory policy for staging buffers."
+            ),
+        )
         omni_config_group.add_argument(
             "--dlo-host-registration-limit-gib",
             type=float,

From 688446d23af1dda5e1c5d40940b0ccca162252a7 Mon Sep 17 00:00:00 2001
From: roG0d <baonudesifeizhai@gmail.com>
Date: Sun, 6 Sep 2026 16:51:41 +0000
Subject: [PATCH 10/10] [Test] Cover combined HWR integrity and registration
 policies

Signed-off-by: roG0d <baonudesifeizhai@gmail.com>
---
 tests/config/test_omni_config.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index 8feed4904b1..e41e127900d 100644
--- a/tests/config/test_omni_config.py
+++ b/tests/config/test_omni_config.py
@@ -225,7 +225,8 @@ def test_hwr_validation_rejects_unsupported_values(validation):
 
 
 @pytest.mark.parametrize("registration_mode", ["auto", "disabled"])
-def test_hwr_registration_policy_reaches_offloader(tmp_path: Path, registration_mode: str):
+@pytest.mark.parametrize("validation", ["manifest_and_metadata", "full_checksum"])
+def test_hwr_registration_policy_reaches_offloader(tmp_path: Path, registration_mode: str, validation: str):
     from vllm_omni.diffusion.data import OmniDiffusionConfig
     from vllm_omni.diffusion.offloader.base import OffloadConfig
 
@@ -235,12 +236,15 @@ def test_hwr_registration_policy_reaches_offloader(tmp_path: Path, registration_
         "host_weight_runtime_mode": "preferred",
         "host_weight_runtime_root": str(tmp_path / "unused-store"),
         "dlo_host_registration_mode": registration_mode,
+        "host_weight_runtime_validation": validation,
     }
     config = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2", cli_overrides=options)
     stage = config.stage_by_id(0)
     assert isinstance(stage, VllmOmniDiffusionStageConfig)
     assert stage.diffusion_config.dlo_host_registration_mode == registration_mode
+    assert stage.diffusion_config.host_weight_runtime_validation == validation
     offline = OmniDiffusionConfig(**options)
+    assert offline.host_weight_runtime_validation == validation
     transport = OffloadConfig.from_od_config(offline)
     assert transport.dlo_host_registration_mode == registration_mode
     assert transport.pin_cpu_memory