[Diffusion] Make no-AllGather host weight cache automatic

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -43,13 +43,14 @@ 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, an opt-in Phase B
-host weight 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. 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.
+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
+directly H2D-copyable when the platform provides a backend, avoiding the
+recurrent host packing otherwise required by mmap. 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
@@ -128,14 +129,14 @@ flowchart TB
 | `DiffusionWorker` | Calling backend teardown before destroying the distributed environment | Managing backend-owned mappings or registrations directly |
 
 The configuration fields are routed to the component that can enforce their
-contract. `dlo_enable_host_weight_cache`, its cache root, and lock timeout are
-consumed by the loader. `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. `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 model-specific global state.
+contract. The cache root and lock timeout are consumed by the loader.
+`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.
+`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
+model-specific global state.
 
 ### Shared startup and lifecycle
 
@@ -424,8 +425,8 @@ 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 host weight cache.
-Unsupported layouts keep one private runtime copy per process.
+parameters. Supported CPU layouts then enter the host weight cache
+automatically. Unsupported layouts keep one private runtime copy per process.
 
 This mode means:
 
@@ -443,9 +444,10 @@ This mode means:
 ### Post-load host weight cache
 
 The Phase B cache preserves the same ownership boundary. DLO neither derives
-cache keys nor interprets model loaders. When enabled with no-AllGather, every
-rank first completes the ordinary loader, custom callbacks, post-load
-processing, strict validation, and calibration. The loader then:
+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:
 
 1. discovers final CPU DiT parameters and persistent buffers;
 2. rejects unsupported non-contiguous, aliased, quantized, distributed, or
@@ -481,9 +483,10 @@ 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 opt-in and steady-state-oriented. It requires
-roughly one local-disk model copy per runtime identity, has no eviction, and
-still performs ordinary loading plus full content-hash passes in every rank.
+This version is deliberately steady-state-oriented. It requires roughly one
+local-disk model copy per runtime identity, has no eviction, and still performs
+ordinary loading plus full content-hash passes in every rank. Cache publication
+or validation failure falls back to the already-valid ordinary host tensors.
 Skip-load-on-hit is a separate follow-up after loader side effects can be
 modeled safely.
 

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -43,7 +43,6 @@ vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --tensor-parallel-size 2 \
   --dlo-no-use-allgather \
-  --dlo-enable-host-weight-cache \
   --dlo-host-weight-cache-pin-limit-gib 80
 ```
 
@@ -53,14 +52,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-enable-host-weight-cache` | Share final no-AllGather DiT weights through a node-local mmap cache | `false` |
 | `--dlo-host-weight-cache-pin-limit-gib GIB` | Per-worker registration budget; zero uses bounded host staging | `0` |
 | `--dlo-resident-layers N` | Keep eligible leading DiT blocks resident in HBM | `0` |
 
-The host weight 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.
+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
+`~/.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.
 
 ## Operational notes
 

tests/config/test_omni_config.py

@@ -908,6 +908,19 @@ def test_structured_diffusion_config_rejects_invalid_compile_granularity():
         omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
 
 
+@pytest.mark.parametrize(
+    ("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):
+    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,
+        )
+
+
 def test_from_pipeline_config_matches_stage_config_to_omegaconf_behavior_for_representative_stage():
     pipeline = _resolve_pipeline_or_skip("qwen3_tts")
     legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0]

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -375,8 +375,9 @@ 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(monkeypatch):
+def test_dlo_plan_fallback_runs_ordinary_loader_then_cache(monkeypatch):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
 
     od_config = SimpleNamespace(
         dtype=torch.float32,
@@ -401,8 +402,14 @@ def test_dlo_plan_fallback_runs_ordinary_loader(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")
+
+    monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache)
+
     assert loader.load_model(load_device="cpu") is model
-    assert calls == ["load", "process"]
+    assert calls == ["load", "process", "cache"]
     assert loader.take_host_weight_plan() is None
 
 
@@ -427,7 +434,6 @@ 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_enable_host_weight_cache=True,
         dlo_host_weight_cache_dir=str(tmp_path),
         dlo_host_weight_cache_lock_timeout=1.0,
         model="unused",
@@ -509,7 +515,6 @@ 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_enable_host_weight_cache=True,
         dlo_host_weight_cache_dir=str(tmp_path),
         dlo_host_weight_cache_pin_limit_gib=1.0,
         model="unused",
@@ -562,7 +567,6 @@ 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_enable_host_weight_cache=True,
         dlo_host_weight_cache_dir=str(tmp_path),
         model="unused",
     )

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -328,7 +328,6 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--omni",
             "--enable-distributed-layerwise-offload",
             "--dlo-no-use-allgather",
-            "--dlo-enable-host-weight-cache",
             "--dlo-host-weight-cache-pin-limit-gib",
             "80",
             "--dlo-resident-layers",
@@ -342,12 +341,10 @@ 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_enable_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_enable_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

@@ -687,7 +687,6 @@ class _DiffusionConfigProjection:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
-    dlo_enable_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)
@@ -753,12 +752,10 @@ class _DiffusionConfigProjection:
         from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
         from vllm_omni.quantization import build_quant_config
 
-        if self.dlo_enable_host_weight_cache and not self.enable_distributed_layerwise_offload:
-            raise ValueError("dlo_enable_host_weight_cache requires distributed layerwise offload")
-        if self.dlo_enable_host_weight_cache and self.dlo_use_allgather:
-            raise ValueError("dlo_enable_host_weight_cache currently requires dlo_use_allgather=False")
-        if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_enable_host_weight_cache:
-            raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_enable_host_weight_cache=True")
+        if self.dlo_host_weight_cache_pin_limit_gib and (
+            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.tf_model_config is None:
             self.tf_model_config = TransformerConfig()

vllm_omni/config/stage_config.py

@@ -382,7 +382,6 @@ class StageDeployConfig:
 
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
-    dlo_enable_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,11 +708,6 @@ 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. This is
-    # 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_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.
@@ -952,12 +947,10 @@ class OmniDiffusionConfig:
             raise ValueError("dlo_host_weight_cache_lock_timeout must be positive")
         if not math.isfinite(self.dlo_host_weight_cache_pin_limit_gib) or self.dlo_host_weight_cache_pin_limit_gib < 0:
             raise ValueError("dlo_host_weight_cache_pin_limit_gib must be non-negative")
-        if self.dlo_enable_host_weight_cache and not self.enable_distributed_layerwise_offload:
-            raise ValueError("dlo_enable_host_weight_cache requires distributed layerwise offload")
-        if self.dlo_enable_host_weight_cache and self.dlo_use_allgather:
-            raise ValueError("dlo_enable_host_weight_cache currently requires dlo_use_allgather=False")
-        if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_enable_host_weight_cache:
-            raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_enable_host_weight_cache=True")
+        if self.dlo_host_weight_cache_pin_limit_gib and (
+            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.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

@@ -444,7 +444,6 @@ class DiffusersPipelineLoader:
                 _prefer_registered_host_weight_cache = (
                     _dist_offload
                     and not _use_ag
-                    and bool(getattr(self.od_config, "dlo_enable_host_weight_cache", False))
                     and float(getattr(self.od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0)) > 0
                 )
 
@@ -549,10 +548,7 @@ class DiffusersPipelineLoader:
             return
         if not getattr(self.od_config, "enable_distributed_layerwise_offload", False):
             return
-        if not getattr(self.od_config, "dlo_enable_host_weight_cache", False):
-            return
         if getattr(self.od_config, "dlo_use_allgather", True):
-            logger.info("DLO host weight cache is no-AllGather only; leaving AllGather storage unchanged")
             return
 
         parallel_config = self.parallel_config

vllm_omni/engine/arg_utils.py

@@ -545,7 +545,6 @@ class OrchestratorArgs:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
-    dlo_enable_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,7 +1040,6 @@ 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_enable_host_weight_cache": kwargs.get("dlo_enable_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,15 +722,6 @@ class OmniServeCommand(CLISubcommand):
                 "independently (including existing TP shards)."
             ),
         )
-        omni_config_group.add_argument(
-            "--dlo-enable-host-weight-cache",
-            action="store_true",
-            help=(
-                "After ordinary loading, publish or join a node-local mmap cache "
-                "for final DiT weights. This first version requires "
-                "--dlo-no-use-allgather."
-            ),
-        )
         omni_config_group.add_argument(
             "--dlo-host-weight-cache-pin-limit-gib",
             type=float,