Merge remote-tracking branch 'origin/main' into qwen-image-2.1

# Conflicts:
#	vllm_omni/diffusion/executor/multiproc_executor.py
#	vllm_omni/diffusion/offloader/sequential_backend.py

tests/diffusion/offloader/test_backend_plan_contract.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Executable contracts for the two behaviors the generic cutover changed.
+
+Placement now follows resolved block identity, and the backends run on the
+resolved plan alone. Declaration parsing, selection errors, rollback,
+residency and enable/disable cycles are covered by ``test_plan_resolver.py``,
+``test_layerwise_backend.py`` and ``test_sequential_backend.py``.
+"""
+
+from typing import ClassVar
+
+import pytest
+import torch
+from torch import nn
+
+from tests.diffusion.offloader.helpers import patch_offload_runtime
+from vllm_omni.diffusion.offloader import layerwise_backend, sequential_backend
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
+from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan
+from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model]
+
+STRATEGIES = [OffloadStrategy.MODEL_LEVEL, OffloadStrategy.LAYER_WISE]
+CPU = torch.device("cpu")
+
+
+class _Block(nn.Linear):
+    def __init__(self):
+        super().__init__(4, 4)
+        self.register_buffer("scale", torch.tensor(0.5))
+
+    def forward(self, x):
+        return super().forward(x).tanh() * self.scale
+
+
+class _Stack(nn.Module):
+    """Two block containers plus non-block state that must stay resident."""
+
+    def __init__(self):
+        super().__init__()
+        self.blocks = nn.ModuleList([_Block() for _ in range(3)])
+        self.tail = nn.ModuleList([_Block() for _ in range(2)])
+        self.bias = nn.Parameter(torch.ones(4))
+        self.register_buffer("scale", torch.tensor(0.25))
+        self.proj = nn.Linear(4, 4)
+
+    def forward(self, x):
+        for block in (*self.blocks, *self.tail):
+            x = block(x)
+        return self.proj(x) * self.scale + self.bias
+
+
+class _Pipeline(nn.Module):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = ["resident"]
+    _offload_plan = OffloadPlan(
+        block_attrs={"transformer": ("blocks", "tail")},
+        encoder_block_attrs={"text_encoder": ("blocks", "tail")},
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _Stack()
+        self.text_encoder = _Stack()
+        self.image_encoder = nn.Linear(4, 4)
+        self.vae = nn.Linear(4, 4)
+        self.resident = nn.Linear(4, 4)
+
+    def forward(self, x):
+        return self.resident(self.vae(self.transformer(self.text_encoder(self.image_encoder(x)))))
+
+
+@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)])
+def execution_device(request, monkeypatch):
+    if request.param == "cuda":
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA required for real transfer/residency coverage")
+        return torch.device("cuda:0")
+    patch_offload_runtime(monkeypatch, current_omni_platform, synchronize=True)
+    monkeypatch.setattr(current_omni_platform, "get_free_memory", lambda: 0)
+    return CPU
+
+
+def _config(strategy, components):
+    return OffloadConfig(strategy=strategy, components=components, pin_cpu_memory=False)
+
+
+def _backend(config, device):
+    kind = ModelLevelOffloadBackend if config.strategy is OffloadStrategy.MODEL_LEVEL else LayerWiseOffloadBackend
+    return kind(config, device)
+
+
+def _assert_no_offload_hooks(pipeline):
+    for module in pipeline.modules():
+        registry = getattr(module, "_hook_registry", None)
+        if registry is not None:
+            assert registry.get_hook("sequential_offload") is None
+            assert registry.get_hook("layerwise_offload") is None
+        assert not getattr(module, "_omni_layerwise_enabled", False)
+
+
+@torch.inference_mode()
+def test_alias_of_a_streamed_block_is_not_placed(execution_device):
+    """Placement moves DiT state by resolved block identity, not by attribute.
+
+    An attribute aliasing a streamed block must stay with its ring: copying it
+    to the device makes the ring record the device as that block's home and
+    strands those weights there after teardown.
+    """
+    pipeline = _Pipeline()
+    pipeline.transformer.proj.weight = pipeline.transformer.blocks[1].weight
+    original = pipeline.transformer.proj.weight.clone()
+    backend = _backend(_config(OffloadStrategy.LAYER_WISE, frozenset({"dit"})), execution_device)
+
+    backend.enable(pipeline)
+    try:
+        assert pipeline.transformer.proj.weight is pipeline.transformer.blocks[1].weight
+        assert pipeline.transformer.proj.weight.device == CPU
+        # Non-block state around the alias is still placed.
+        assert pipeline.transformer.proj.bias.device == execution_device
+        assert pipeline.transformer.bias.device == execution_device
+    finally:
+        backend.disable()
+
+    assert pipeline.transformer.proj.weight.device == CPU
+    torch.testing.assert_close(pipeline.transformer.proj.weight, original)
+
+
+@pytest.mark.parametrize("strategy", STRATEGIES)
+@torch.inference_mode()
+def test_backends_execute_the_resolved_plan_alone(execution_device, strategy, monkeypatch):
+    pipeline = _Pipeline().to(execution_device)
+    x = torch.randn(2, 4, device=execution_device)
+    expected = pipeline(x)
+    config = _config(strategy, frozenset({"dit", "text_encoder"}))
+    resolved = resolve_offload_plan(pipeline, config)
+
+    def forbidden_read(*args, **kwargs):
+        pytest.fail("Backend reinterpreted topology after plan resolution")
+
+    # Once resolved, neither declarations nor selector helpers are an input to
+    # backend execution. Transport options stay readable.
+    monkeypatch.setattr(config, "offloads", forbidden_read)
+    monkeypatch.setattr(config, "should_offload_encoder", forbidden_read)
+    monkeypatch.setattr(pipeline, "_offload_plan", None)
+    backend_module = sequential_backend if strategy is OffloadStrategy.MODEL_LEVEL else layerwise_backend
+    monkeypatch.setattr(backend_module, "resolve_offload_plan", lambda *_: resolved)
+
+    backend = _backend(config, execution_device)
+    try:
+        backend.enable(pipeline)
+        torch.testing.assert_close(pipeline(x), expected)
+    finally:
+        backend.disable()
+    _assert_no_offload_hooks(pipeline)

tests/diffusion/offloader/test_diffusion_cpu_offload.py

@@ -9,14 +9,14 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.helpers.mark import hardware_test
+from tests.helpers.mark import hardware_marks, hardware_test
 from tests.helpers.monitor import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
 AUDIO_MODEL: dict[str, dict[str, int | None]] = {
-    "stabilityai/stable-audio-open-1.0": {"cuda": 100, "rocm": None},
+    "stabilityai/stable-audio-open-1.0": {"cuda": 150, "rocm": None},
 }
 
 IMAGE_MODELS: dict[str, dict[str, int | None]] = {
@@ -27,11 +27,19 @@ MODELS: dict[str, dict[str, int | None]] = {**AUDIO_MODEL, **IMAGE_MODELS}
 
 MODEL_MARKS = {
     "riverclouds/qwen_image_random": pytest.mark.core_model,
-    "stabilityai/stable-audio-open-1.0": pytest.mark.full_model,
+    "stabilityai/stable-audio-open-1.0": [
+        pytest.mark.full_model,
+        *hardware_marks(res={"cuda": ["L4", "B200"], "rocm": "MI325"}),
+    ],
 }
 
 _GATED_MODELS = {"stabilityai/stable-audio-open-1.0"}
 
+# Models whose thresholds are calibrated against max_memory_allocated. All
+# other models keep asserting on the polled device-wide figure until their
+# thresholds have been re-measured on that metric on both CUDA and ROCm.
+_ALLOCATED_THRESHOLD_MODELS = {"stabilityai/stable-audio-open-1.0"}
+
 
 AUDIO_MODEL_PARAMS: dict[str, dict[str, Any]] = {
     "runner_params": {},
@@ -47,7 +55,7 @@ IMAGE_MODEL_PARAMS: dict[str, dict[str, Any]] = {
 }
 
 
-def inference(model_name: str, offload: bool = True) -> tuple[float, Any]:
+def inference(model_name: str, offload: bool = True) -> tuple[float, float, Any]:
     gc.collect()
     current_omni_platform.empty_cache()
     device_index = current_omni_platform.current_device()
@@ -78,12 +86,17 @@ def inference(model_name: str, offload: bool = True) -> tuple[float, Any]:
             ),
         )
     peak = monitor.peak_used_mb
+    # Exact in-process peak of live tensors. Unlike the polled device-wide
+    # figure above, this excludes caching-allocator slack and CUDA context,
+    # which differ between the offload and no-offload paths (the offload hook
+    # calls empty_cache() on every swap) and drift between runs.
+    peak_allocated = current_omni_platform.max_memory_allocated(device=device_index) / (1024**2)
     monitor.stop()
 
     gc.collect()
     current_omni_platform.empty_cache()
 
-    return peak, output
+    return peak, peak_allocated, output
 
 
 def check_audio_determinism(audio1, audio2, atol=1e-2):
@@ -109,9 +122,9 @@ def check_audio_determinism(audio1, audio2, atol=1e-2):
 )
 def test_cpu_offload_diffusion_model(model_name: str):
     try:
-        offload_peak_memory, output_offload = inference(model_name, offload=True)
+        offload_peak_memory, offload_peak_allocated, output_offload = inference(model_name, offload=True)
         cleanup_dist_env_and_memory()
-        no_offload_peak_memory, output_no_offload = inference(model_name, offload=False)
+        no_offload_peak_memory, no_offload_peak_allocated, output_no_offload = inference(model_name, offload=False)
     except ValueError as exc:
         # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing.
         if "Access to model" in str(exc) and "is restricted" in str(exc):
@@ -122,16 +135,14 @@ def test_cpu_offload_diffusion_model(model_name: str):
         pytest.fail(f"Inference failed: {exc}")
     except Exception:
         pytest.fail("Inference failed")
-    print(f"Offload peak memory: {offload_peak_memory} MB")
-    print(f"No offload peak memory: {no_offload_peak_memory} MB")
+    print(f"Offload peak memory: {offload_peak_memory} MB (allocated: {offload_peak_allocated:.1f} MB)")
+    print(f"No offload peak memory: {no_offload_peak_memory} MB (allocated: {no_offload_peak_allocated:.1f} MB)")
 
     if model_name == "stabilityai/stable-audio-open-1.0":
         audio_offload = output_offload[0].multimodal_output.get("audio")
         audio_no_offload = output_no_offload[0].multimodal_output.get("audio")
         check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2)
 
-    # Set platform-specific VRAM saving thresholds to account
-    # for varying runtime memory overhead and fragmentation between CUDA and ROCm.
     is_rocm = torch.version.hip is not None
     platform = "rocm" if is_rocm else "cuda"
     threshold = MODELS[model_name][platform]
@@ -139,7 +150,22 @@ def test_cpu_offload_diffusion_model(model_name: str):
         pytest.skip(f"Threshold not defined for {platform} on {model_name}")
     assert threshold is not None
 
-    assert offload_peak_memory + threshold < no_offload_peak_memory, (
-        f"Offload peak memory {offload_peak_memory} MB should be less than "
-        f"no offload peak memory {no_offload_peak_memory} MB by {threshold} MB"
-    )
+    if model_name in _ALLOCATED_THRESHOLD_MODELS:
+        # Threshold is a lower bound on the peak *allocated* saving, i.e. the
+        # weights the offloader keeps off-GPU at the moment of peak usage
+        # (stable-audio: the fp16 T5 encoder, ~209 MB). max_memory_allocated
+        # is used instead of the polled device-wide figure because the latter
+        # includes caching-allocator slack that differs between the two paths
+        # and drifts by tens of MB between runs.
+        assert offload_peak_allocated + threshold < no_offload_peak_allocated, (
+            f"Offload peak allocated memory {offload_peak_allocated:.1f} MB should be less than "
+            f"no offload peak allocated memory {no_offload_peak_allocated:.1f} MB by {threshold} MB"
+        )
+    else:
+        # Threshold is calibrated against the polled device-wide figure and
+        # accounts for runtime memory overhead and fragmentation differences
+        # between CUDA and ROCm.
+        assert offload_peak_memory + threshold < no_offload_peak_memory, (
+            f"Offload peak memory {offload_peak_memory} MB should be less than "
+            f"no offload peak memory {no_offload_peak_memory} MB by {threshold} MB"
+        )

tests/diffusion/offloader/test_diffusion_layerwise_offload.py

@@ -1,7 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+import concurrent.futures
 import gc
+import multiprocessing as mp
+import os
 from typing import Any
 
 import numpy as np
@@ -9,7 +12,7 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.helpers.mark import hardware_test
+from tests.helpers.mark import hardware_marks
 from tests.helpers.monitor import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -31,8 +34,14 @@ IMAGE_VIDEO_MODELS: dict[str, dict[str, int | None]] = {
 MODELS: dict[str, dict[str, int | None]] = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
 
 MODEL_MARKS = {
-    "riverclouds/qwen_image_random": pytest.mark.core_model,
-    "stabilityai/stable-audio-open-1.0": pytest.mark.full_model,
+    "riverclouds/qwen_image_random": [
+        pytest.mark.core_model,
+        *hardware_marks(res={"cuda": "L4", "rocm": "MI325"}),
+    ],
+    "stabilityai/stable-audio-open-1.0": [
+        pytest.mark.full_model,
+        *hardware_marks(res={"cuda": ["L4", "B200"], "rocm": "MI325"}),
+    ],
 }
 
 AUDIO_MODEL_PARAMS: dict[str, dict[str, Any]] = {
@@ -45,32 +54,98 @@ IMAGE_VIDEO_MODELS_PARAMS: dict[str, dict[str, Any]] = {
     "sampler_params": {"height": 480, "width": 640, "num_frames": 5},
 }
 
+_OFFLOAD_STATE_PROBE = "vllm_omni_layerwise_offload_test"
+# A second CI build can reverse the same exact-head measurement without a
+# source edit; the default keeps the existing baseline-first workload.
+_MEASUREMENT_ORDER_ENV = "VLLM_OMNI_OFFLOAD_TEST_ORDER"
+_MEASUREMENT_ORDERS: dict[str, tuple[bool, bool]] = {
+    "baseline-first": (False, True),
+    "offload-first": (True, False),
+}
 
-def check_audio_determinism(audio1, audio2, atol=1e-2):
-    device = current_omni_platform.device_type
-    if isinstance(audio1, np.ndarray):
-        audio1 = torch.from_numpy(audio1).to(device)
-    if isinstance(audio2, np.ndarray):
-        audio2 = torch.from_numpy(audio2).to(device)
 
-    if not torch.allclose(audio1, audio2, atol=atol):
-        diff = torch.abs(audio1 - audio2)
-        print(f"Max difference: {diff.max().item()}")
-        print(f"Mean difference: {diff.mean().item()}")
+class OffloadStateProbe:
+    """Test-only worker extension exposing the configured offloader state."""
+
+    model_runner: Any
+
+    def get_offload_state_for_test(self) -> dict[str, Any]:
+        backend = getattr(self.model_runner, "offload_backend", None)
+        block_groups = getattr(backend, "_blocks", ())
+        group_sizes = [len(group) for group in block_groups] if isinstance(block_groups, (list, tuple)) else []
+
+        return {
+            "probe": _OFFLOAD_STATE_PROBE,
+            "backend_type": (
+                None if backend is None else f"{backend.__class__.__module__}.{backend.__class__.__qualname__}"
+            ),
+            "requested": bool(getattr(self.model_runner.od_config, "enable_layerwise_offload", False)),
+            "enabled": bool(backend is not None and backend.is_enabled()),
+            "block_group_count": len(group_sizes),
+            "block_group_sizes": group_sizes,
+            "block_count": sum(group_sizes),
+        }
+
+
+def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float = 1e-2) -> bool:
+    if not np.allclose(audio1, audio2, atol=atol):
+        diff = np.abs(audio1 - audio2)
+        print(f"Max difference: {diff.max()}")
+        print(f"Mean difference: {diff.mean()}")
         raise AssertionError(f"Audio outputs differ beyond tolerance atol={atol}")
     return True
 
 
+def _device_used_mb(device_index: int) -> float:
+    current_omni_platform.synchronize()
+    with current_omni_platform.device(device_index):
+        free_bytes, total_bytes = current_omni_platform.mem_get_info()
+    return (total_bytes - free_bytes) / (1024**2)
+
+
+def _extract_audio(output: Any) -> np.ndarray | None:
+    if not output:
+        return None
+    multimodal_output = getattr(output[0], "multimodal_output", None)
+    if not isinstance(multimodal_output, dict):
+        return None
+    audio = multimodal_output.get("audio")
+    if audio is None:
+        return None
+    if isinstance(audio, torch.Tensor):
+        return audio.detach().cpu().numpy()
+    return np.asarray(audio)
+
+
+def _collect_offload_states(value: Any) -> list[dict[str, Any]]:
+    if isinstance(value, dict):
+        return [value] if value.get("probe") == _OFFLOAD_STATE_PROBE else []
+    if isinstance(value, (list, tuple)):
+        states: list[dict[str, Any]] = []
+        for item in value:
+            states.extend(_collect_offload_states(item))
+        return states
+    return []
+
+
+def _measurement_order() -> tuple[bool, bool]:
+    """Return the requested order without changing the default CI workload."""
+    configured = os.environ.get(_MEASUREMENT_ORDER_ENV, "baseline-first")
+    try:
+        return _MEASUREMENT_ORDERS[configured]
+    except KeyError as exc:
+        allowed = ", ".join(_MEASUREMENT_ORDERS)
+        raise ValueError(f"{_MEASUREMENT_ORDER_ENV} must be one of: {allowed}; got {configured!r}") from exc
+
+
 def run_inference(
     model_name: str,
     layerwise_offload: bool = False,
     num_inference_steps: int = 3,
-) -> tuple[float, Any]:
+) -> dict[str, Any]:
     current_omni_platform.empty_cache()
     device_index = current_omni_platform.current_device()
-    with current_omni_platform.device(device_index):
-        free_bytes, total_bytes = current_omni_platform.mem_get_info()
-    initial_used_mb = (total_bytes - free_bytes) / (1024**2)
+    initial_used_mb = _device_used_mb(device_index)
 
     if model_name in AUDIO_MODEL:
         params = AUDIO_MODEL_PARAMS
@@ -80,10 +155,17 @@ def run_inference(
     with OmniRunner(
         model_name,
         enable_layerwise_offload=layerwise_offload,
+        worker_extension_cls=f"{OffloadStateProbe.__module__}.{OffloadStateProbe.__qualname__}",
         # TODO: we might want to add overlapped feature e2e tests
         # cache_backend="cache_dit",
         **params["runner_params"],
     ) as runner:
+        offload_states = _collect_offload_states(
+            runner.omni.engine.collective_rpc(method="get_offload_state_for_test", timeout=60)
+        )
+        if not offload_states:
+            raise AssertionError("The offload-state worker probe returned no diffusion worker results")
+
         # Measure steady-state inference memory, not model construction. Enabling
         # layerwise offload first loads the model and then replaces each block's
         # device storage with CPU-backed weights.  Monitoring that transition
@@ -108,20 +190,135 @@ def run_inference(
         finally:
             monitor.stop()
 
+        audio = _extract_audio(output)
+        del output
+
     # DeviceMemoryMonitor reports absolute device usage. Subtract this run's
-    # starting usage so process-wide compiler/workspace caches retained from a
-    # previous run do not make the second measurement order-dependent.
-    peak = max(0.0, monitor.peak_used_mb - initial_used_mb)
+    # starting usage. Each mode runs in its own spawned process below, so model
+    # objects and process-local allocator/compiler/backend workspace state cannot
+    # carry into the other measurement. Shared filesystem caches can persist;
+    # _MEASUREMENT_ORDER_ENV enables a reversed-order validation run for that.
+    peak_used_mb = monitor.peak_used_mb
+    incremental_peak_mb = max(0.0, peak_used_mb - initial_used_mb)
 
+    del runner
     gc.collect()
+    cleanup_dist_env_and_memory()
     current_omni_platform.empty_cache()
+    post_cleanup_used_mb = _device_used_mb(device_index)
 
-    return peak, output
+    return {
+        "initial_used_mb": initial_used_mb,
+        "peak_used_mb": peak_used_mb,
+        "incremental_peak_mb": incremental_peak_mb,
+        "post_cleanup_used_mb": post_cleanup_used_mb,
+        "audio": audio,
+        "offload_states": offload_states,
+    }
+
+
+def run_inference_isolated(
+    model_name: str,
+    layerwise_offload: bool = False,
+    num_inference_steps: int = 3,
+) -> dict[str, Any]:
+    """Run one measurement in a fresh interpreter and accelerator context."""
+    with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=mp.get_context("spawn")) as executor:
+        return executor.submit(run_inference, model_name, layerwise_offload, num_inference_steps).result()
+
+
+def _assert_offload_state(measurement: dict[str, Any], *, expected_enabled: bool) -> None:
+    states = measurement["offload_states"]
+    assert states, "Expected at least one offload-state result"
+    for state in states:
+        assert state["requested"] is expected_enabled, f"Unexpected offload request state: {state}"
+        assert state["enabled"] is expected_enabled, f"Unexpected offload state: {state}"
+        if expected_enabled:
+            assert state["backend_type"], f"Enabled offloader has no backend type: {state}"
+            assert state["block_group_count"] > 0, f"Enabled offloader has no block groups: {state}"
+            assert state["block_count"] > 0, f"Enabled offloader has no managed blocks: {state}"
+        else:
+            assert state["block_group_count"] == 0, f"Disabled offloader retained block groups: {state}"
+            assert state["block_count"] == 0, f"Disabled offloader retained managed blocks: {state}"
+
+
+def _assert_audio_outputs(
+    model_name: str,
+    audio_no_offload: np.ndarray | None,
+    audio_offload: np.ndarray | None,
+) -> None:
+    if model_name not in AUDIO_MODEL:
+        return
+
+    assert audio_no_offload is not None, "Baseline Stable Audio inference returned no audio"
+    assert audio_offload is not None, "Layerwise-offloaded Stable Audio inference returned no audio"
+    # Match the sibling cpu-offload test's tolerance: layerwise offload moves
+    # blocks across the PCIe bus on a side stream, which can perturb cuBLAS
+    # algorithm selection and produce ~ULP-level drift larger than 1e-3.
+    check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2)
+
+
+def _print_measurement(label: str, measurement: dict[str, Any]) -> None:
+    print(
+        f"{label}: initial={measurement['initial_used_mb']:.1f} MB, "
+        f"peak={measurement['peak_used_mb']:.1f} MB, "
+        f"incremental_peak={measurement['incremental_peak_mb']:.1f} MB, "
+        f"post_cleanup={measurement['post_cleanup_used_mb']:.1f} MB"
+    )
+    print(f"{label} offload state: {measurement['offload_states']}")
 
 
 @pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
-@pytest.mark.parametrize("model_name", list(MODELS.keys()))
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("configured", "expected"),
+    [
+        (None, (False, True)),
+        ("baseline-first", (False, True)),
+        ("offload-first", (True, False)),
+    ],
+)
+def test_measurement_order(monkeypatch, configured: str | None, expected: tuple[bool, bool]) -> None:
+    if configured is None:
+        monkeypatch.delenv(_MEASUREMENT_ORDER_ENV, raising=False)
+    else:
+        monkeypatch.setenv(_MEASUREMENT_ORDER_ENV, configured)
+
+    assert _measurement_order() == expected
+
+
+@pytest.mark.diffusion
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_measurement_order_rejects_unknown_value(monkeypatch) -> None:
+    monkeypatch.setenv(_MEASUREMENT_ORDER_ENV, "unknown")
+
+    with pytest.raises(ValueError, match=_MEASUREMENT_ORDER_ENV):
+        _measurement_order()
+
+
+@pytest.mark.diffusion
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("audio_no_offload", "audio_offload", "message"),
+    [
+        (None, None, "Baseline Stable Audio inference returned no audio"),
+        (None, np.zeros(1), "Baseline Stable Audio inference returned no audio"),
+        (np.zeros(1), None, "Layerwise-offloaded Stable Audio inference returned no audio"),
+    ],
+)
+def test_stable_audio_requires_both_outputs(audio_no_offload, audio_offload, message) -> None:
+    with pytest.raises(AssertionError, match=message):
+        _assert_audio_outputs("stabilityai/stable-audio-open-1.0", audio_no_offload, audio_offload)
+
+
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "model_name",
+    [pytest.param(name, marks=MODEL_MARKS[name]) for name in MODELS],
+)
 def test_layerwise_offload_diffusion_model(model_name: str):
     """Test that layerwise offloading reduces GPU memory usage.
 
@@ -130,14 +327,15 @@ def test_layerwise_offload_diffusion_model(model_name: str):
     offloader keeps only a single transformer block on GPU at a time, with
     prefetching for compute-memory overlap.
     """
+    measurements: dict[bool, dict[str, Any]] = {}
     try:
-        # Run without layerwise offloading (baseline)
-        no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False)
-        cleanup_dist_env_and_memory()
-
-        # Run with layerwise offloading (1 layer on device)
-        layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True)
-        cleanup_dist_env_and_memory()
+        measurement_order = _measurement_order()
+        print(
+            "Measurement order: "
+            + " -> ".join("layerwise-offload" if enabled else "baseline" for enabled in measurement_order)
+        )
+        for enabled in measurement_order:
+            measurements[enabled] = run_inference_isolated(model_name, layerwise_offload=enabled)
     except ValueError as exc:
         # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing.
         if "Access to model" in str(exc) and "is restricted" in str(exc):
@@ -145,20 +343,18 @@ def test_layerwise_offload_diffusion_model(model_name: str):
                 f"Skipping: gated HF repo {model_name!r} inaccessible "
                 f"({exc}). See docs/contributing/ci/hf_credentials.md."
             )
-        pytest.fail(f"Inference failed: {exc}")
-    except Exception:
-        pytest.fail("Inference failed")
+        raise
 
-    print(f"Layerwise offload peak memory (1 GPU layer): {layerwise_offload_peak_memory} MB")
-    print(f"No offload peak memory: {no_offload_peak_memory} MB")
+    no_offload = measurements[False]
+    layerwise_offload = measurements[True]
+    _assert_offload_state(no_offload, expected_enabled=False)
+    _assert_offload_state(layerwise_offload, expected_enabled=True)
+    _print_measurement("No offload", no_offload)
+    _print_measurement("Layerwise offload", layerwise_offload)
 
-    if model_name == "stabilityai/stable-audio-open-1.0":
-        audio_offload = output_offload[0].multimodal_output.get("audio")
-        audio_no_offload = output_no_offload[0].multimodal_output.get("audio")
-        # Match the sibling cpu-offload test's tolerance: layerwise offload moves
-        # blocks across the PCIe bus on a side stream, which can perturb cuBLAS
-        # algorithm selection and produce ~ULP-level drift larger than 1e-3.
-        check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2)
+    audio_no_offload = no_offload["audio"]
+    audio_offload = layerwise_offload["audio"]
+    _assert_audio_outputs(model_name, audio_no_offload, audio_offload)
 
     is_rocm = torch.version.hip is not None
     platform = "rocm" if is_rocm else "cuda"
@@ -170,6 +366,8 @@ def test_layerwise_offload_diffusion_model(model_name: str):
 
     # Verify that layerwise offloading significantly reduces memory usage
     # Passes only if the actual savings meets the expected savings
+    no_offload_peak_memory = no_offload["incremental_peak_mb"]
+    layerwise_offload_peak_memory = layerwise_offload["incremental_peak_mb"]
     actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory
     assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, (
         f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB "

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -32,6 +32,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     TensorBinding,
     build_checkpoint_mmap_plan,
 )
+from vllm_omni.diffusion.offloader import plan_resolver
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
 from vllm_omni.diffusion.offloader.block_discovery import (
     get_blocks_attr_names,
@@ -52,7 +53,7 @@ from vllm_omni.diffusion.offloader.offload_plan import (
     OffloadPlan,
     get_offload_plan,
 )
-from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan
+from vllm_omni.diffusion.offloader.plan_resolver import ResolvedComponent, resolve_offload_plan
 from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
 from vllm_omni.host_weight_runtime import MappedHostRegion
 from vllm_omni.platforms import current_omni_platform
@@ -1235,6 +1236,78 @@ class _MmapPostLoadPipeline(nn.Module):
 
 
 class TestMmapWeightLoading:
+    @pytest.mark.parametrize("initial_device", ["cpu", "meta"])
+    @pytest.mark.parametrize("use_allgather", [False, True])
+    def test_mmap_keeps_undeclared_nested_blocks_streamed(
+        self, tmp_path, patched_offload_runtime, monkeypatch, initial_device, use_allgather
+    ):
+        class Refiner(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+
+            def forward(self, x):
+                for block in self.layers:
+                    x = block(x)
+                return x
+
+        class Transformer(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+                self.refiner = Refiner()
+
+            def forward(self, x):
+                x = self.refiner(x)
+                for block in self.blocks:
+                    x = block(x)
+                return x
+
+        pipeline = nn.Module()
+        pipeline.transformer = Transformer()
+        weights = {name: param.detach().clone() for name, param in pipeline.named_parameters()}
+        x = torch.randn(1, 2)
+        with torch.no_grad():
+            expected = pipeline.transformer(x)
+        checkpoint_file = tmp_path / "model.safetensors"
+        save_file(weights, str(checkpoint_file))
+        # Exercise the real mmap path with a small model instead of allocating
+        # the 1 GiB needed to trigger the production compatibility threshold.
+        monkeypatch.setattr(plan_resolver, "_NESTED_OFFLOAD_THRESHOLD_MB", 0)
+        if initial_device == "meta":
+            pipeline.transformer.to_empty(device="meta")
+        plan = HostWeightPlan(
+            backing_kind="checkpoint_mmap",
+            bindings={name: TensorBinding(checkpoint_key=name, file_path=str(checkpoint_file)) for name in weights},
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=use_allgather,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+
+        for _ in range(2):
+            try:
+                backend.enable(pipeline)
+                assert len(backend._all_hook_groups) == 2
+                with torch.no_grad():
+                    torch.testing.assert_close(pipeline.transformer(x), expected)
+            finally:
+                backend.disable()
+            for name, parameter in pipeline.named_parameters():
+                assert not parameter.is_meta
+                torch.testing.assert_close(parameter, weights[name])
+            for module in pipeline.modules():
+                registry = getattr(module, "_hook_registry", None)
+                if registry is not None:
+                    assert registry.get_hook("distributed_layerwise_offload") is None
+
     def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime):
         pipeline = _MmapPostLoadPipeline()
         weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()}
@@ -1250,10 +1323,7 @@ class TestMmapWeightLoading:
             ),
             torch.device("cpu"),
         )
-        modules = SimpleNamespace(
-            dits=[pipeline.transformer],
-            dit_names=["transformer"],
-        )
+        dits = (ResolvedComponent(path="transformer", module=pipeline.transformer, selected=True),)
         plan = HostWeightPlan(
             backing_kind="checkpoint_mmap",
             bindings={
@@ -1265,7 +1335,7 @@ class TestMmapWeightLoading:
             },
         )
 
-        backend._load_weights_via_mmap(pipeline, modules, plan)
+        backend._load_weights_via_mmap(pipeline, dits, plan)
 
         assert pipeline.transformer.post_load_calls == 1
         assert pipeline.transformer.time_embedder.weight.dtype == torch.float32
@@ -1717,6 +1787,128 @@ class TestOffloadPlan:
         for name, parameter in pipeline.named_parameters():
             torch.testing.assert_close(parameter, expected_parameters[name])
 
+    @pytest.mark.parametrize("shared_owner", ["parent", "encoder", "sibling"])
+    def test_nested_duplicate_ownership_is_rejected_before_any_hook(self, patched_offload_runtime, shared_owner):
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(
+                offload_submodules={"refiner": "blocks", "other_refiner": "blocks"},
+                encoder_block_attrs={"text_encoder": ("encoder.block",)},
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.transformer.refiner = _SingleBlockModel(num_blocks=2)
+                self.transformer.other_refiner = _SingleBlockModel(num_blocks=2)
+                self.text_encoder = _PlainEncoder()
+                owners = {
+                    "parent": self.transformer.blocks,
+                    "encoder": self.text_encoder.encoder.block,
+                    "sibling": self.transformer.other_refiner.blocks,
+                }
+                self.transformer.refiner.blocks = owners[shared_owner]
+
+        pipeline = Pipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False),
+            torch.device("cpu"),
+        )
+        storage = {name: parameter.data_ptr() for name, parameter in pipeline.named_parameters()}
+        try:
+            with pytest.raises(ValueError, match="claimed by both"):
+                backend.enable(pipeline)
+            assert not backend.enabled
+            assert not backend._all_hook_groups
+            assert {name: parameter.data_ptr() for name, parameter in pipeline.named_parameters()} == storage
+            assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)
+            for module in pipeline.modules():
+                assert getattr(module, "_hook_registry", None) is None
+        finally:
+            backend.disable()
+
+    def test_unstageable_submodule_is_rejected_before_any_hook(self, patched_offload_runtime):
+        """Plan-dependent rejection happens before earlier components are hooked."""
+
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(
+                offload_submodules={"refiner": "missing"},
+                encoder_component_types={"text_encoder": "text_encoder"},
+                encoder_block_attrs={"text_encoder": ("encoder.block",)},
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                # Declared for submodule offload but without a block list or the
+                # load_to_device/offload_to_cpu lifecycle.
+                self.transformer.refiner = nn.Linear(2, 2)
+                self.text_encoder = _PlainEncoder()
+
+        pipeline = Pipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False),
+            torch.device("cpu"),
+        )
+
+        # Storage identity proves the encoder was never staged and rolled back.
+        encoder_storage = [block.weight.data_ptr() for block in pipeline.text_encoder.encoder.block]
+
+        with pytest.raises(ValueError, match="must implement load_to_device"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        assert [block.weight.data_ptr() for block in pipeline.text_encoder.encoder.block] == encoder_storage
+        assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)
+        for block in pipeline.transformer.blocks:
+            assert getattr(block, "_hook_registry", None) is None
+
+    def test_undeclared_submodule_block_scan_warns(self, patched_offload_runtime, monkeypatch):
+        """The size-triggered attribute scan still works and is deprecated."""
+
+        class Refiner(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+
+        class Transformer(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+                self.refiner = Refiner()
+
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+        warnings: list[str] = []
+
+        class _Recorder:
+            def warning(self, message, *args):
+                warnings.append(message % args if args else message)
+
+            def __getattr__(self, _name):
+                return lambda *args, **kwargs: None
+
+        monkeypatch.setattr(plan_resolver, "logger", _Recorder())
+        monkeypatch.setattr(plan_resolver, "_NESTED_OFFLOAD_THRESHOLD_MB", 0)
+        plan_resolver._warn_nested_block_scan.cache_clear()
+        plan_resolver._warn_legacy_discovery.cache_clear()
+        pipeline = Pipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False),
+            torch.device("cpu"),
+        )
+
+        backend.enable(pipeline)
+
+        assert len(backend._all_hook_groups) == 2
+        assert sum("OffloadPlan.offload_submodules" in message for message in warnings) == 1
+
+        backend.disable()
+
 
 class TestMmapValidation:
     """Tests for loader preflight and backend plan realization."""
@@ -2101,7 +2293,7 @@ class TestMmapValidation:
             ),
             torch.device("cpu"),
         )
-        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+        dits = (ResolvedComponent(path="transformer", module=pipeline.transformer, selected=True),)
         result = build_checkpoint_mmap_plan(
             pipeline,
             dit_modules=(("transformer", pipeline.transformer),),
@@ -2113,7 +2305,7 @@ class TestMmapValidation:
         )
 
         assert result.plan is not None
-        backend._load_weights_via_mmap(pipeline, modules, result.plan)
+        backend._load_weights_via_mmap(pipeline, dits, result.plan)
         assert pipeline.transformer.validate_called, "validate_loaded_weights should be called"
 
 

tests/diffusion/offloader/test_layerwise_backend.py

@@ -24,6 +24,11 @@ from tests.diffusion.offloader.helpers import (
     patch_offload_runtime,
 )
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.block_discovery import (
+    get_blocks_attr_names,
+    get_blocks_from_dit,
+    set_blocks_attr_names,
+)
 from vllm_omni.diffusion.offloader.layerwise_backend import (
     LayerWiseOffloadBackend,
     LayerwiseOffloadHook,
@@ -209,21 +214,21 @@ class _NoAttrsModel(nn.Module):
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):
         model = _SingleBlockModel(num_blocks=3)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == ["blocks"]
         assert len(blocks) == 3
         assert all(isinstance(b, _DummyBlock) for b in blocks)
 
     def test_get_blocks_from_dit_multi_block_attrs(self):
         model = _MultiBlockModel(num_transformer=2, num_single=3)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"}
         assert len(blocks) == 5
         assert all(isinstance(b, _DummyBlock) for b in blocks)
 
     def test_get_blocks_from_dit_empty_blocks(self):
         model = _EmptyBlocksModel()
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == []
         assert blocks == []
 
@@ -233,17 +238,17 @@ class TestGetBlocksFromDit:
             AttributeError,
             match="Attribute 'nonexistent_blocks' declared in _layerwise_offload_blocks_attrs does not exist",
         ):
-            LayerWiseOffloadBackend.get_blocks_from_dit(model)
+            get_blocks_from_dit(model)
 
     def test_get_blocks_from_dit_no_attrs_defined(self):
         model = _NoAttrsModel(num_blocks=3)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == []
         assert blocks == []
 
     def test_get_blocks_from_dit_deprecated_single_attr(self):
         model = _DeprecatedSingleAttrModel(num_blocks=2)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == ["blocks"]
         assert len(blocks) == 2
 
@@ -251,17 +256,17 @@ class TestGetBlocksFromDit:
 class TestGetBlocksAttrNames:
     def test_get_blocks_attr_names_new_format(self):
         model = _MultiBlockModel()
-        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        attrs = get_blocks_attr_names(model)
         assert attrs == ["transformer_blocks", "single_transformer_blocks"]
 
     def test_get_blocks_attr_names_no_attrs(self):
         model = _NoAttrsModel()
-        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        attrs = get_blocks_attr_names(model)
         assert attrs == []
 
     def test_set_blocks_attr_names(self):
         model = _NoAttrsModel()
-        LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"])
+        set_blocks_attr_names(model, ["new_blocks"])
         assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs")
         assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"]