[Core] Add explicit HWR registration bypass and progress logs

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

docs/design/feature/host_weight_runtime.md

@@ -313,16 +313,6 @@ 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 filesystem lock acquisition. It does not
 cancel synchronous validation, a producer that has already started, or atomic
 publication. A hung in-process producer therefore blocks its owning process and

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

docs/design/module/host_weight_runtime.md

@@ -311,41 +311,9 @@ 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.
-
-`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.
+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.
 
 For tmpfs, artifact bytes are host-memory consumption and may consume swap;
 they must not be reported operationally as ordinary disk capacity. The store

tests/config/test_omni_config.py

@@ -193,6 +193,39 @@ def test_resolve_execution_mode_rejects_unknown_execution_type():
         omni_config_module._resolve_execution_mode("unknown_execution_type")
 
 
+@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")

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)

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -96,7 +96,7 @@ class TestDistributedLayerwiseOffloadHook:
         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 @@ class TestMmapWeightLoading:
 
         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 @@ class TestMmapWeightLoading:
                 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 @@ class TestMmapWeightLoading:
 
         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 @@ class TestMmapWeightLoading:
         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 @@ class TestDistributedComponentSelection:
             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 @@ class TestDistributedComponentSelection:
         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(

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -507,6 +507,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",
         ]
@@ -522,6 +524,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():

tests/host_weight_runtime/test_build_admission.py

@@ -1,305 +0,0 @@
-# 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]

vllm_omni/config/omni_config.py

@@ -739,6 +739,7 @@ class _DiffusionConfigProjection:
     dlo_resident_layers: int = Field(default=0, ge=0)
     host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
     host_weight_runtime_root: str | None = None
+    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"
@@ -895,6 +896,7 @@ class _DiffusionConfigProjection:
         )
         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,

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
+    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

vllm_omni/diffusion/data.py

@@ -119,8 +119,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)
@@ -833,6 +840,7 @@ class OmniDiffusionConfig:
     host_weight_runtime_root: str | None = None
     # 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
@@ -1244,6 +1252,7 @@ class OmniDiffusionConfig:
         )
         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,

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 @@ class 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 @@ class 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,
         )

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -6,12 +6,14 @@ from __future__ import annotations
 
 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 @@ from .host_registration import (
     HostRegistrationError,
 )
 
+logger = init_logger(__name__)
+
 _CUDA_HOST_REGISTER_READ_ONLY = 0x08
 _CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113
 
@@ -187,7 +191,9 @@ class CudaHostRegistration:
 
         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 @@ class CudaHostRegistration:
                         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] = []

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -1080,6 +1080,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         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 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         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(

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
+    dlo_host_registration_mode: str = "auto"
     dlo_host_registration_limit_gib: float = 0.0
     boundary_ratio: float | None = None
     flow_shift: float | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1145,6 +1145,7 @@ class AsyncOmniEngine:
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
             "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
             "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
+            "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": (

vllm_omni/entrypoints/cli/serve.py

@@ -780,6 +780,17 @@ class OmniServeCommand(CLISubcommand):
                 "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,

vllm_omni/host_weight_runtime/config.py

@@ -8,7 +8,6 @@ import math
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
-from typing import Literal
 
 
 class RuntimeMode(str, Enum):
@@ -67,7 +66,6 @@ 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 (
@@ -80,10 +78,6 @@ class CapacityPolicy:
             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)

vllm_omni/host_weight_runtime/filesystem/store.py

@@ -627,7 +627,7 @@ class FilesystemHostWeightStore:
                 self.policy_version = 1
 
     def _capacity_document(self, *, policy_version: int) -> dict[str, object]:
-        document: dict[str, object] = {
+        return {
             "schema_version": 1,
             "policy_version": policy_version,
             "max_artifact_bytes": self.capacity_policy.max_artifact_bytes,
@@ -635,9 +635,6 @@ class FilesystemHostWeightStore:
             "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"
@@ -1016,40 +1013,6 @@ class FilesystemHostWeightStore:
         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: