[Core] Add explicit HWR registration bypass and progress logs

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

benchmarks/host_weight_runtime/README.md

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

benchmarks/host_weight_runtime/safetensors_retention.py

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

docs/design/feature/host_weight_runtime.md

@@ -42,11 +42,6 @@ V1 does not include:
 
 ## Motivation and use cases
 
-For investigating CPU memory retained by dependency tensor materialization,
-see the [standalone safetensors diagnostic](https://github.com/vllm-project/vllm-omni/blob/main/benchmarks/host_weight_runtime/README.md).
-It distinguishes repeated dependency calls from reuse of cached views and does
-not establish per-request HWR leakage.
-
 Model loading can create the same final host representation repeatedly. This
 is especially expensive when loading performs checkpoint decoding, tensor
 renaming, TP slicing, quantization, packing, or scale construction before GPU

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

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():

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,