[Diffusion] Make DLO host weight cache opt-in

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -44,10 +44,12 @@ DLO weight collective.
 Host storage is selected separately from the transfer protocol. The loader can
 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 automatically normalizes supported 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. An opt-in host-registration budget can make the complete final mapping
+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.
@@ -130,9 +132,12 @@ 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
+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 deliberately selects the transform-complete host weight cache;
-the backend later consumes the same budget for registration.
+positive value also deliberately selects the transform-complete host weight
+cache; the backend later consumes the same budget for registration.
 `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
@@ -425,8 +430,10 @@ 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 then enter the host weight cache
-automatically. Unsupported layouts keep one private runtime copy per process.
+parameters. Those ordinary pinned tensors remain the default. When the host
+weight cache or a positive registration budget is explicitly selected,
+supported CPU layouts enter the cache; unsupported layouts keep one private
+runtime copy per process.
 
 This mode means:
 
@@ -447,7 +454,8 @@ The Phase B cache preserves the same ownership boundary. DLO neither derives
 cache keys nor interprets model loaders. No-AllGather DLO first attempts a
 compatible direct-checkpoint mmap plan. When that plan is unavailable, every
 rank completes the ordinary loader, custom callbacks, post-load processing,
-strict validation, and calibration. The loader then automatically:
+strict validation, and calibration. When cache use or registration is
+explicitly selected, the loader:
 
 1. discovers final CPU DiT parameters and persistent buffers;
 2. rejects unsupported non-contiguous, aliased, quantized, distributed, or

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -38,6 +38,12 @@ vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --dlo-no-use-allgather
 
+# Share final host weights with bounded pinned staging
+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
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
@@ -52,12 +58,15 @@ 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-host-weight-cache-pin-limit-gib GIB` | Per-worker registration budget; zero uses bounded host staging | `0` |
+| `--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-resident-layers N` | Keep eligible leading DiT blocks resident in HBM | `0` |
 
-No-AllGather DLO automatically uses a compatible checkpoint mmap plan, or
-builds/joins the final-layout host weight cache after ordinary loading when a
-checkpoint plan is unavailable. The cache uses
+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
 `~/.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.
@@ -68,6 +77,10 @@ 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
+  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.
 - CUDA is the first registration backend. Platforms without an equivalent

tests/config/test_omni_config.py

@@ -912,12 +912,24 @@ def test_structured_diffusion_config_rejects_invalid_compile_granularity():
     ("enable_dlo", "use_allgather"),
     [(False, False), (True, True)],
 )
-def test_structured_diffusion_config_restricts_host_weight_registration_to_no_allgather_dlo(enable_dlo, use_allgather):
+@pytest.mark.parametrize(
+    ("field_name", "field_value"),
+    [
+        ("dlo_use_host_weight_cache", True),
+        ("dlo_host_weight_cache_pin_limit_gib", 1.0),
+    ],
+)
+def test_structured_diffusion_config_restricts_host_weight_cache_to_no_allgather_dlo(
+    enable_dlo,
+    use_allgather,
+    field_name,
+    field_value,
+):
     with pytest.raises(ValueError, match="requires no-AllGather distributed layerwise offload"):
         omni_config_module._DiffusionConfigProjection(
             enable_distributed_layerwise_offload=enable_dlo,
             dlo_use_allgather=use_allgather,
-            dlo_host_weight_cache_pin_limit_gib=1.0,
+            **{field_name: field_value},
         )
 
 

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -375,7 +375,7 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
     assert loader.take_host_weight_plan() is plan
 
 
-def test_dlo_plan_fallback_runs_ordinary_loader_then_cache(monkeypatch):
+def test_dlo_plan_fallback_retains_ordinary_loader_weights_by_default(monkeypatch):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
     import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
 
@@ -402,14 +402,13 @@ def test_dlo_plan_fallback_runs_ordinary_loader_then_cache(monkeypatch):
         lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
     )
 
-    def build_cache(_pipeline, **_kwargs):
-        calls.append("cache")
-        return HostWeightPlanResult(None, "unsupported test layout", "unsupported_layout")
+    def build_cache(*_args, **_kwargs):
+        pytest.fail("host weight cache must require explicit selection")
 
     monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache)
 
     assert loader.load_model(load_device="cpu") is model
-    assert calls == ["load", "process", "cache"]
+    assert calls == ["load", "process"]
     assert loader.take_host_weight_plan() is None
 
 
@@ -434,6 +433,7 @@ def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypat
         quantization_config=None,
         enable_distributed_layerwise_offload=True,
         dlo_use_allgather=False,
+        dlo_use_host_weight_cache=True,
         dlo_host_weight_cache_dir=str(tmp_path),
         dlo_host_weight_cache_lock_timeout=1.0,
         model="unused",
@@ -567,6 +567,7 @@ def test_dlo_host_weight_cache_failure_retains_ordinary_weights(monkeypatch, tmp
         quantization_config=None,
         enable_distributed_layerwise_offload=True,
         dlo_use_allgather=False,
+        dlo_use_host_weight_cache=True,
         dlo_host_weight_cache_dir=str(tmp_path),
         model="unused",
     )

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -328,6 +328,7 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--omni",
             "--enable-distributed-layerwise-offload",
             "--dlo-no-use-allgather",
+            "--dlo-use-host-weight-cache",
             "--dlo-host-weight-cache-pin-limit-gib",
             "80",
             "--dlo-resident-layers",
@@ -341,10 +342,12 @@ def test_serve_cli_forwards_distributed_offload_residency():
 
     assert args.enable_distributed_layerwise_offload is True
     assert args.dlo_use_allgather is False
+    assert args.dlo_use_host_weight_cache is True
     assert args.dlo_host_weight_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_use_host_weight_cache"] is True
     assert engine_args["dlo_host_weight_cache_dir"] is None
     assert engine_args["dlo_host_weight_cache_lock_timeout"] == 600.0
     assert engine_args["dlo_host_weight_cache_pin_limit_gib"] == 80

vllm_omni/config/omni_config.py

@@ -688,6 +688,7 @@ class _DiffusionConfigProjection:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_use_host_weight_cache: bool = False
     dlo_host_weight_cache_dir: str | None = None
     dlo_host_weight_cache_lock_timeout: float = Field(default=600.0, gt=0)
     dlo_host_weight_cache_pin_limit_gib: float = Field(default=0.0, ge=0)
@@ -762,6 +763,8 @@ class _DiffusionConfigProjection:
             not self.enable_distributed_layerwise_offload or self.dlo_use_allgather
         ):
             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.tf_model_config is None:
             self.tf_model_config = TransformerConfig()

vllm_omni/config/stage_config.py

@@ -382,6 +382,7 @@ class StageDeployConfig:
 
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
+    dlo_use_host_weight_cache: bool | None = None
     dlo_host_weight_cache_dir: str | None = None
     dlo_host_weight_cache_lock_timeout: float | None = None
     dlo_host_weight_cache_pin_limit_gib: float | None = None

vllm_omni/diffusion/data.py

@@ -708,6 +708,10 @@ 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.
+    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.
@@ -951,6 +955,8 @@ class OmniDiffusionConfig:
             not self.enable_distributed_layerwise_offload or self.dlo_use_allgather
         ):
             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.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

@@ -550,6 +550,14 @@ 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
+        ):
+            logger.info(
+                "DLO host weight cache not selected; retaining ordinary host tensors after checkpoint mmap fallback"
+            )
+            return
 
         parallel_config = self.parallel_config
         tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1))

vllm_omni/engine/arg_utils.py

@@ -545,6 +545,7 @@ class OrchestratorArgs:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_use_host_weight_cache: bool = False
     dlo_host_weight_cache_dir: str | None = None
     dlo_host_weight_cache_lock_timeout: float = 600.0
     dlo_host_weight_cache_pin_limit_gib: float = 0.0

vllm_omni/engine/async_omni_engine.py

@@ -1040,6 +1040,7 @@ class AsyncOmniEngine:
             "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
+            "dlo_use_host_weight_cache": kwargs.get("dlo_use_host_weight_cache", False),
             "dlo_host_weight_cache_dir": kwargs.get("dlo_host_weight_cache_dir", None),
             "dlo_host_weight_cache_lock_timeout": kwargs.get("dlo_host_weight_cache_lock_timeout", 600.0),
             "dlo_host_weight_cache_pin_limit_gib": kwargs.get("dlo_host_weight_cache_pin_limit_gib", 0.0),

vllm_omni/entrypoints/cli/serve.py

@@ -722,6 +722,14 @@ class OmniServeCommand(CLISubcommand):
                 "independently (including existing TP shards)."
             ),
         )
+        omni_config_group.add_argument(
+            "--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."
+            ),
+        )
         omni_config_group.add_argument(
             "--dlo-host-weight-cache-pin-limit-gib",
             type=float,