[Refactor][Diffusion] Read the resolved offload policy outside the config layer (#7327)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/models/magi2/test_pipeline_magi2.py

@@ -78,10 +78,15 @@ class _ParallelStub:
 @dataclass
 class _TopologyStub:
     parallel_config: _ParallelStub = field(default_factory=_ParallelStub)
+    diffusion_offload_config: object = None
     enable_cpu_offload: bool = False
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_resident_layers: int = 0
+    dlo_host_registration_limit_gib: float = 0.0
+    host_weight_runtime_mode: str = "disabled"
+    pin_cpu_memory: bool = True
     quantization_config: object = None
     cache_backend: str = "none"
     custom_pipeline_args: dict[str, object] = field(default_factory=dict)
@@ -383,26 +388,33 @@ def test_native_topology_rejects_nondivisible_tensor_parallelism():
 
 
 def test_native_topology_accepts_single_device_layerwise_with_cpu_staging():
-    config = _topology_config(
-        enable_layerwise_offload=True,
-        additional_config={},
+    _validate_native_topology(
+        _topology_config(
+            enable_layerwise_offload=True,
+            additional_config={},
+        )
+    )
+    _validate_native_topology(
+        _topology_config(
+            enable_cpu_offload=True,
+            enable_layerwise_offload=True,
+            additional_config={},
+        )
     )
-    _validate_native_topology(config)
 
-    config.enable_cpu_offload = True
-    _validate_native_topology(config)
-
-    config.enable_layerwise_offload = False
-    with pytest.raises(ValueError, match="Combine --enable-cpu-offload"):
-        _validate_native_topology(config)
+    with pytest.raises(ValueError, match="already stages its auxiliary components"):
+        _validate_native_topology(
+            _topology_config(
+                enable_cpu_offload=True,
+                additional_config={},
+            )
+        )
 
 
 def test_native_topology_requires_dlo_rank_local_mode():
-    config = _topology_config(enable_distributed_layerwise_offload=True)
     with pytest.raises(ValueError, match="dlo-no-use-allgather"):
-        _validate_native_topology(config)
-    config.dlo_use_allgather = False
-    _validate_native_topology(config)
+        _validate_native_topology(_topology_config(enable_distributed_layerwise_offload=True))
+    _validate_native_topology(_topology_config(enable_distributed_layerwise_offload=True, dlo_use_allgather=False))
 
 
 def test_native_topology_accepts_four_device_hsdp_cfg_and_vae_layouts():

tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py

@@ -13,6 +13,21 @@ class _OffloadAbort(BaseException):
     pass
 
 
+def _legacy_distributed_offload_config() -> SimpleNamespace:
+    return SimpleNamespace(
+        diffusion_offload_config=None,
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        dlo_host_registration_limit_gib=0.0,
+        host_weight_runtime_mode="disabled",
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(data_parallel_size=1),
+    )
+
+
 def test_no_allgather_example_selects_distributed_backend():
     from examples.offline_inference.minimax_h3.dlo_lifecycle import engine_kwargs
     from vllm_omni.diffusion.offloader.base import OffloadConfig
@@ -100,10 +115,7 @@ def test_manual_component_failure_forces_retained_cache_release(mocker):
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
-    pipeline.od_config = mocker.Mock()
-    pipeline.od_config.diffusion_offload_config = None
-    pipeline.od_config.enable_layerwise_offload = False
-    pipeline.od_config.enable_distributed_layerwise_offload = True
+    pipeline.od_config = _legacy_distributed_offload_config()
     pipeline._model_cpu_offload_modules = []
     pipeline._dlo_component_cache = mocker.Mock()
     component = mocker.Mock()
@@ -123,10 +135,7 @@ def test_manual_component_offload_failure_forces_retained_cache_release(mocker):
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
-    pipeline.od_config = mocker.Mock()
-    pipeline.od_config.diffusion_offload_config = None
-    pipeline.od_config.enable_layerwise_offload = False
-    pipeline.od_config.enable_distributed_layerwise_offload = True
+    pipeline.od_config = _legacy_distributed_offload_config()
     pipeline._model_cpu_offload_modules = []
     pipeline._dlo_component_cache = mocker.Mock()
     component = mocker.Mock()