[Diffusion] Prefer host-cache registration when pinned

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

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