[Diffusion] Register runtime cache for direct DLO H2D

Avoid recurrent mmap-to-pinned packing in no-AllGather DLO by registering transform-complete runtime-cache mappings under an explicit per-worker budget. Copy mapped tensors directly into the rotating device buffers while preserving bounded-staging fallback and ordered shutdown cleanup.

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -46,7 +46,9 @@ produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
 otherwise it uses the ordinary loader. In no-AllGather mode, an opt-in Phase B
 runtime cache can normalize those final ordinary-loader DiT tensors into an
 immutable node-local mmap entry. Consequently, replicas can share either
-proven-compatible checkpoint pages or equivalent final runtime layouts.
+proven-compatible checkpoint pages or equivalent final runtime layouts. An
+opt-in CUDA registration budget can make the complete final mapping directly
+H2D-copyable, avoiding the recurrent host packing otherwise required by mmap.
 
 The Phase A shared-mmap support boundary is TP1. TP greater than one is an
 ordinary-loader compatibility path: DLO can consume the resulting TP-local
@@ -138,12 +140,13 @@ error retains the already-valid ordinary tensors.
 DLO maps this plan without moving the DiT to `meta` and without rerunning any
 post-load hook. It changes tensor storage while preserving `Parameter` objects,
 then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing
-streaming hooks. The source mappings remain open through DLO disable and feed
+streaming hooks. The source mappings remain open through DLO disable. With the
-the existing bounded two-slot host staging path. Because file-backed pages are
+default zero registration budget they feed the existing bounded two-slot host
-not pinned, every streamed block is first packed from mmap storage into one of
+staging path. Every streamed block is first packed from pageable mmap storage
-those pinned slots and is then copied to the device. This removes model-sized
+into one of those pinned slots and is then copied to the device. This removes
-private pinned storage, but adds a recurrent host-to-host copy and can amplify
+model-sized private pinned storage, but adds a recurrent host-to-host copy and
-TP synchronization waits when several ranks compete for host bandwidth.
+can amplify TP synchronization waits when several ranks compete for host
+bandwidth.
 
 This version is deliberately opt-in and steady-state-oriented. It requires
 roughly one local-disk model copy per runtime identity, has no eviction, and
@@ -151,6 +154,43 @@ still performs ordinary loading plus full content-hash passes in every rank.
 Skip-load-on-hit is a separate follow-up after loader side effects can be
 modeled safely.
 
+### Registered runtime-cache transfer
+
+`dlo_runtime_cache_pin_limit_gib > 0` opts a CUDA worker into complete
+runtime-cache registration. It also changes loader selection: no-AllGather
+uses the ordinary loader and final runtime cache even when a TP1 direct-
+checkpoint plan is available. A checkpoint binding can carry a deferred
+adapter such as grouped-QKV reordering. Registering that raw source would not
+remove the adapter's recurrent CPU allocation/copy, whereas the runtime cache
+contains the transform-complete final tensors.
+
+After mapping and validating every final tensor, DLO groups source storages by
+backing file, page-aligns their address ranges, and coalesces overlapping or
+adjacent ranges within that mapping. The complete byte count is checked
+against the per-worker budget before any CUDA call. It then registers every
+range with the read-only CUDA host-registration flag and verifies that PyTorch
+recognizes every mapped tensor as pinned. Failure rolls back any ranges already
+registered and preserves the existing two-slot staging path. Partial direct
+transfer is deliberately unsupported. A rollback failure aborts worker
+initialization so a still-registered range is retained until process/context
+teardown rather than being unsafely unmapped.
+
+On success, the hooks copy each immutable tensor view directly into its offset
+in the existing flattened rotating device buffer on the copy stream. Parameter
+repointing, two-block HBM residency, and H2D/compute overlap are unchanged. The
+implementation currently issues more H2D copy operations than block packing,
+but avoids the much larger recurrent CPU copy. A packed cache schema should be
+considered only if launch-count measurements justify its added format and
+publication complexity.
+
+Registration is per process and CUDA context; physical file-cache pages remain
+shared and are not copied into one private host allocation per worker. The
+operator must nevertheless budget page-locked memory and satisfy OS/CUDA
+registration limits for every worker. During shutdown, DLO first synchronizes
+outstanding device work, unregisters all ranges, and then closes the safetensors
+mappings. Worker teardown invokes this cleanup before destroying the
+distributed environment and CUDA context.
+
 ### AllGather path
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -203,6 +243,10 @@ the persistent private full-model copy per pure-DP process, but each process
 still packs and transfers every complete block. Sharing is node-local; each
 node has its own page cache.
 
+For a runtime-cache mapping, a positive registration budget replaces those
+host slots with direct copies from the shared final tensor views. A zero budget
+or failed registration retains the bounded staging behavior.
+
 When direct mmap preflight fails, the regular model loader remains responsible
 for preparing each rank's weights, including TP-local tensors or HSDP-managed
 parameters. Supported CPU layouts may then enter the opt-in runtime cache.
@@ -225,9 +269,9 @@ This mode means:
 
 | Parallelism | DLO + AllGather | DLO without AllGather |
 |---|---|---|
-| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. |
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. CUDA workers may register the complete final mapping for direct H2D. |
-| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active without a DLO weight collective. Runtime-cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. |
+| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active without a DLO weight collective. Runtime-cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. Registered transfer does not remove ordinary SP activation/attention collectives. |
-| **TP > 1** | Outside the Phase A direct-mmap scope. The loader preserves TP-local layouts and DLO may apply DP/SP host sharding to ordinary runtime tensors. | The ordinary TP-aware loader produces rank-local tensors. Runtime-cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. |
+| **TP > 1** | Outside the Phase A direct-mmap scope. The loader preserves TP-local layouts and DLO may apply DP/SP host sharding to ordinary runtime tensors. | The ordinary TP-aware loader produces rank-local tensors. Runtime-cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. Registering those final mappings removes staging-driven TP skew but does not remove ordinary TP collectives. |
 | **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
 
 ### Combined dimensions
@@ -286,6 +330,8 @@ Current source-level validation includes:
   rejection;
 - loader ordering through final post-load mutation and DLO remapping without
   rerunning post-load hooks;
+- all-or-nothing page-aligned registration, budget rejection, partial-failure
+  rollback, direct-copy staging bypass, and explicit shutdown cleanup;
 - resident-layer requests requiring no-AllGather;
 - DP request-wave validation for denoising-step compatibility;
 - sharding, double-buffer, AllGather-size, and heterogeneous-block regression
@@ -397,13 +443,48 @@ rose from 3.75 to 32.49 seconds, while communication-kernel time rose from
 summed across ranks and can overlap; they identify the mechanism, not wall
 time. Matching private/cache outputs were byte-identical.
 
-The result confirms that runtime-cache v1 is a host-capacity option, not a
+The result confirms that zero-budget runtime-cache staging is a host-capacity
-latency-neutral replacement for private pinned weights. Broader SP, model,
+option, not a latency-neutral replacement for private pinned weights. Broader
-platform, and production-quality validation remains tracked in
+SP, model, platform, and production-quality validation remains tracked in
 [issue #6231](https://github.com/vllm-project/vllm-omni/issues/6231). The
 historical B300 rows above predate the runtime cache and must not be treated as
 Phase B memory validation.
 
+### L20X registered-runtime-cache smoke
+
+The read-only registered path was then tested with the same MiniMax-H3
+workload. These are functional and mechanism-validation measurements, not a
+production benchmark. The TP1 pair used one request per wave. Both modes moved
+exactly 121.86 GiB H2D in the profiled request and produced byte-identical video
+and audio for both measured seeds.
+
+| TP1 no-AllGather mode | Mean wave | Throughput | Request-time host PSS | `Private_Dirty` | Engine peak HBM | CPU `aten::copy_` | GPU compute |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| Checkpoint mmap + bounded staging | 33.87 s | 0.0295 req/s | 133.65 GiB | 71.16 GiB | 13,226 MB | 21.91 s | 419.28 ms |
+| Final runtime cache + registered direct H2D | 4.88 s | 0.2055 req/s | 129.41 GiB | 66.93 GiB | 13,226 MB | 1.64 s | 419.05 ms |
+
+Registration covered 61.73 GiB in 13 page-aligned ranges. Mean wave latency
+fell by 85.6% and throughput increased 6.94x. GPU compute and H2D payload were
+unchanged, identifying recurrent host transform/packing as the removed work.
+The direct path issued 2,425 H2D events versus 1,945 for packed staging, but
+that launch increase was much smaller than the eliminated CPU cost.
+
+Two independent TP2 engines also registered the same two TP-coordinate cache
+entries across four workers. Compared with the zero-budget cache, mean wave
+latency improved from 9.21/12.55 seconds to 3.80/3.86 seconds for engines A/B;
+the private-loader reference was 3.61/3.72 seconds. Registered request-time
+host PSS was 224.26 GiB versus 232.23 GiB staged and 366.20 GiB private.
+Aggregate CPU `aten::copy_` time was 3.83 seconds registered, 32.49 seconds
+staged, and 3.75 seconds private; aggregate GPU compute and the 264.24 GiB H2D
+payload were effectively unchanged. Ordinary TP communication remained, but
+the extra wait caused by staging-driven rank skew largely disappeared.
+
+All workers explicitly unregistered before releasing their mappings, the
+system's acquired-minus-released pin counter returned to its pre-run value,
+and every GPU allocation was released. Cold runtime-cache publication remained
+expensive because it still performs ordinary loading and multiple full-content
+passes; that startup problem is outside the registered transfer scope.
+
 ## Recommendations
 
 - Use **DP + DLO AllGather** for the supported throughput and host-memory
@@ -413,8 +494,10 @@ Phase B memory validation.
 - Use **no-AllGather** when independent replica execution is required. TP1
   direct-mmap deployments can share checkpoint pages per node. Enable the
   runtime cache only when independent replicas have repeated final layouts and
-  host capacity is more important than the added mmap-to-pinned staging cost.
+  benchmark the target CPU/memory topology, especially for matching
-  Benchmark the target CPU/memory topology, especially for matching
+  coordinates across multiple TP engines. On CUDA, provide an explicit
-  coordinates across multiple TP engines.
+  registration budget large enough for the complete final layout when
+  recurrent staging cost is unacceptable; leave it at zero when page-locking
+  that layout is not operationally acceptable.
 - Prefer **HSDP alone** for production HSDP deployments until the combined
   HSDP + DLO no-AllGather path has broader end-to-end coverage.

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -7,7 +7,9 @@ each rank streams a complete block independently. Compatible TP1 deployments
 can share checkpoint-backed host pages among processes on the same node;
 otherwise DLO streams the ordinary loader's rank-local tensors. The optional
 runtime-weight cache extends node-local sharing to final ordinary-loader
-layouts, including matching TP coordinates across independent engines.
+layouts, including matching TP coordinates across independent engines. On
+CUDA, an explicit registration budget can make those shared mappings directly
+H2D-copyable and avoid recurrent mmap-to-pinned host packing.
 
 See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
 for the implementation contract and compatibility matrix.
@@ -48,7 +50,8 @@ vllm serve /path/to/model --omni \
   --tensor-parallel-size 2 \
   --dlo-no-use-allgather \
   --dlo-enable-runtime-cache \
-  --dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights
+  --dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights \
+  --dlo-runtime-cache-pin-limit-gib 80
 
 # Sequence parallel deployment
 vllm serve /path/to/model --omni \
@@ -65,6 +68,7 @@ omni = Omni(
     dlo_use_allgather=False,
     dlo_enable_runtime_cache=True,
     dlo_runtime_cache_dir="/var/cache/vllm-omni/dlo-runtime-weights",
+    dlo_runtime_cache_pin_limit_gib=80,
 )
 ```
 
@@ -79,6 +83,7 @@ omni = Omni(
 | `--dlo-enable-runtime-cache` | Normalize final ordinary-loader DiT weights into a node-local mmap cache; requires no-AllGather | `false` |
 | `--dlo-runtime-cache-dir PATH` | Shared local-disk cache root | `~/.cache/vllm-omni/dlo-runtime-weights` |
 | `--dlo-runtime-cache-lock-timeout SECONDS` | Maximum wait for another process publishing the same layout | `600` |
+| `--dlo-runtime-cache-pin-limit-gib GIB` | Per-worker limit for registering the complete read-only runtime-cache mapping for direct H2D; zero keeps bounded host staging | `0` |
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
 
 ## Host-weight loading
@@ -95,14 +100,21 @@ DLO may still consume those TP-local tensors, but this is a compatibility path:
 it does not share checkpoint-backed runtime weights across DP replicas and
 provides no shared-mmap host-memory guarantee.
 
-When `--dlo-enable-runtime-cache` is set with no-AllGather, that fallback gains
+When `--dlo-enable-runtime-cache` is set with no-AllGather, ordinary loading
-a second storage opportunity. Every rank still completes ordinary loading,
+gains a second storage opportunity. Every rank completes loader callbacks,
-loader callbacks, post-load processing, validation, and calibration. It then
+post-load processing, validation, and calibration. It then hashes the final
-hashes the final CPU DiT parameters and persistent buffers. One process per
+CPU DiT parameters and persistent buffers. One process per equivalent runtime
-equivalent runtime layout publishes sharded safetensors plus a checksummed
+layout publishes sharded safetensors plus a checksummed manifest; peers
-manifest; peers validate and mmap the same files. DLO replaces only tensor
+validate and mmap the same files. DLO replaces only tensor storage, preserves
-storage, preserves the final `Parameter` objects, and releases the private
+the final `Parameter` objects, and releases the private allocator-backed copies
-allocator-backed copies before installing its streaming hooks.
+before installing its streaming hooks.
+
+A positive `--dlo-runtime-cache-pin-limit-gib` deliberately selects this final
+runtime layout even when a direct TP1 checkpoint plan is available. Raw
+checkpoint mappings may require deferred loader transforms such as grouped-QKV
+reordering; registering them would still leave recurrent host work. Publishing
+the fully transformed layout once makes every streamed source directly
+copyable.
 
 The runtime layout identity excludes DP rank and SP rank. It includes TP world
 size/rank and conservative SP implementation/world-size guards. Therefore two
@@ -110,13 +122,27 @@ independent `TP=2` engines normally create two entries: both TP0 workers map
 one entry and both TP1 workers map the other. Cache sharing does not create a
 process group or an inference-time collective.
 
-At inference time, cached tensors remain pageable mmap views. DLO packs each
+With a zero registration budget, cached tensors remain pageable mmap views.
-streamed block into one of two bounded pinned host slots before the asynchronous
+DLO packs each streamed block into one of two bounded pinned host slots before
-H2D copy. The private ordinary-loader path instead creates model-sized pinned
+the asynchronous H2D copy. This lowers steady-state PSS but adds recurrent
-storage once during initialization. The cache therefore exchanges recurrent
+host-to-host copies; competing TP ranks or engines can turn DDR contention and
-host-to-host copies and DDR bandwidth for lower steady-state PSS; competing TP
+staging skew into collective wait time.
-ranks or engines can turn staging skew into collective wait time. Treat this as
+
-a host-capacity mode and benchmark it on the target CPU/memory topology.
+With a positive budget on CUDA, DLO instead attempts to register the complete
+page-aligned cache mapping as read-only pinned host memory. Registration is
+all-or-nothing: the total mapped range must fit the per-worker budget, and any
+CUDA registration or verification failure rolls back partial registrations and
+uses the bounded staging path. If CUDA cannot roll back a partial registration,
+worker initialization aborts instead of unmapping memory that CUDA may still
+own. On success, each immutable tensor view is copied directly into the
+existing rotating device buffers on the H2D stream, so no private host staging
+slots are allocated.
+
+CUDA registration is process-local, but it does not clone file contents. The
+same physical page-cache pages remain shared among equivalent workers; each
+worker does, however, hold its own registration reference and must satisfy the
+host's page-locking limits. Shutdown synchronizes outstanding GPU work,
+unregisters every range, and only then closes the mmap handles.
 
 The mmap plan skips only dedicated DiT weight sources. Other component sources,
 such as a text encoder loaded through the shared diffusion loader, continue to
@@ -165,6 +191,13 @@ entry is published. Operators may remove entries or the complete cache root
 only after all workers using its mmap files have stopped. The feature is
 opt-in for this reason.
 
+For direct H2D, set the registration limit above the complete page-aligned
+layout size reported in the worker log. Partial registration is intentionally
+unsupported because mixing direct and staged tensors would complicate buffer
+lifetime and make performance topology-dependent. Registration is CUDA-only;
+an insufficient budget, unsupported runtime, OS/CUDA pinning limit, or
+registration error is reported and falls back to bounded staging.
+
 ## Declarative topology
 
 Models may declare an `OffloadPlan` instead of embedding offload logic:
@@ -196,7 +229,9 @@ must enter each collective.
 - Direct checkpoint mmap currently requires TP1. TP greater than one falls
   back to the ordinary TP-aware loader. With the optional runtime cache,
   equivalent processes at the same TP coordinate can then share that final
-  rank-local layout in no-AllGather mode.
+  rank-local layout in no-AllGather mode. A positive registration budget also
+  chooses the final runtime cache at TP1 so deferred checkpoint transforms are
+  not repeated during inference.
 - HSDP plus AllGather is rejected to avoid double sharding. HSDP without
   AllGather has limited end-to-end validation.
 - Per-tensor online FP8 linears use the ordinary loader and can run with either
@@ -212,16 +247,19 @@ must enter each collective.
 
 The runtime cache improves steady-state host PSS, not startup peak: every rank
 loads private weights before remapping, and validation adds full content-hash
-passes. Skip-load-on-cache-hit and cache eviction remain follow-up work in
+passes. Direct H2D removes recurrent staging but does not change that startup
+contract. Skip-load-on-cache-hit and cache eviction remain follow-up work in
 [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
 
-It is also not latency-neutral. In the four-L20X MiniMax-H3 smoke documented in
+Zero-budget staging is also not latency-neutral. In the four-L20X MiniMax-H3
-the L20X runtime-cache matrix in the
+smoke documented in the L20X runtime-cache matrix in the
 [feature design](../../../design/feature/offloader/distributed_layerwise_offload.md#l20x-runtime-cache-smoke-matrix),
 two independent TP2 engines reduced request-time host PSS by 36.6% with the
 cache, but recurrent mmap-to-pinned staging reduced combined throughput by
-69.3%. HBM and GPU compute were unchanged. Use the cache only when that
+69.3%. The registered path recovered near-private per-engine latency while
-measured memory/latency tradeoff is appropriate for the deployment.
+retaining the shared-memory benefit; a matched TP1 run improved mean wave
+latency from 33.87 seconds to 4.88 seconds with byte-identical outputs. See the
+feature design for the full smoke-test caveats and metrics.
 
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -492,6 +492,59 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch,
     assert loader.take_host_weight_plan() is runtime_plan
 
 
+def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=1,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        dlo_runtime_cache_pin_limit_gib=1.0,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+    runtime_plan = HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings={},
+        runtime_layout_key="layout-key",
+        post_load_complete=True,
+    )
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: events.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: events.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+
+    def unexpected_checkpoint_plan(*_args, **_kwargs):
+        pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan")
+
+    def build_cache(_pipeline, **_kwargs):
+        events.append("cache")
+        return HostWeightPlanResult(runtime_plan)
+
+    monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert events == ["load", "process", "cache"]
+    assert loader.take_host_weight_plan() is runtime_plan
+
+
 def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
     import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
+from vllm_omni.diffusion.offloader.cuda_host_registration import (
+    CudaHostRegistration,
+    CudaHostRegistrationBudgetError,
+    CudaHostRegistrationCleanupError,
+    CudaHostRegistrationError,
+    _coalesce_ranges,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _FakeStorage:
+    def __init__(self, pointer: int, size: int) -> None:
+        self._pointer = pointer
+        self._size = size
+
+    def data_ptr(self) -> int:
+        return self._pointer
+
+    def nbytes(self) -> int:
+        return self._size
+
+
+class _FakeTensor:
+    device = SimpleNamespace(type="cpu")
+
+    def __init__(self, pointer: int, size: int, *, pinned: bool = True) -> None:
+        self._storage = _FakeStorage(pointer, size)
+        self._pinned = pinned
+
+    def untyped_storage(self) -> _FakeStorage:
+        return self._storage
+
+    def numel(self) -> int:
+        return self._storage.nbytes()
+
+    def is_pinned(self) -> bool:
+        return self._pinned
+
+
+class _FakeRuntime:
+    def __init__(
+        self,
+        register_results: list[int | Exception],
+        unregister_results: list[int | Exception] | None = None,
+    ) -> None:
+        self._register_results = iter(register_results)
+        self._unregister_results = iter(unregister_results or [])
+        self.registered: list[tuple[int, int, int]] = []
+        self.unregistered: list[int] = []
+
+    def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
+        self.registered.append((pointer, size, flags))
+        result = next(self._register_results)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    def cudaHostUnregister(self, pointer: int) -> int:
+        self.unregistered.append(pointer)
+        result = next(self._unregister_results, 0)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    @staticmethod
+    def cudaGetErrorString(error: int) -> str:
+        return f"error-{error}"
+
+
+def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
+    assert _coalesce_ranges(
+        [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
+        page_size=4096,
+    ) == (
+        registration_module._AddressRange(0x1000, 0x3000),
+        registration_module._AddressRange(0x9000, 0xA000),
+    )
+
+
+def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationBudgetError, match="exceeding"):
+        CudaHostRegistration.create(
+            {"weights": [_FakeTensor(0x1003, 4096)]},  # type: ignore[list-item]
+            max_bytes=4096,
+        )
+
+    assert runtime.registered == []
+
+
+def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="error-7"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, RuntimeError("driver rejected mapping")])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="driver rejected mapping"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7], unregister_results=[9])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationCleanupError, match="rollback errors"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+
+def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {
+            "first": [_FakeTensor(0x1003, 4096), _FakeTensor(0x2800, 1024)],  # type: ignore[list-item]
+            "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+        },
+        max_bytes=12288,
+    )
+
+    assert registration.total_bytes == 12288
+    assert registration.region_count == 2
+    assert registration.close() == []
+    assert runtime.unregistered == [0x9000, 0x1000]
+    assert registration.close() == []
+
+
+def test_registration_keeps_adjacent_file_mappings_separate(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    registration = CudaHostRegistration.create(
+        {
+            "first.safetensors": [_FakeTensor(0x1000, 4096)],  # type: ignore[list-item]
+            "second.safetensors": [_FakeTensor(0x2000, 4096)],  # type: ignore[list-item]
+        },
+        max_bytes=8192,
+    )
+
+    assert registration.region_count == 2
+    assert runtime.registered == [
+        (0x1000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+        (0x2000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+    ]
+    assert registration.close() == []
+
+
+def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0], unregister_results=[9, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
+        max_bytes=4096,
+    )
+
+    assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"]
+    assert registration.close() == []
+    assert runtime.unregistered == [0x1000, 0x1000]

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -546,6 +546,31 @@ class TestDistributedLayerwiseOffloadHook:
         assert staging[0][torch.float32].numel() == 12
         assert staging[1][torch.float32].numel() == 12
 
+    def test_registered_mmap_copies_directly_without_staging(self, patched_offload_runtime):
+        current_block = nn.Linear(2, 2, bias=False)
+        next_block = nn.Linear(2, 2, bias=False)
+        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+        next_block.weight.data.copy_(expected)
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        hook.initialize_hook(current_block)
+        hook.registered_mmap = True
+
+        def fail_staging(_slot):
+            raise AssertionError("registered mmap must bypass host staging")
+
+        hook._stage_mmap_sources = fail_staging  # type: ignore[method-assign]
+        hook.prefetch_layer(slot=0, non_blocking=True)
+
+        assert torch.equal(next_block.weight, expected)
+
 
 class TestPinnedResidentLayerGroup:
     def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -636,6 +661,24 @@ class TestPinnedResidentLayerGroup:
         group.offload()
         assert all(block.weight.numel() == 0 for block in blocks)
 
+    def test_registered_mmap_resident_layers_bypass_staging(self, patched_offload_runtime):
+        block = nn.Linear(2, 2, bias=False)
+        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+        block.weight.data.copy_(expected)
+        group = PinnedResidentLayerGroup(
+            [block],
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        group.registered_mmap = True
+        group._cpu_staging_buffers.clear()
+
+        group.load()
+
+        assert torch.equal(block.weight, expected)
+
     def test_shared_allgather_output_is_narrowed_to_current_block(
         self, dist_group, patched_offload_runtime, monkeypatch
     ):
@@ -995,6 +1038,52 @@ class TestMmapWeightLoading:
         assert len(cleanup_calls) == 2
         backend.disable()
 
+    def test_runtime_cache_registration_uses_budget_and_releases_before_mmap(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        patched_offload_runtime,
+    ):
+        events: list[str] = []
+
+        class Registration:
+            total_bytes = 4096
+            region_count = 1
+
+            @staticmethod
+            def close() -> list[str]:
+                events.append("unregister")
+                return []
+
+        create_calls: list[tuple[dict[str, list[torch.Tensor]], int]] = []
+
+        def create(sources, *, max_bytes):
+            create_calls.append((sources, max_bytes))
+            return Registration()
+
+        monkeypatch.setattr(dist_backend_module.CudaHostRegistration, "create", staticmethod(create))
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dlo_use_allgather=False,
+                dlo_runtime_cache_pin_limit_gib=1.5,
+            ),
+            torch.device("cuda"),
+        )
+        sources = {"runtime.safetensors": [torch.ones(1)]}
+        backend._runtime_cache_mapped_sources = sources
+
+        assert backend._try_register_runtime_cache_mmap()
+        assert create_calls == [(sources, int(1.5 * 1024**3))]
+
+        backend.enabled = True
+        backend._using_rank_local_mmap = True
+        backend._mmap_file_cache = {}
+        backend._release_mmap_handles = lambda: events.append("mmap")  # type: ignore[method-assign]
+        backend.disable()
+
+        assert events == ["unregister", "mmap"]
+
 
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):

tests/diffusion/test_worker_wrapper_base.py

@@ -12,11 +12,13 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """
 
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
 
+import vllm_omni.diffusion.worker.diffusion_worker as worker_module
 from vllm_omni.diffusion.worker.diffusion_worker import (
     CustomPipelineWorkerExtension,
     DiffusionWorker,
@@ -204,6 +206,29 @@ class TestWorkerWrapperBaseDelegation:
         wrapper.worker.shutdown.assert_called_once()
         assert result is None
 
+    def test_worker_shutdown_disables_offloader_before_distributed_teardown(
+        self,
+        mocker: MockerFixture,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        worker = object.__new__(DiffusionWorker)
+        offload_backend = mocker.Mock()
+        kv_transfer_manager = mocker.Mock()
+        worker.model_runner = SimpleNamespace(
+            offload_backend=offload_backend,
+            kv_transfer_manager=kv_transfer_manager,
+        )
+        events: list[str] = []
+        offload_backend.disable.side_effect = lambda: events.append("offloader")
+        kv_transfer_manager.shutdown_prefetch.side_effect = lambda: events.append("kv_transfer")
+        monkeypatch.setattr(worker_module, "destroy_distributed_env", lambda: events.append("distributed"))
+
+        worker.shutdown()
+
+        offload_backend.disable.assert_called_once_with()
+        kv_transfer_manager.shutdown_prefetch.assert_called_once_with()
+        assert events == ["offloader", "kv_transfer", "distributed"]
+
 
 # -------------------------------------------------------------------------
 # Tests: execute_method

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -333,6 +333,8 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "/var/cache/vllm-omni-dlo",
             "--dlo-runtime-cache-lock-timeout",
             "30",
+            "--dlo-runtime-cache-pin-limit-gib",
+            "80",
             "--dlo-resident-layers",
             "20",
         ]
@@ -347,12 +349,14 @@ def test_serve_cli_forwards_distributed_offload_residency():
     assert args.dlo_enable_runtime_cache is True
     assert args.dlo_runtime_cache_dir == "/var/cache/vllm-omni-dlo"
     assert args.dlo_runtime_cache_lock_timeout == 30
+    assert args.dlo_runtime_cache_pin_limit_gib == 80
     assert args.dlo_resident_layers == 20
     assert engine_args["enable_distributed_layerwise_offload"] is True
     assert engine_args["dlo_use_allgather"] is False
     assert engine_args["dlo_enable_runtime_cache"] is True
     assert engine_args["dlo_runtime_cache_dir"] == "/var/cache/vllm-omni-dlo"
     assert engine_args["dlo_runtime_cache_lock_timeout"] == 30
+    assert engine_args["dlo_runtime_cache_pin_limit_gib"] == 80
     assert engine_args["dlo_resident_layers"] == 20
 
 

vllm_omni/config/omni_config.py

@@ -690,6 +690,7 @@ class _DiffusionConfigProjection:
     dlo_enable_runtime_cache: bool = False
     dlo_runtime_cache_dir: str | None = None
     dlo_runtime_cache_lock_timeout: float = Field(default=600.0, gt=0)
+    dlo_runtime_cache_pin_limit_gib: float = Field(default=0.0, ge=0)
     dlo_resident_layers: int = Field(default=0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
@@ -756,6 +757,8 @@ class _DiffusionConfigProjection:
             raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload")
         if self.dlo_enable_runtime_cache and self.dlo_use_allgather:
             raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False")
+        if self.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True")
 
         if self.tf_model_config is None:
             self.tf_model_config = TransformerConfig()

vllm_omni/config/stage_config.py

@@ -385,6 +385,7 @@ class StageDeployConfig:
     dlo_enable_runtime_cache: bool | None = None
     dlo_runtime_cache_dir: str | None = None
     dlo_runtime_cache_lock_timeout: float | None = None
+    dlo_runtime_cache_pin_limit_gib: float | None = None
     dlo_resident_layers: int | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None

vllm_omni/diffusion/data.py

@@ -708,13 +708,18 @@ class OmniDiffusionConfig:
     # This avoids AllGather synchronization, while host memory follows the
     # loader's existing rank-local layout instead of adding a second DP shard.
     dlo_use_allgather: bool = True
-    # Build or join a node-local mmap cache after ordinary loading when direct
+    # Build or join a node-local mmap cache after ordinary loading. This is
-    # checkpoint mmap is unavailable. The first version is no-AllGather only.
+    # used when direct checkpoint mmap is unavailable, or when direct CUDA
+    # registration requires a transform-complete final layout. The first
+    # version is no-AllGather only.
     dlo_enable_runtime_cache: bool = False
     # Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-runtime-weights.
     dlo_runtime_cache_dir: str | None = None
     # Maximum wait for the per-layout POSIX writer lock.
     dlo_runtime_cache_lock_timeout: float = 600.0
+    # Maximum file-backed memory a worker may register with CUDA. Zero keeps
+    # the bounded mmap-to-pinned staging path.
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
 
@@ -945,10 +950,14 @@ class OmniDiffusionConfig:
     def __post_init__(self):
         if not math.isfinite(self.dlo_runtime_cache_lock_timeout) or self.dlo_runtime_cache_lock_timeout <= 0:
             raise ValueError("dlo_runtime_cache_lock_timeout must be positive")
+        if not math.isfinite(self.dlo_runtime_cache_pin_limit_gib) or self.dlo_runtime_cache_pin_limit_gib < 0:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib must be non-negative")
         if self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload:
             raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload")
         if self.dlo_enable_runtime_cache and self.dlo_use_allgather:
             raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False")
+        if self.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True")
         if self.diffusion_compile_granularity not in {"regional", "full"}:
             raise ValueError(
                 "diffusion_compile_granularity must be 'regional' or 'full', "

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -441,10 +441,16 @@ class DiffusersPipelineLoader:
                 _dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1))
                 _sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1))
                 _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size
+                _prefer_registered_runtime_cache = (
+                    _dist_offload
+                    and not _use_ag
+                    and bool(getattr(self.od_config, "dlo_enable_runtime_cache", False))
+                    and float(getattr(self.od_config, "dlo_runtime_cache_pin_limit_gib", 0.0)) > 0
+                )
 
                 plan_result = None
                 weight_sources = self._get_weight_sources(model)
-                if _dist_offload:
+                if _dist_offload and not _prefer_registered_runtime_cache:
                     modules = ModuleDiscovery.discover(model)
                     plan_result = build_checkpoint_mmap_plan(
                         model,
@@ -456,6 +462,16 @@ class DiffusersPipelineLoader:
                         online_quantization=_has_online_quant,
                     )
                     self.host_weight_plan = plan_result.plan
+                elif _prefer_registered_runtime_cache:
+                    # Direct checkpoint bindings can carry deferred transforms
+                    # (for example grouped-QKV reordering), so registering the
+                    # raw checkpoint mapping would not make every recurrent H2D
+                    # source directly copyable. Materialize once, apply every
+                    # post-load mutation, then publish/register the final layout.
+                    logger.info(
+                        "DLO runtime-cache registration requested; using the ordinary loader "
+                        "once to publish a transform-complete final mmap layout"
+                    )
 
                 _skip_load = self.host_weight_plan is not None
 

vllm_omni/diffusion/offloader/base.py

@@ -31,6 +31,9 @@ class OffloadConfig:
     # blocks from the loader-selected host backing with H2D only.
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0  # leading DiT layers kept on device
+    # Per-worker budget for registering shared runtime-cache mappings with CUDA.
+    # Zero retains bounded host staging.
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
 
     @classmethod
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -101,6 +104,7 @@ class OffloadConfig:
         # requirements (concurrent requests, dummy run skip).
         dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
         dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
+        dlo_runtime_cache_pin_limit_gib = float(getattr(od_config, "dlo_runtime_cache_pin_limit_gib", 0.0))
         if dlo_resident_layers < 0:
             raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
         if dlo_resident_layers and dlo_use_allgather:
@@ -136,6 +140,7 @@ class OffloadConfig:
             dp_size=dp_size,
             dlo_use_allgather=dlo_use_allgather,
             dlo_resident_layers=dlo_resident_layers,
+            dlo_runtime_cache_pin_limit_gib=dlo_runtime_cache_pin_limit_gib,
         )
 
 

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bounded CUDA registration for existing read-only host mappings."""
+
+from __future__ import annotations
+
+import mmap
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+_CUDA_HOST_REGISTER_READ_ONLY = 0x08
+
+
+class CudaHostRegistrationError(RuntimeError):
+    """Raised when an existing host range cannot be registered safely."""
+
+
+class CudaHostRegistrationBudgetError(CudaHostRegistrationError):
+    """Raised before mutation when the complete mapping exceeds its budget."""
+
+
+class CudaHostRegistrationCleanupError(CudaHostRegistrationError):
+    """Raised when partial registration could not be rolled back safely."""
+
+
+@dataclass(frozen=True)
+class _AddressRange:
+    start: int
+    end: int
+
+    @property
+    def size(self) -> int:
+        return self.end - self.start
+
+
+def _coalesce_ranges(
+    ranges: Sequence[tuple[int, int]],
+    page_size: int = mmap.PAGESIZE,
+) -> tuple[_AddressRange, ...]:
+    """Page-align and merge overlapping ranges from one backing mapping."""
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+
+    aligned: list[_AddressRange] = []
+    for start, size in ranges:
+        if start <= 0 or size < 0:
+            raise ValueError(f"invalid host range start={start}, size={size}")
+        if size == 0:
+            continue
+        aligned_start = start - start % page_size
+        end = start + size
+        aligned_end = ((end + page_size - 1) // page_size) * page_size
+        aligned.append(_AddressRange(aligned_start, aligned_end))
+
+    merged: list[_AddressRange] = []
+    for region in sorted(aligned, key=lambda item: (item.start, item.end)):
+        if merged and region.start <= merged[-1].end:
+            previous = merged[-1]
+            merged[-1] = _AddressRange(previous.start, max(previous.end, region.end))
+        else:
+            merged.append(region)
+    return tuple(merged)
+
+
+def _tensor_regions(
+    sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
+) -> tuple[_AddressRange, ...]:
+    """Resolve storage spans without merging unrelated file mappings."""
+    regions: list[_AddressRange] = []
+    for mapping_name, tensors in sources_by_mapping.items():
+        ranges: list[tuple[int, int]] = []
+        for tensor in tensors:
+            if tensor.device.type != "cpu":
+                raise CudaHostRegistrationError(
+                    f"CUDA host registration requires CPU storage, but {mapping_name!r} contains {tensor.device}"
+                )
+            storage = tensor.untyped_storage()
+            ranges.append((storage.data_ptr(), storage.nbytes()))
+        regions.extend(_coalesce_ranges(ranges))
+    return tuple(regions)
+
+
+def _error_message(runtime: Any, error: Any) -> str:
+    try:
+        message = runtime.cudaGetErrorString(error)
+    except Exception:
+        return str(error)
+    if isinstance(message, bytes):
+        return message.decode(errors="replace")
+    return str(message)
+
+
+class CudaHostRegistration:
+    """Own CUDA registrations for already-existing file-backed CPU tensors."""
+
+    def __init__(
+        self,
+        runtime: Any,
+        regions: tuple[_AddressRange, ...],
+    ) -> None:
+        self._runtime = runtime
+        self._regions = regions
+        self._closed = False
+
+    @classmethod
+    def create(
+        cls,
+        sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
+        *,
+        max_bytes: int,
+    ) -> CudaHostRegistration:
+        if max_bytes <= 0:
+            raise CudaHostRegistrationBudgetError("CUDA host-registration budget is disabled")
+        regions = _tensor_regions(sources_by_mapping)
+        total_bytes = sum(region.size for region in regions)
+        if total_bytes > max_bytes:
+            raise CudaHostRegistrationBudgetError(
+                f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
+            )
+        if not regions:
+            raise CudaHostRegistrationError("no non-empty host ranges were available for registration")
+        if not torch.cuda.is_available():
+            raise CudaHostRegistrationError("CUDA is not available")
+
+        try:
+            runtime = torch.cuda.cudart()
+        except Exception as exc:
+            raise CudaHostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc
+
+        registered: list[_AddressRange] = []
+        try:
+            for region in regions:
+                error = runtime.cudaHostRegister(
+                    region.start,
+                    region.size,
+                    _CUDA_HOST_REGISTER_READ_ONLY,
+                )
+                if int(error) != 0:
+                    raise CudaHostRegistrationError(
+                        "cudaHostRegister(read-only) failed for "
+                        f"[{region.start:#x}, {region.end:#x}): {_error_message(runtime, error)}"
+                    )
+                registered.append(region)
+
+            unpinned = [
+                mapping_name
+                for mapping_name, tensors in sources_by_mapping.items()
+                if any(tensor.numel() and not tensor.is_pinned() for tensor in tensors)
+            ]
+            if unpinned:
+                raise CudaHostRegistrationError(
+                    f"CUDA registration succeeded but PyTorch did not recognize pinned storage for {unpinned[:3]}"
+                )
+        except Exception as exc:
+            rollback_errors: list[str] = []
+            for region in reversed(registered):
+                try:
+                    error = runtime.cudaHostUnregister(region.start)
+                    if int(error) != 0:
+                        rollback_errors.append(
+                            f"cudaHostUnregister({region.start:#x}) failed: {_error_message(runtime, error)}"
+                        )
+                except Exception as rollback_exc:
+                    rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}")
+            if rollback_errors:
+                raise CudaHostRegistrationCleanupError(
+                    f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}"
+                ) from exc
+            if isinstance(exc, CudaHostRegistrationError):
+                raise
+            raise CudaHostRegistrationError(f"CUDA host registration raised: {exc}") from exc
+
+        return cls(runtime, regions)
+
+    @property
+    def total_bytes(self) -> int:
+        return sum(region.size for region in self._regions)
+
+    @property
+    def region_count(self) -> int:
+        return len(self._regions)
+
+    def close(self) -> list[str]:
+        """Unregister every range, returning errors after best-effort cleanup."""
+        if self._closed:
+            return []
+        errors: list[str] = []
+        failed: list[_AddressRange] = []
+        for region in reversed(self._regions):
+            try:
+                error = self._runtime.cudaHostUnregister(region.start)
+                if int(error) != 0:
+                    errors.append(
+                        f"cudaHostUnregister({region.start:#x}) failed: {_error_message(self._runtime, error)}"
+                    )
+                    failed.append(region)
+            except Exception as exc:
+                errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}")
+                failed.append(region)
+        self._regions = tuple(reversed(failed))
+        self._closed = not self._regions
+        return errors

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -18,6 +18,7 @@ This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
 
 from __future__ import annotations
 
+import time
 from itertools import chain
 from typing import Any
 
@@ -34,6 +35,11 @@ from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig
 from .block_discovery import get_blocks_from_dit
+from .cuda_host_registration import (
+    CudaHostRegistration,
+    CudaHostRegistrationCleanupError,
+    CudaHostRegistrationError,
+)
 from .module_collector import ModuleDiscovery
 from .offload_plan import (
     OffloadPlan,
@@ -93,6 +99,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         self.rank = rank
         self.pin_memory = pin_memory
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.tensor_transforms = tensor_transforms or {}
 
         self.copy_stream = copy_stream or current_omni_platform.Stream()
@@ -444,6 +451,26 @@ class DistributedLayerwiseOffloadHook(ModelHook):
     #  Prefetch: H2D + AllGather (overlapped on dedicated streams)      #
     # ------------------------------------------------------------------ #
 
+    @staticmethod
+    def _resolve_mmap_source(
+        source_info: dict[str, Any],
+        meta: dict[str, Any],
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        source = source_info["tensor"]
+        transform = source_info["transform"]
+        if callable(transform):
+            source = transform(source)
+        if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
+            raise ValueError(
+                "mmap weight transform changed tensor layout for "
+                f"{source_info['name']!r}: expected dtype={dtype}, "
+                f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
+                f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
+                f"stride={source.stride()}"
+            )
+        return source
+
     @staticmethod
     def _pack_mmap_sources(
         cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
@@ -456,18 +483,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             destination = slot_buffers[dtype][:total_numel]
             sources = cpu_sources[dtype]
             for source_info, meta in zip(sources, metas, strict=True):
-                source = source_info["tensor"]
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
-                transform = source_info["transform"]
-                if callable(transform):
-                    source = transform(source)
-                if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
-                    raise ValueError(
-                        "mmap weight transform changed tensor layout for "
-                        f"{source_info['name']!r}: expected dtype={dtype}, "
-                        f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
-                        f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
-                        f"stride={source.stride()}"
-                    )
                 start = meta["offset"]
                 physical_storage = destination[start : start + meta["numel"]]
                 if source.is_contiguous():
@@ -481,6 +497,32 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             staged[dtype] = destination
         return staged
 
+    @staticmethod
+    def _copy_mmap_sources_to_device(
+        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
+        metadata: dict[torch.dtype, list[dict[str, Any]]],
+        gpu_buffers: dict[torch.dtype, torch.Tensor],
+        *,
+        non_blocking: bool,
+    ) -> None:
+        """Copy registered source views directly into flattened device buffers."""
+        for dtype, metas in metadata.items():
+            destination = gpu_buffers[dtype]
+            sources = cpu_sources[dtype]
+            for source_info, meta in zip(sources, metas, strict=True):
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
+                start = meta["offset"]
+                physical_storage = destination[start : start + meta["numel"]]
+                async_copy = non_blocking and source.is_pinned()
+                if source.is_contiguous():
+                    physical_storage.copy_(source.flatten(), non_blocking=async_copy)
+                else:
+                    torch.as_strided(
+                        physical_storage,
+                        size=source.shape,
+                        stride=source.stride(),
+                    ).copy_(source, non_blocking=async_copy)
+
     def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]:
         """Pack this block's mmap views into one bounded host staging slot."""
         previous_copy = self.cpu_staging_events[slot]
@@ -514,17 +556,27 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
 
         if self.dp_size <= 1 or self.dp_group is None:
-            cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
+            if self.rank_local_mmap and self.registered_mmap:
-            with current_omni_platform.stream(self.copy_stream):
+                with current_omni_platform.stream(self.copy_stream):
-                for dtype, cpu_shard in cpu_weights.items():
+                    self._copy_mmap_sources_to_device(
-                    gw = gpu_weights[dtype]
+                        self.cpu_sources,
-                    async_copy = non_blocking and cpu_shard.is_pinned()
+                        self.metadata,
-                    gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
+                        gpu_weights,
-                evt.record(self.copy_stream)
+                        non_blocking=non_blocking,
-            if self.rank_local_mmap:
+                    )
-                # The CPU slot may be overwritten only after this H2D copy has
+                    evt.record(self.copy_stream)
-                # finished.  The shared event protects reuse by another hook.
+            else:
-                self.cpu_staging_events[slot] = evt
+                cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
+                with current_omni_platform.stream(self.copy_stream):
+                    for dtype, cpu_shard in cpu_weights.items():
+                        gw = gpu_weights[dtype]
+                        async_copy = non_blocking and cpu_shard.is_pinned()
+                        gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
+                    evt.record(self.copy_stream)
+                if self.rank_local_mmap:
+                    # The CPU slot may be overwritten only after this H2D copy has
+                    # finished.  The shared event protects reuse by another hook.
+                    self.cpu_staging_events[slot] = evt
         else:
             gpu_shards: dict[torch.dtype, torch.Tensor] = {}
             shard_bufs = self.gpu_shard_buffers[slot]
@@ -734,6 +786,7 @@ class PinnedResidentLayerGroup:
         self.copy_stream = copy_stream
         self.loaded = False
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.pin_memory = pin_memory
         self._states: list[dict[str, Any]] = []
         self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
@@ -804,6 +857,14 @@ class PinnedResidentLayerGroup:
         ready = current_omni_platform.Event()
         with current_omni_platform.stream(self.copy_stream):
             for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
+                if self.rank_local_mmap and self.registered_mmap:
+                    DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device(
+                        state["cpu_sources"],
+                        state["metadata"],
+                        block_buffers,
+                        non_blocking=True,
+                    )
+                    continue
                 if self.rank_local_mmap:
                     slot = index % 2
                     previous_copy = self._cpu_staging_events[slot]
@@ -901,8 +962,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._resident_layer_group: PinnedResidentLayerGroup | None = None
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.host_weight_plan = host_weight_plan
         self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
+        self._runtime_cache_mapped_sources: dict[str, list[torch.Tensor]] = {}
+        self._cuda_host_registration: CudaHostRegistration | None = None
 
     def load_resident_layers(self) -> None:
         """Load the model-declared leading blocks for the denoise stage."""
@@ -1135,6 +1199,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         file_cache: dict[str, Any] = {}
         replacements: list[tuple[torch.Tensor, torch.Tensor]] = []
+        mapped_sources: dict[str, list[torch.Tensor]] = {}
         try:
             # Validate and map every tensor before mutating the pipeline. A
             # missing/corrupt shard therefore leaves the ordinary tensors intact.
@@ -1169,6 +1234,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                         f"runtime={tuple(target.shape)}/{target.dtype}"
                     )
                 replacements.append((target, tensor))
+                mapped_sources.setdefault(binding.file_path, []).append(tensor)
         except Exception:
             file_cache.clear()
             raise
@@ -1181,12 +1247,74 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         self._mmap_transforms_by_tensor_id.clear()
         self._mmap_file_cache = file_cache
+        self._runtime_cache_mapped_sources = mapped_sources
         logger.info(
             "Remapped %d final DiT tensors from runtime cache %s",
             len(replacements),
             plan.runtime_layout_key[:12] if plan.runtime_layout_key else "unknown",
         )
 
+    def _try_register_runtime_cache_mmap(self) -> bool:
+        """Register complete cache mappings when the operator budget permits."""
+        pin_limit_gib = self.config.dlo_runtime_cache_pin_limit_gib
+        if pin_limit_gib <= 0:
+            return False
+        if self.device.type != "cuda":
+            logger.warning(
+                "DLO runtime-cache registration is CUDA-only; using bounded host staging on %s",
+                self.device,
+            )
+            return False
+        if not self.config.pin_cpu_memory:
+            logger.warning("DLO runtime-cache registration requires pin_cpu_memory=True; using bounded host staging")
+            return False
+        if not self._runtime_cache_mapped_sources:
+            logger.warning("DLO runtime-cache registration found no mapped sources; using bounded host staging")
+            return False
+
+        max_bytes = int(pin_limit_gib * 1024**3)
+        started = time.perf_counter()
+        try:
+            registration = CudaHostRegistration.create(
+                self._runtime_cache_mapped_sources,
+                max_bytes=max_bytes,
+            )
+        except CudaHostRegistrationCleanupError:
+            # Falling back could unmap a range that CUDA still owns. Abort the
+            # worker so process/context teardown provides the final cleanup.
+            logger.exception("DLO runtime-cache registration rollback failed")
+            raise
+        except CudaHostRegistrationError as exc:
+            logger.warning(
+                "DLO runtime-cache direct H2D unavailable (%s); using bounded host staging",
+                exc,
+            )
+            return False
+
+        self._cuda_host_registration = registration
+        logger.info(
+            "Registered %.2f GiB of shared runtime-cache mappings in %d range(s) for direct H2D in %.3f s",
+            registration.total_bytes / 1024**3,
+            registration.region_count,
+            time.perf_counter() - started,
+        )
+        return True
+
+    def _release_registered_mmap(self) -> bool:
+        registration = self._cuda_host_registration
+        if registration is None:
+            return True
+        errors = registration.close()
+        if errors:
+            logger.error(
+                "DLO runtime-cache host unregistration failed; retaining mmap handles for retry/process exit: %s",
+                errors[:3],
+            )
+            return False
+        self._cuda_host_registration = None
+        logger.info("Unregistered DLO runtime-cache host mappings")
+        return True
+
     def _init_dp_group(self) -> None:
         """Reuse the process group initialized by parallel_state.
 
@@ -1550,8 +1678,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}")
             if self._using_rank_local_mmap:
                 logger.info(
-                    "DLO rank-local mmap storage enabled: file-backed pages are "
+                    "DLO rank-local mmap storage enabled: file-backed pages are node-shared; "
-                    "node-shared; each worker owns only two bounded host staging slots"
+                    "transfer setup will select registered direct H2D or bounded host staging"
                 )
         else:
             remaining_meta = [
@@ -1723,6 +1851,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         if not self._all_hook_groups:
             self.enabled = bool(self._resident_blocks)
+            if self.enabled and self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "runtime_cache":
+                self._using_registered_mmap = self._try_register_runtime_cache_mmap()
+                self._runtime_cache_mapped_sources.clear()
+                assert self._resident_layer_group is not None
+                self._resident_layer_group.registered_mmap = self._using_registered_mmap
+                if self._using_registered_mmap:
+                    self._resident_layer_group._cpu_staging_buffers.clear()
             if self._using_mmap and not self.enabled:
                 self._release_mmap_handles()
             return
@@ -1738,9 +1873,18 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         unified_shard_buffers = None
         if self.dp_size > 1:
             unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
+
+        if self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "runtime_cache":
+            self._using_registered_mmap = self._try_register_runtime_cache_mmap()
+            self._runtime_cache_mapped_sources.clear()
+            for hook in all_hooks:
+                hook.registered_mmap = self._using_registered_mmap
+            if self._resident_layer_group is not None:
+                self._resident_layer_group.registered_mmap = self._using_registered_mmap
+
         unified_cpu_staging = None
         cpu_staging_events = None
-        if self._using_rank_local_mmap:
+        if self._using_rank_local_mmap and not self._using_registered_mmap:
             unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                 all_hooks,
                 self._resident_layer_group,
@@ -1822,6 +1966,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
     def _release_mmap_handles(self) -> None:
         """Release safetensors mmap file handles."""
         self._mmap_transforms_by_tensor_id.clear()
+        self._runtime_cache_mapped_sources.clear()
         if hasattr(self, "_mmap_file_cache"):
             self._mmap_file_cache.clear()
             del self._mmap_file_cache
@@ -1846,6 +1991,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         if self._using_rank_local_mmap:
             current_omni_platform.synchronize()
+        registration_released = self._release_registered_mmap()
 
         for blocks in self._blocks:
             for block in blocks:
@@ -1861,9 +2007,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups.clear()
         self._resident_blocks.clear()
         self._resident_layer_group = None
-        self._release_mmap_handles()
+        if registration_released:
+            self._release_mmap_handles()
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.enabled = False
         logger.info("Distributed layer-wise offloading disabled")
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -859,11 +859,18 @@ class DiffusionWorker:
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
-        if self.model_runner is not None:
+        try:
-            mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+            if self.model_runner is not None:
-            if mgr is not None:
+                mgr = getattr(self.model_runner, "kv_transfer_manager", None)
-                mgr.shutdown_prefetch()
+                try:
-        destroy_distributed_env()
+                    offload_backend = getattr(self.model_runner, "offload_backend", None)
+                    if offload_backend is not None:
+                        offload_backend.disable()
+                finally:
+                    if mgr is not None:
+                        mgr.shutdown_prefetch()
+        finally:
+            destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:

vllm_omni/engine/arg_utils.py

@@ -548,6 +548,7 @@ class OrchestratorArgs:
     dlo_enable_runtime_cache: bool = False
     dlo_runtime_cache_dir: str | None = None
     dlo_runtime_cache_lock_timeout: float = 600.0
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
     dlo_resident_layers: int = 0
     boundary_ratio: float | None = None
     flow_shift: float | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1043,6 +1043,7 @@ class AsyncOmniEngine:
             "dlo_enable_runtime_cache": kwargs.get("dlo_enable_runtime_cache", False),
             "dlo_runtime_cache_dir": kwargs.get("dlo_runtime_cache_dir", None),
             "dlo_runtime_cache_lock_timeout": kwargs.get("dlo_runtime_cache_lock_timeout", 600.0),
+            "dlo_runtime_cache_pin_limit_gib": kwargs.get("dlo_runtime_cache_pin_limit_gib", 0.0),
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 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

@@ -746,6 +746,15 @@ class OmniServeCommand(CLISubcommand):
             default=600.0,
             help="Seconds to wait for another process publishing the same runtime layout.",
         )
+        omni_config_group.add_argument(
+            "--dlo-runtime-cache-pin-limit-gib",
+            type=float,
+            default=0.0,
+            help=(
+                "Maximum GiB of read-only runtime-cache mappings each CUDA worker "
+                "may register for direct H2D. Zero keeps bounded pinned staging."
+            ),
+        )
         omni_config_group.add_argument(
             "--dlo-resident-layers",
             type=int,