[Diffusion] Prefer host-cache registration when pinned

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -46,13 +46,14 @@ produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
 otherwise it uses the ordinary loader. In no-AllGather mode, the Phase B host
 weight cache can explicitly normalize supported final ordinary-loader DiT
 tensors into an immutable node-local mmap entry. Ordinary pinned weights remain
-the default after checkpoint mmap fallback because bounded cache staging can
-reduce throughput. Replicas can share either proven-compatible checkpoint pages
-or, when selected, equivalent final runtime layouts. An opt-in
-host-registration budget can make the complete final mapping
-directly H2D-copyable when the platform provides a backend, avoiding the
-recurrent host packing otherwise required by mmap. CUDA is the first
-implementation.
+the default after checkpoint mmap fallback when the cache is not selected; cache
+publication and validation add startup and disk lifecycle costs. Replicas can
+share either proven-compatible checkpoint pages or, when selected, equivalent
+final runtime layouts. Under the existing
+pinned-memory policy, the backend attempts to make the complete final mapping
+directly H2D-copyable when the platform provides registration, avoiding the
+recurrent host packing otherwise required by mmap. An optional positive limit
+adds a per-worker safety ceiling. CUDA is the first implementation.
 
 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
@@ -107,7 +108,7 @@ flowchart TB
     Discovery["ModuleDiscovery + optional OffloadPlan"] --> Backend
     Contract -.->|"passed once by the runner"| Backend
 
-    Backend -.->|"host weight cache + positive budget"| Registration["host_registration<br/>platform-neutral lifecycle"]
+    Backend -.->|"host weight cache + pinned-memory policy"| Registration["host_registration<br/>platform-neutral lifecycle"]
     Registration --> CUDA["CUDA registration backend"]
     Backend --> Hooks["DistributedLayerwiseOffloadHook<br/>one circular hook per streamed block"]
     Forward["pipeline.forward"] --> Hooks
@@ -119,7 +120,7 @@ flowchart TB
 | Component | Owns | Must not own |
 | --- | --- | --- |
 | CLI, engine arguments, and `OmniDiffusionConfig` | User policy, defaults, validation, and propagation to every worker | Checkpoint inspection, cache publication, or buffer allocation |
-| `OffloadConfig.from_od_config()` | Offload-strategy selection, effective DLO group size, AllGather policy, resident-layer count, and registration budget | Host-weight compatibility decisions |
+| `OffloadConfig.from_od_config()` | Offload-strategy selection, effective DLO group size, AllGather policy, resident-layer count, pinned-memory policy, and optional registration ceiling | Host-weight compatibility decisions |
 | `DiffusersPipelineLoader` | Ordinary loading lifecycle and the decision to return an optional `HostWeightPlan` | DLO process groups or block-stream scheduling |
 | Checkpoint mmap planner | Proving that checkpoint tensors can represent the required runtime layout before ordinary DiT materialization is skipped | Host weight cache identity or transfer policy |
 | `host_weight_cache` | Normalizing final ordinary-loader tensors into an immutable node-local entry and returning a validated plan | Model mutation, host registration, H2D, or collectives |
@@ -133,11 +134,14 @@ flowchart TB
 The configuration fields are routed to the component that can enforce their
 contract. The cache root and lock timeout are consumed by the loader.
 `dlo_use_host_weight_cache` explicitly selects final-layout cache publication
-after direct-checkpoint mmap fallback. It defaults to false so the ordinary
+instead of direct-checkpoint mmap. It defaults to false so the ordinary
 pinned-loader path remains available without requiring an opt-out.
-`dlo_host_weight_cache_pin_limit_gib` is also visible while loading because a
-positive value also deliberately selects the transform-complete host weight
-cache; the backend later consumes the same budget for registration.
+With `pin_cpu_memory=True`, selecting the host weight cache deliberately chooses
+the transform-complete layout even when direct checkpoint mmap is available,
+because only that final layout can be registered without deferred transforms.
+The backend consumes `dlo_host_weight_cache_pin_limit_gib` as an optional
+registration ceiling; zero applies no additional ceiling and does not select
+the cache by itself.
 `dlo_use_allgather` controls loader plan routing because the host weight cache
 is no-AllGather-only, and it selects the backend transfer protocol. This is
 coordination through typed configuration and `HostWeightPlan`, not through
@@ -246,19 +250,22 @@ sequenceDiagram
     participant P as platform registration
     participant H as block hooks
 
-    L->>L: checkpoint mmap preflight
-    alt loader selects a compatible checkpoint plan
-        L-->>B: checkpoint plan via runner
-    else ordinary loading
+    alt host weight cache explicitly selected
         L->>L: load, mutate, validate, and calibrate final tensors
-        opt host weight cache enabled
-            L->>C: publish or join final-layout entry
-            C-->>L: validated host weight cache plan or fail-closed fallback
+        L->>C: publish or join final-layout entry
+        C-->>L: validated host weight cache plan or fail-closed fallback
+        L-->>B: cache plan or materialized fallback via runner
+    else automatic checkpoint selection
+        L->>L: checkpoint mmap preflight
+        alt loader selects a compatible checkpoint plan
+            L-->>B: checkpoint plan via runner
+        else ordinary loading
+            L->>L: load, mutate, validate, and calibrate final tensors
+            L-->>B: materialized tensors via runner
         end
-        L-->>B: plan or materialized tensors via runner
     end
     B->>B: consume complete rank-local block sources
-    opt host weight cache plan and positive registration budget
+    opt host weight cache plan and pinned-memory policy
         B->>P: register the complete immutable mapping
         P-->>B: direct-H2D source or bounded-staging fallback
     end
@@ -267,7 +274,7 @@ sequenceDiagram
         alt registered host weight cache mapping
             H->>H: copy complete next block directly to device
         else pageable mmap mapping
-            H->>H: pack into a pinned staging slot, then copy to device
+            H->>H: pack into a bounded staging slot, then copy to device
         else ordinary runtime tensors
             H->>H: copy complete next block to device
         end
@@ -287,8 +294,8 @@ the source of the complete-block H2D copy.
 | Ordinary runtime tensors | AllGather group larger than one | Each rank retains a pinned local weight shard | shard H2D → DLO AllGather → device slot |
 | Checkpoint mmap | AllGather group larger than one | Each rank prepares its persistent shard, then releases the source mapping | shard H2D → DLO AllGather → device slot |
 | Ordinary runtime tensors | Rank-local / no-AllGather | Each rank retains the complete loader-produced local layout | complete-block H2D → device slot |
-| Checkpoint or host weight cache mmap, zero registration budget | Rank-local / no-AllGather | Immutable mapping plus two bounded pinned host-staging slots | mmap views → pinned staging slot → device slot |
-| Host weight cache mmap, positive registration budget | Rank-local / no-AllGather | Complete mapping is registered by the platform backend | registered mmap views → device slot |
+| Checkpoint mmap, or host weight cache with pinned memory disabled/registration unavailable | Rank-local / no-AllGather | Immutable mapping plus two bounded host-staging slots | mmap views → staging slot → device slot |
+| Host weight cache mmap with pinned memory enabled and registration available | Rank-local / no-AllGather | Complete mapping is registered by the platform backend | registered mmap views → device slot |
 
 The three axes therefore remain independent:
 
@@ -453,14 +460,15 @@ 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 host weight 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.
+For a host weight cache mapping, the default pinned-memory policy attempts to
+replace those host slots with direct copies from the shared final tensor views.
+Disabling pinned memory or a failed registration retains 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. Those ordinary pinned tensors remain the default. When the host
-weight cache or a positive registration budget is explicitly selected,
+When direct mmap preflight fails or the host weight cache is selected, the
+regular model loader remains responsible for preparing each rank's weights,
+including TP-local tensors or HSDP-managed parameters. Ordinary pinned tensors
+remain the default without cache selection. When the host weight cache is
+explicitly selected,
 supported CPU layouts enter the cache; unsupported layouts keep one private
 runtime copy per process.
 
@@ -512,13 +520,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. With the
-default zero registration budget they feed the existing bounded two-slot host
-staging path. Every streamed block is first packed from pageable mmap storage
-into one of those pinned slots and is then copied to the device. This removes
-model-sized private pinned storage, but adds a recurrent host-to-host copy and
-can amplify TP synchronization waits when several ranks compete for host
-bandwidth.
+streaming hooks. The source mappings remain open through DLO disable. With
+pinned memory disabled, or when registration is unavailable, they feed the
+existing bounded two-slot host-staging path. Every streamed block is first
+packed from pageable mmap storage into a staging slot and is then copied to the
+device. This removes model-sized private pinned storage, but adds a recurrent
+host-to-host copy and can amplify TP synchronization waits when several ranks
+compete for host bandwidth.
 
 This version is deliberately steady-state-oriented. It requires roughly one
 local-disk model copy per runtime identity, has no eviction, and still performs
@@ -529,19 +537,20 @@ modeled safely.
 
 ### Registered host weight cache transfer
 
-`dlo_host_weight_cache_pin_limit_gib > 0` opts a worker into registration of
-the complete host weight cache mapping when its platform supports it. It also
-changes loader selection: no-AllGather uses the ordinary loader and
-final-layout host weight cache even when a TP1 direct-checkpoint plan is
-available. A checkpoint binding can carry a deferred adapter such as
+Selecting the host weight cache with `pin_cpu_memory=True` opts a worker into
+registration of the complete mapping when its platform supports it. It also
+changes loader selection: no-AllGather uses the ordinary loader and final-layout
+host weight 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 host weight 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 platform mutation. The CUDA backend
+adjacent ranges within that mapping. When a positive per-worker ceiling is
+configured, the complete byte count is checked against it before any platform
+mutation. The CUDA backend
 first requires `cudaDevAttrHostRegisterReadOnlySupported`, then registers every
 range with the read-only host-registration flag and verifies that PyTorch
 recognizes every mapped tensor as pinned. A device without that capability
@@ -642,10 +651,10 @@ artifacts before making hardware or performance claims. Track that work in
   direct-mmap deployments can share checkpoint pages per node. Enable the
   host weight cache only when independent replicas have repeated final layouts
   and benchmark the target CPU/memory topology, especially for matching
-  coordinates across multiple TP engines. On CUDA, provide an explicit
-  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.
+  coordinates across multiple TP engines. With the default pinned-memory
+  policy, supported CUDA workers prefer registered direct H2D. Configure a
+  positive ceiling when an explicit page-locking bound is required, or disable
+  pinned memory when pageable staging is operationally preferable.
 
 ### HSDP
 

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -38,17 +38,18 @@ vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --dlo-no-use-allgather
 
-# Share final host weights with bounded pinned staging
+# Share final host weights; prefer registered direct H2D under the pinned-memory policy
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --dlo-no-use-allgather \
   --dlo-use-host-weight-cache
 
-# Share final host weights and register up to 80 GiB per worker for direct H2D
+# Apply an 80 GiB per-worker safety ceiling to the same registration attempt
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --tensor-parallel-size 2 \
   --dlo-no-use-allgather \
+  --dlo-use-host-weight-cache \
   --dlo-host-weight-cache-pin-limit-gib 80
 ```
 
@@ -58,15 +59,18 @@ vllm serve /path/to/model --omni \
 | --- | --- | --- |
 | `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
 | `--dlo-no-use-allgather` | Stream complete rank-local blocks independently | `false` |
-| `--dlo-use-host-weight-cache` | Use the final-layout cache after checkpoint mmap fallback | `false` |
-| `--dlo-host-weight-cache-pin-limit-gib GIB` | Per-worker registration budget; positive values also select the cache | `0` |
+| `--dlo-use-host-weight-cache` | Use the final-layout cache and prefer registration when pinned memory is enabled | `false` |
+| `--dlo-host-weight-cache-pin-limit-gib GIB` | Optional per-worker registration ceiling; zero means no additional ceiling | `0` |
 | `--dlo-resident-layers N` | Keep eligible leading DiT blocks resident in HBM | `0` |
 
 No-AllGather DLO automatically uses a compatible checkpoint mmap plan. When
 checkpoint mmap is unavailable, it preserves ordinary pinned loader weights by
-default. `--dlo-use-host-weight-cache` explicitly selects the final-layout
-cache with bounded staging; a positive registration budget also selects the
-cache and attempts direct H2D. The cache uses
+default. `--dlo-use-host-weight-cache` explicitly selects the transform-complete
+final-layout cache. Under the default `pin_cpu_memory=True` policy, DLO attempts
+to register the complete immutable mapping for direct H2D. Setting
+`pin_cpu_memory=False` through structured or programmatic configuration retains
+the staged mmap path instead. A positive pin limit is only an optional safety
+ceiling and does not select the cache. The cache uses
 `~/.cache/vllm-omni/dlo-host-weights`. Programmatic configuration may override
 `dlo_host_weight_cache_dir` and the writer lock timeout; these advanced storage
 controls are intentionally not separate CLI flags.
@@ -77,15 +81,16 @@ controls are intentionally not separate CLI flags.
   directory. Do not use tmpfs or a cross-node filesystem for this version.
 - Cache entries are immutable and validated before use. A cache or registration
   failure keeps the ordinary loader weights or falls back to bounded staging.
-- The bounded-staging cache is a memory-first mode: recurrent pageable-mmap to
+- The bounded-staging fallback is a memory-first mode: recurrent pageable-mmap to
   pinned-slot copies can reduce throughput, especially across concurrent TP
-  engines. Keep ordinary pinned weights unless the host-memory saving is worth
-  that tradeoff, or use a validated registered path.
-- A positive registration budget must cover the complete page-aligned mapping
-  reported in the worker log. Registration is all-or-nothing.
+  engines. Use it intentionally by disabling pinned host memory, or when the
+  active platform cannot register the immutable mapping.
+- When a positive registration ceiling is configured, it must cover the
+  complete page-aligned mapping reported in the worker log. Registration is
+  all-or-nothing. Zero applies no additional ceiling.
 - CUDA is the first registration backend and requires read-only host
   registration support for immutable cache mappings. Platforms or devices
-  without that capability continue to use bounded staging; a positive budget
+  without that capability continue to use bounded staging; pinned-memory policy
   requests direct H2D but does not guarantee it.
 - Registration is process-local but does not duplicate the underlying file
   pages. Each worker must still satisfy its platform and OS page-locking limits.

tests/config/test_omni_config.py

@@ -933,6 +933,15 @@ def test_structured_diffusion_config_restricts_host_weight_cache_to_no_allgather
         )
 
 
+def test_structured_diffusion_config_requires_cache_for_registration_ceiling():
+    with pytest.raises(ValueError, match="requires dlo_use_host_weight_cache=True"):
+        omni_config_module._DiffusionConfigProjection(
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=False,
+            dlo_host_weight_cache_pin_limit_gib=1.0,
+        )
+
+
 @pytest.mark.parametrize(
     "field_name",
     [

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -498,7 +498,8 @@ def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypat
     assert loader.take_host_weight_plan() is runtime_plan
 
 
-def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(monkeypatch, tmp_path):
+@pytest.mark.parametrize("pin_cpu_memory", [True, False])
+def test_dlo_host_weight_cache_prefers_final_layout_over_checkpoint_plan(monkeypatch, tmp_path, pin_cpu_memory):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
     import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
 
@@ -515,8 +516,10 @@ def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(m
         quantization_config=None,
         enable_distributed_layerwise_offload=True,
         dlo_use_allgather=False,
+        dlo_use_host_weight_cache=True,
+        pin_cpu_memory=pin_cpu_memory,
         dlo_host_weight_cache_dir=str(tmp_path),
-        dlo_host_weight_cache_pin_limit_gib=1.0,
+        dlo_host_weight_cache_pin_limit_gib=0.0,
         model="unused",
     )
     loader = DiffusersPipelineLoader(LoadConfig(), od_config)
@@ -536,7 +539,7 @@ def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(m
     loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
 
     def unexpected_checkpoint_plan(*_args, **_kwargs):
-        pytest.fail("registered host weight cache mode must not select the raw checkpoint mmap plan")
+        pytest.fail("selected host weight cache mode must not select the raw checkpoint mmap plan")
 
     def build_cache(_pipeline, **_kwargs):
         events.append("cache")

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -119,6 +119,20 @@ def test_platform_factory_dispatches_cuda(monkeypatch: pytest.MonkeyPatch) -> No
     assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=4096) is sentinel
 
 
+def test_platform_factory_allows_registration_without_additional_budget(monkeypatch: pytest.MonkeyPatch) -> None:
+    sentinel = object()
+    sources = {"weights": []}
+
+    def create(actual_sources, *, max_bytes):
+        assert actual_sources is sources
+        assert max_bytes is None
+        return sentinel
+
+    monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
+
+    assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=None) is sentinel
+
+
 def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
     assert _coalesce_ranges(
         [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
@@ -143,6 +157,21 @@ def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytes
     assert runtime.registered == []
 
 
+def test_registration_without_additional_budget_registers_complete_mapping(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)
+
+    registration = CudaHostRegistration.create(
+        {"weights": [_FakeTensor(0x1003, 4096)]},  # type: ignore[list-item]
+        max_bytes=None,
+    )
+
+    assert registration.total_bytes == 8192
+    assert len(runtime.registered) == 1
+    assert registration.close() == []
+
+
 def test_registration_rejects_unsupported_read_only_mapping(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
     runtime = _FakeRuntime([0], read_only_registration_supported=False)

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -1207,6 +1207,68 @@ class TestMmapWeightLoading:
 
         assert events == ["unregister", "mmap"]
 
+    def test_host_weight_cache_registration_defaults_to_pinned_memory_policy(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        patched_offload_runtime,
+    ):
+        class Registration:
+            total_bytes = 4096
+            region_count = 1
+
+            @staticmethod
+            def close() -> list[str]:
+                return []
+
+        registration = Registration()
+        create_calls: list[tuple[dict[str, list[torch.Tensor]], int | None]] = []
+
+        def register(sources, *, device, max_bytes):
+            assert device == torch.device("cuda")
+            create_calls.append((sources, max_bytes))
+            return registration
+
+        monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dlo_use_allgather=False,
+                dlo_host_weight_cache_pin_limit_gib=0.0,
+            ),
+            torch.device("cuda"),
+        )
+        sources = {"runtime.safetensors": [torch.ones(1)]}
+        backend._host_weight_cache_mapped_sources = sources
+
+        assert backend._try_register_host_weight_cache_mmap()
+        assert create_calls == [(sources, None)]
+        assert backend._host_registration is registration
+
+    def test_host_weight_cache_registration_respects_disabled_pinned_memory(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        patched_offload_runtime,
+    ):
+        monkeypatch.setattr(
+            dist_backend_module,
+            "register_host_mappings",
+            lambda *_args, **_kwargs: pytest.fail("registration must follow pin_cpu_memory=False"),
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+                dlo_host_weight_cache_pin_limit_gib=0.0,
+            ),
+            torch.device("cuda"),
+        )
+        backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
+
+        assert not backend._try_register_host_weight_cache_mmap()
+        assert backend._host_registration is None
+
     def test_host_weight_cache_registration_falls_back_on_unsupported_platform(
         self,
         patched_offload_runtime,

vllm_omni/config/omni_config.py

@@ -765,6 +765,8 @@ class _DiffusionConfigProjection:
             raise ValueError("dlo_host_weight_cache_pin_limit_gib requires no-AllGather distributed layerwise offload")
         if self.dlo_use_host_weight_cache and (not self.enable_distributed_layerwise_offload or self.dlo_use_allgather):
             raise ValueError("dlo_use_host_weight_cache requires no-AllGather distributed layerwise offload")
+        if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_use_host_weight_cache:
+            raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_use_host_weight_cache=True")
 
         if self.tf_model_config is None:
             self.tf_model_config = TransformerConfig()

vllm_omni/diffusion/data.py

@@ -708,16 +708,15 @@ 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
-    # Opt into publishing/joining the final-layout host weight cache when a
-    # compatible direct-checkpoint mmap plan is unavailable. A positive pin
-    # limit also selects the cache because registration requires its layout.
+    # Opt into publishing/joining the transform-complete final-layout host
+    # weight cache instead of direct-checkpoint mmap.
     dlo_use_host_weight_cache: bool = False
     # Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-host-weights.
     dlo_host_weight_cache_dir: str | None = None
     # Maximum wait for the per-layout POSIX writer lock.
     dlo_host_weight_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.
+    # Optional maximum file-backed memory a worker may register with CUDA.
+    # Zero applies no additional ceiling; pin_cpu_memory controls registration.
     dlo_host_weight_cache_pin_limit_gib: float = 0.0
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
@@ -957,6 +956,8 @@ class OmniDiffusionConfig:
             raise ValueError("dlo_host_weight_cache_pin_limit_gib requires no-AllGather distributed layerwise offload")
         if self.dlo_use_host_weight_cache and (not self.enable_distributed_layerwise_offload or self.dlo_use_allgather):
             raise ValueError("dlo_use_host_weight_cache requires no-AllGather distributed layerwise offload")
+        if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_use_host_weight_cache:
+            raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_use_host_weight_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,15 +441,13 @@ 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_host_weight_cache = (
-                    _dist_offload
-                    and not _use_ag
-                    and float(getattr(self.od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0)) > 0
+                _select_host_weight_cache = (
+                    _dist_offload and not _use_ag and bool(getattr(self.od_config, "dlo_use_host_weight_cache", False))
                 )
 
                 plan_result = None
                 weight_sources = self._get_weight_sources(model)
-                if _dist_offload and not _prefer_registered_host_weight_cache:
+                if _dist_offload and not _select_host_weight_cache:
                     modules = ModuleDiscovery.discover(model)
                     plan_result = build_checkpoint_mmap_plan(
                         model,
@@ -461,14 +459,14 @@ class DiffusersPipelineLoader:
                         online_quantization=_has_online_quant,
                     )
                     self.host_weight_plan = plan_result.plan
-                elif _prefer_registered_host_weight_cache:
+                elif _select_host_weight_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 host weight cache registration requested; using the ordinary loader "
+                        "DLO host weight cache selected; using the ordinary loader "
                         "once to publish a transform-complete final mmap layout"
                     )
 
@@ -550,10 +548,7 @@ class DiffusersPipelineLoader:
             return
         if getattr(self.od_config, "dlo_use_allgather", True):
             return
-        if not (
-            getattr(self.od_config, "dlo_use_host_weight_cache", False)
-            or float(getattr(self.od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0)) > 0
-        ):
+        if not getattr(self.od_config, "dlo_use_host_weight_cache", False):
             logger.info(
                 "DLO host weight cache not selected; retaining ordinary host tensors after checkpoint mmap fallback"
             )

vllm_omni/diffusion/offloader/base.py

@@ -31,8 +31,8 @@ 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 host weight cache mappings with CUDA.
-    # Zero retains bounded host staging.
+    # Optional per-worker ceiling for registering shared host weight cache mappings.
+    # Zero applies no additional ceiling; pin_cpu_memory controls registration.
     dlo_host_weight_cache_pin_limit_gib: float = 0.0
 
     @classmethod

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -182,13 +182,11 @@ class CudaHostRegistration:
         cls,
         sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
         *,
-        max_bytes: int,
+        max_bytes: int | None,
     ) -> CudaHostRegistration:
-        if max_bytes <= 0:
-            raise HostRegistrationBudgetError("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:
+        if max_bytes is not None and total_bytes > max_bytes:
             raise HostRegistrationBudgetError(
                 f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
             )

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -1269,20 +1269,16 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         )
 
     def _try_register_host_weight_cache_mmap(self) -> bool:
-        """Register complete cache mappings when the operator budget permits."""
+        """Register complete cache mappings under the pinned-memory policy."""
         pin_limit_gib = self.config.dlo_host_weight_cache_pin_limit_gib
-        if pin_limit_gib <= 0:
-            return False
         if not self.config.pin_cpu_memory:
-            logger.warning(
-                "DLO host weight cache registration requires pin_cpu_memory=True; using bounded host staging"
-            )
+            logger.info("DLO host weight cache registration disabled by pin_cpu_memory=False; using host staging")
             return False
         if not self._host_weight_cache_mapped_sources:
             logger.warning("DLO host weight cache registration found no mapped sources; using bounded host staging")
             return False
 
-        max_bytes = int(pin_limit_gib * 1024**3)
+        max_bytes = int(pin_limit_gib * 1024**3) if pin_limit_gib > 0 else None
         started = time.perf_counter()
         try:
             registration = register_host_mappings(

vllm_omni/diffusion/offloader/host_registration.py

@@ -38,12 +38,13 @@ def register_host_mappings(
     sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
     *,
     device: torch.device,
-    max_bytes: int,
+    max_bytes: int | None,
 ) -> HostRegistration:
     """Register mappings with the active platform or report unsupported use.
 
     CUDA is the first implementation. Other platforms retain DLO's bounded
     host-staging path until they provide an equivalent registration backend.
+    ``max_bytes=None`` applies no additional ceiling to the complete mapping.
     """
     if device.type == "cuda":
         from .cuda_host_registration import CudaHostRegistration

vllm_omni/entrypoints/cli/serve.py

@@ -726,8 +726,9 @@ class OmniServeCommand(CLISubcommand):
             "--dlo-use-host-weight-cache",
             action="store_true",
             help=(
-                "Use the node-local final-layout host weight cache when direct checkpoint mmap is unavailable. "
-                "This trades recurrent host staging work for shared host pages unless direct registration succeeds."
+                "Use the node-local final-layout host weight cache for no-AllGather DLO. "
+                "With the default pinned-memory policy, supported platforms register the immutable mapping "
+                "for direct H2D; otherwise DLO uses bounded host staging."
             ),
         )
         omni_config_group.add_argument(
@@ -735,9 +736,9 @@ class OmniServeCommand(CLISubcommand):
             type=float,
             default=0.0,
             help=(
-                "Maximum GiB of read-only host weight cache mappings each worker may "
-                "register for direct H2D when the platform supports it. Zero keeps "
-                "bounded pinned staging."
+                "Optional maximum GiB of read-only host weight cache mappings each worker may "
+                "register for direct H2D. Zero applies no additional registration ceiling; "
+                "this option requires --dlo-use-host-weight-cache."
             ),
         )
         omni_config_group.add_argument(